Index: blacs.py =================================================================== --- blacs.py (revision 6052) +++ blacs.py (working copy) @@ -90,6 +90,7 @@ from gpaw.mpi import SerialCommunicator, serial_comm from gpaw.matrix_descriptor import MatrixDescriptor from gpaw.utilities.blas import gemm, r2k, gemmdot +from gpaw.utilities.lapack import general_diagonalize from gpaw.utilities.blacs import scalapack_inverse_cholesky, \ scalapack_diagonalize_ex, scalapack_general_diagonalize_ex, \ scalapack_diagonalize_dc, scalapack_general_diagonalize_dc, \ @@ -488,78 +489,6 @@ comm.barrier() -class SLDenseLinearAlgebra: - """ScaLAPACK Dense Linear Algebra. - - This class is instantiated in LCAO. Not for casual use, at least for now. - - Requires two distributors and three descriptors for initialization - as well as grid descriptors and band descriptors. Distributors are - for cols2blocks (1D -> 2D BLACS grid) and blocks2cols (2D -> 1D - BLACS grid). ScaLAPACK operations must occur on 2D BLACS grid for - performance and scalability. - - _general_diagonalize is "hard-coded" for LCAO. - Expects both Hamiltonian and Overlap matrix to be on the 2D BLACS grid. - This is done early on to save memory. - """ - - def __init__(self, gd, bd, cols2blocks, blocks2cols, timer=nulltimer): - self.gd = gd - self.bd = bd - assert cols2blocks.dstdescriptor == blocks2cols.srcdescriptor - self.indescriptor = cols2blocks.srcdescriptor - self.blockdescriptor = cols2blocks.dstdescriptor - self.outdescriptor = blocks2cols.dstdescriptor - self.cols2blocks = cols2blocks - self.blocks2cols = blocks2cols - self.timer = timer - - def diagonalize(self, H_mm, C_nM, eps_n, S_mm): - # C_nM needs to be simultaneously compatible with: - # 1. outdescriptor - # 2. broadcast with gd.comm - # We will does this with a dummy buffer C2_nM - indescriptor = self.indescriptor - outdescriptor = self.outdescriptor - blockdescriptor = self.blockdescriptor - - dtype = S_mm.dtype - eps_M = np.empty(C_nM.shape[-1]) # empty helps us debug - subM, subN = outdescriptor.gshape - - C_mm = blockdescriptor.zeros(dtype=dtype) - self.timer.start('General diagonalize ex') - blockdescriptor.general_diagonalize_ex(H_mm, S_mm.copy(), C_mm, eps_M, - UL='L', iu=self.bd.nbands) - self.timer.stop('General diagonalize ex') - - # Make C_nM compatible with the redistributor - self.timer.start('Redistribute coefs') - if outdescriptor: - C2_nM = C_nM - else: - C2_nM = outdescriptor.empty(dtype=dtype) - assert outdescriptor.check(C2_nM) - self.blocks2cols.redistribute_submatrix(C_mm, C2_nM, subM, subN) - self.timer.stop('Redistribute coefs') - - self.timer.start('Send coefs to domains') - if outdescriptor: # grid masters only - assert self.gd.comm.rank == 0 - bd = self.bd - # grid master with bd.rank = 0 - # scatters to other grid masters - # NOTE: If the origin of the blacs grid - # ever shifts this will not work - bd.distribute(eps_M[:bd.nbands], eps_n) - else: - assert self.gd.comm.rank != 0 - - self.gd.comm.broadcast(C_nM, 0) - self.gd.comm.broadcast(eps_n, 0) - self.timer.stop('Send coefs to domains') - class SLDenseLinearAlgebra2: """ScaLAPACK Dense Linear Algebra. @@ -716,30 +645,37 @@ # ------------------------------------------------------------------- -class BlacsLayouts: - def __init__(self, gd, bd, ncpus, mcpus, blocksize, timer=nulltimer): +class KohnShamLayouts: + def __init__(self, gd, bd, timer=nulltimer): + assert gd.comm.parent is bd.comm.parent # must have same parent comm + self.world = bd.comm.parent self.gd = gd self.bd = bd self.timer = timer + #def get_diagonalizer(self): + # raise RuntimeError('Virtual member function should not be called.') + + def _diagonalize(self, H_NN, S_NN, eps_N): + raise RuntimeError('Virtual member function should not be called.') + + +class BlacsLayouts(KohnShamLayouts): + def __init__(self, gd, bd, ncpus, mcpus, blocksize, timer=nulltimer): + KohnShamLayouts.__init__(self, gd, bd, timer) + bcommsize = self.bd.comm.size gcommsize = self.gd.comm.size - - assert gd.comm.parent is bd.comm.parent # must have same parent comm - world = bd.comm.parent - shiftks = world.rank - world.rank % (bcommsize * gcommsize) + shiftks = self.world.rank - self.world.rank % (bcommsize * gcommsize) column_ranks = shiftks + np.arange(bcommsize) * gcommsize block_ranks = shiftks + np.arange(bcommsize * gcommsize) - self.columncomm = world.new_communicator(column_ranks) #XXX rename? - self.blockcomm = world.new_communicator(block_ranks) + self.columncomm = self.world.new_communicator(column_ranks) #XXX rename? + self.blockcomm = self.world.new_communicator(block_ranks) assert ncpus * mcpus <= bcommsize * gcommsize self.blockgrid = BlacsGrid(self.blockcomm, mcpus, ncpus) - def get_diagonalizer(self): - raise RuntimeError('Virtual member function should not be called.') - class BlacsBandLayouts(BlacsLayouts): # This class 'describes' all the realspace Blacs-related layouts def __init__(self, gd, bd, mcpus, ncpus, blocksize, timer=nulltimer): @@ -765,19 +701,34 @@ # Only redistribute filled out half for Hermitian matrices self.Nn2nn = Redistributor(self.blockcomm, self.Nndescriptor, self.nndescriptor) - self.Nn2nn_lower = Redistributor(self.blockcomm, self.Nndescriptor, - self.nndescriptor, 'L') + #self.Nn2nn = Redistributor(self.blockcomm, self.Nndescriptor, + # self.nndescriptor, 'L') #XXX faster but... # Resulting matrix will be used in dgemm which is symmetry obvlious self.nn2nN = Redistributor(self.blockcomm, self.nndescriptor, self.nNdescriptor) - def get_diagonalizer(self): - return SLDenseLinearAlgebra2(self.gd, self.bd, self.Nn2nn, self.nn2nN, - self.timer) + #def get_diagonalizer(self): + # return SLDenseLinearAlgebra2(self.gd, self.bd, self.Nn2nn, self.nn2nN, + # self.timer) class BlacsOrbitalLayouts(BlacsLayouts): + """ScaLAPACK Dense Linear Algebra. + + This class is instantiated in LCAO. Not for casual use, at least for now. + + Requires two distributors and three descriptors for initialization + as well as grid descriptors and band descriptors. Distributors are + for cols2blocks (1D -> 2D BLACS grid) and blocks2cols (2D -> 1D + BLACS grid). ScaLAPACK operations must occur on 2D BLACS grid for + performance and scalability. + + _general_diagonalize is "hard-coded" for LCAO. + Expects both Hamiltonian and Overlap matrix to be on the 2D BLACS grid. + This is done early on to save memory. + """ #XXX rewrite this docstring a bit! + # This class 'describes' all the LCAO Blacs-related layouts def __init__(self, gd, bd, nao, mcpus, ncpus, blocksize, timer=nulltimer): BlacsLayouts.__init__(self, gd, bd, mcpus, ncpus, blocksize, timer) @@ -828,10 +779,50 @@ self.mm2nM = Redistributor(self.blockcomm, self.mmdescriptor, self.nM_unique_descriptor) - def get_diagonalizer(self): - return SLDenseLinearAlgebra(self.gd, self.bd, self.mM2mm, self.mm2nM, - self.timer) + def diagonalize(self, H_mm, C_nM, eps_n, S_mm): + # C_nM needs to be simultaneously compatible with: + # 1. outdescriptor + # 2. broadcast with gd.comm + # We will does this with a dummy buffer C2_nM + indescriptor = self.mM2mm.srcdescriptor #cols2blocks + outdescriptor = self.mm2nM.dstdescriptor #blocks2cols + blockdescriptor = self.mM2mm.dstdescriptor #cols2blocks + dtype = S_mm.dtype + eps_M = np.empty(C_nM.shape[-1]) # empty helps us debug + subM, subN = outdescriptor.gshape + + C_mm = blockdescriptor.zeros(dtype=dtype) + self.timer.start('General diagonalize ex') + blockdescriptor.general_diagonalize_ex(H_mm, S_mm.copy(), C_mm, eps_M, + UL='L', iu=self.bd.nbands) + self.timer.stop('General diagonalize ex') + + # Make C_nM compatible with the redistributor + self.timer.start('Redistribute coefs') + if outdescriptor: + C2_nM = C_nM + else: + C2_nM = outdescriptor.empty(dtype=dtype) + assert outdescriptor.check(C2_nM) + self.mm2nM.redistribute_submatrix(C_mm, C2_nM, subM, subN) #blocks2cols + self.timer.stop('Redistribute coefs') + + self.timer.start('Send coefs to domains') + if outdescriptor: # grid masters only + assert self.gd.comm.rank == 0 + # grid master with bd.rank = 0 + # scatters to other grid masters + # NOTE: If the origin of the blacs grid + # ever shifts this will not work + self.bd.distribute(eps_M[:self.bd.nbands], eps_n) + else: + assert self.gd.comm.rank != 0 + + self.gd.comm.broadcast(C_nM, 0) + self.gd.comm.broadcast(eps_n, 0) + self.timer.stop('Send coefs to domains') + def distribute_overlap_matrix(self, S_qmM, root=0): # Some MPI implementations need a lot of memory to do large # reductions. To avoid trouble, we do comm.sum on smaller blocks @@ -884,10 +875,9 @@ return self.calculate_density_matrix(f_n, C_nM, rho_mM) -class OrbitalLayouts: - def __init__(self, gd, bd, nao): - self.gd = gd # XXX shouldn't be necessary - self.bd = bd +class OrbitalLayouts(KohnShamLayouts): + def __init__(self, gd, bd, nao, timer=nulltimer): + KohnShamLayouts.__init__(self, gd, bd, timer) self.mMdescriptor = MatrixDescriptor(nao, nao) self.nMdescriptor = MatrixDescriptor(bd.mynbands, nao) @@ -898,11 +888,32 @@ self.nao = nao self.orbital_comm = serial_comm - def get_diagonalizer(self): - from gpaw.lcao.eigensolver import LapackDiagonalizer - diagonalizer = LapackDiagonalizer(self.gd, self.bd) - return diagonalizer + def diagonalize(self, H_MM, C_nM, eps_n, S_MM): + eps_M = np.empty(C_nM.shape[-1]) + info = self._diagonalize(H_MM, S_MM.copy(), eps_M) + if info != 0: + raise RuntimeError('Failed to diagonalize: %d' % info) + + nbands = self.bd.nbands + if self.bd.rank == 0: + self.gd.comm.broadcast(H_MM[:nbands], 0) + self.gd.comm.broadcast(eps_M[:nbands], 0) + self.bd.distribute(H_MM[:nbands], C_nM) + self.bd.distribute(eps_M[:nbands], eps_n) + + def _diagonalize(self, H_MM, S_MM, eps_M): + # Only one processor really does any work. + if self.gd.comm.rank == 0 and self.bd.comm.rank == 0: + return general_diagonalize(H_MM, eps_M, S_MM) + else: + return 0 + def estimate_memory(self, mem, dtype): + nao = self.setups.nao + itemsize = mem.itemsize[dtype] + mem.subnode('eps [M]', self.nao * mem.floatsize) + mem.subnode('H [MM]', self.nao * self.nao * itemsize) + def distribute_overlap_matrix(self, S_qMM, root=0): self.gd.comm.sum(S_qMM, root) return S_qMM Index: wavefunctions.py =================================================================== --- wavefunctions.py (revision 6052) +++ wavefunctions.py (working copy) @@ -450,8 +450,7 @@ def set_eigensolver(self, eigensolver): WaveFunctions.set_eigensolver(self, eigensolver) - eigensolver.initialize(self.gd, self.dtype, self.setups.nao, - self.ksl.get_diagonalizer()) + eigensolver.initialize(self.gd, self.dtype, self.setups.nao, self.ksl) def set_positions(self, spos_ac): self.timer.start('Basic WFS set positions') @@ -558,8 +557,6 @@ # ATLAS can't handle uninitialized output array: #rho_MM.fill(42) - #from gpaw.blacs import BlacsOrbitalLayouts - #if isinstance(self.ksl, BlacsOrbitalLayouts): self.timer.start('Calculate density matrix') rho_MM = self.ksl.calculate_density_matrix(f_n, C_nM, rho_MM) self.timer.stop('Calculate density matrix') @@ -977,11 +974,10 @@ self.timer.stop('Set positions (LCAO WFS)') eigensolver = get_eigensolver('lcao', 'lcao') - diagonalizer = lcaowfs.ksl.get_diagonalizer() eigensolver.initialize(self.gd, self.dtype, self.setups.nao, - diagonalizer) + lcaowfs.ksl) # XXX when density matrix is properly distributed, be sure to # update the density here also eigensolver.iterate(hamiltonian, lcaowfs) Index: lcao/eigensolver.py =================================================================== --- lcao/eigensolver.py (revision 6052) +++ lcao/eigensolver.py (working copy) @@ -1,50 +1,11 @@ import numpy as np from gpaw.utilities import unpack -from gpaw.utilities.lapack import general_diagonalize from gpaw.utilities.blas import gemm from gpaw.utilities import scalapack from gpaw import sl_diagonalize, extra_parameters import gpaw.mpi as mpi -class BaseDiagonalizer: - def __init__(self, gd, bd): - self.gd = gd - self.bd = bd - - def diagonalize(self, H_MM, C_nM, eps_n, S_MM): - eps_M = np.empty(C_nM.shape[-1]) - info = self._diagonalize(H_MM, S_MM.copy(), eps_M) - if info != 0: - raise RuntimeError('Failed to diagonalize: %d' % info) - - nbands = self.bd.nbands - if self.bd.rank == 0: - self.gd.comm.broadcast(H_MM[:nbands], 0) - self.gd.comm.broadcast(eps_M[:nbands], 0) - self.bd.distribute(H_MM[:nbands], C_nM) - self.bd.distribute(eps_M[:nbands], eps_n) - - def _diagonalize(self, H_MM, S_MM, eps_M): - raise NotImplementedError - - def estimate_memory(self, mem, dtype): - nao = self.setups.nao - itemsize = mem.itemsize[dtype] - mem.subnode('eps [M]', self.nao * mem.floatsize) - mem.subnode('H [MM]', self.nao * self.nao * itemsize) - - -class LapackDiagonalizer(BaseDiagonalizer): - """Serial diagonalizer.""" - def _diagonalize(self, H_MM, S_MM, eps_M): - # Only one processor really does any work. - if self.gd.comm.rank == 0 and self.bd.comm.rank == 0: - return general_diagonalize(H_MM, eps_M, S_MM) - else: - return 0 - - class LCAO: """Eigensolver for LCAO-basis calculation""" @@ -120,7 +81,7 @@ kpt.eps_n[0] = 42 - diagonalizationstring = self.diagonalizer.__class__.__name__ + diagonalizationstring = repr(self.diagonalizer) wfs.timer.start(diagonalizationstring) self.diagonalizer.diagonalize(H_MM, kpt.C_nM, kpt.eps_n, S_MM) wfs.timer.stop(diagonalizationstring) @@ -140,4 +101,4 @@ def estimate_memory(self, mem, dtype): pass - # self.diagonalizer.estimate_memory(mem, dtype) + # self.diagonalizer.estimate_memory(mem, dtype) #XXX enable this Index: paw.py =================================================================== --- paw.py (revision 6052) +++ paw.py (working copy) @@ -525,10 +525,9 @@ if sl_diagonalize is None: from gpaw import sl_diagonalize from gpaw.blacs import BlacsOrbitalLayouts - ncpus, mcpus, blocksize = sl_diagonalize[:3] + mcpus, ncpus, blocksize = sl_diagonalize[:3] ksl = BlacsOrbitalLayouts(gd, self.bd, nao, - ncpus, mcpus, blocksize, - self.timer) + mcpus, ncpus, blocksize, self.timer) else: # XXX This is actually a non-BLACS object. The class should be # defined somewhere else