rendered paste bodyIndex: 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_qMMIndex: 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 thisIndex: 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