Refactor MPI load balancing code into a separate module.

Also use common MPI work distribution framework for computation of
polarization contributions.