rendered paste bodyIndex: blacs.py===================================================================--- blacs.py (revision 6096)+++ blacs.py (working copy)@@ -86,10 +86,12 @@ import numpy as np -from gpaw import sl_diagonalize from gpaw.mpi import SerialCommunicator, serial_comm from gpaw.matrix_descriptor import MatrixDescriptor+from gpaw.utilities import uncamelcase from gpaw.utilities.blas import gemm, r2k, gemmdot+from gpaw.utilities.lapack import general_diagonalize, \+ diagonalize, sldiagonalize from gpaw.utilities.blacs import scalapack_inverse_cholesky, \ scalapack_diagonalize_ex, scalapack_general_diagonalize_ex, \ scalapack_diagonalize_dc, scalapack_general_diagonalize_dc, \@@ -488,78 +490,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 +646,114 @@ # ------------------------------------------------------------------- -class BlacsLayouts:- def __init__(self, gd, bd, ncpus, mcpus, blocksize, timer=nulltimer):+def get_kohn_sham_layouts(mode, use_blacs, gd, bd, **kwargs):+ """Create Kohn-Sham layouts object."""+ if not isinstance(mode, str):+ return None #XXX+ name = {'fd': 'BandLayouts', 'lcao': 'OrbitalLayouts'}[mode]+ if use_blacs:+ assert 'mcpus' in kwargs+ assert 'ncpus' in kwargs+ assert 'blocksize' in kwargs+ name = 'Blacs' + name+ ksl = {'BandLayouts': BandLayouts,+ 'BlacsBandLayouts': BlacsBandLayouts,+ 'BlacsOrbitalLayouts': BlacsOrbitalLayouts,+ 'OrbitalLayouts': OrbitalLayouts,+ }[name](gd, bd, **kwargs)+ assert isinstance(ksl, KohnShamLayouts)+ assert isinstance(ksl, BlacsLayouts) == use_blacs, (ksl, use_blacs)+ return ksl+++class KohnShamLayouts:+ using_blacs = False++ 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+ self._kwargs = {'timer': timer} + #def get_diagonalizer(self):+ # raise RuntimeError('Virtual member function should not be called.')++ def get_keywords(self):+ return self._kwargs.copy() # just a shallow copy...++ def diagonalize(self, *args, **kwargs):+ raise RuntimeError('Virtual member function should not be called.')++ def __repr__(self):+ return uncamelcase(self.__class__.__name__)+++class BlacsLayouts(KohnShamLayouts):+ using_blacs = True++ def __init__(self, gd, bd, ncpus, mcpus, blocksize, timer=nulltimer):+ KohnShamLayouts.__init__(self, gd, bd, timer)+ self._kwargs.update({'mcpus': mcpus, 'ncpus': ncpus, 'blocksize': blocksize})+ 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 BandLayouts(KohnShamLayouts): + def diagonalize(self, H_NN, eps_n):+ nbands = self.bd.nbands+ eps_N = np.empty(nbands)+ info = self._diagonalize(H_NN, eps_N)+ if info != 0:+ raise RuntimeError('Failed to diagonalize: %d' % info)++ if self.gd.comm.rank == 0:+ self.bd.distribute(eps_N, eps_n)+ self.bd.comm.broadcast(H_NN, 0)++ self.gd.comm.broadcast(H_NN, 0)+ self.gd.comm.broadcast(eps_n, 0)++ def _diagonalize(self, H_NN, eps_N):+ """Serial diagonalizer."""+ # Only one processor really does any work.+ if self.gd.comm.rank == 0 and self.bd.comm.rank == 0:+ return diagonalize(H_NN, eps_N)+ else:+ return 0+++class OldSLBandLayouts(BandLayouts): #old SL before BLACS grids. TODO delete!+ """Original ScaLAPACK diagonalizer using + redundantly distributed arrays."""+ def __init__(self, gd, bd, timer=nulltimer, root=0):+ BandLayouts.__init__(self, gd, bd, timer)+ bcommsize = self.bd.comm.size+ gcommsize = self.gd.comm.size+ shiftks = self.world.rank - self.world.rank % (bcommsize * gcommsize)+ block_ranks = shiftks + np.arange(bcommsize * gcommsize)+ self.blockcomm = self.world.new_communicator(block_ranks)+ self.root = root+ # Keep buffers?++ def _diagonalize(self, H_NN, eps_N):+ # Work is done on BLACS grid, but one processor still collects+ # all eigenvectors. Only processors on the BLACS grid return+ # meaningful values of info.+ return sldiagonalize(H_NN, eps_N, self.blockcomm, root=self.root)++ class BlacsBandLayouts(BlacsLayouts): # This class 'describes' all the realspace Blacs-related layouts def __init__(self, gd, bd, mcpus, ncpus, blocksize, timer=nulltimer):@@ -765,22 +779,53 @@ # 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) + def diagonalize(self, H_nn, eps_n):+ nbands = self.bd.nbands+ eps_N = np.empty(nbands)+ C_nn = self._diagonalize(H_nn, eps_N)+ if self.gd.comm.rank == 0:+ self.bd.distribute(eps_N, eps_n)+ self.gd.comm.broadcast(eps_n, 0)+ return C_nn + def _diagonalize(self, H_nn, eps_N):+ """Parallel diagonalizer."""+ C_nn = self.nndescriptor.empty(dtype=H_nn.dtype)+ self.nndescriptor.diagonalize_dc(H_nn, C_nn, eps_N, 'L')+ return C_nn++ 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)+ self._kwargs.update(nao=nao) #XXX should only be done by OrbitalLayouts! nbands = bd.nbands mynbands = bd.mynbands@@ -828,10 +873,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 +969,10 @@ 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._kwargs.update(nao=nao) self.mMdescriptor = MatrixDescriptor(nao, nao) self.nMdescriptor = MatrixDescriptor(bd.mynbands, nao) @@ -898,11 +983,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: hs_operators.py===================================================================--- hs_operators.py (revision 6096)+++ hs_operators.py (working copy)@@ -26,7 +26,7 @@ async = True hermitian = True - def __init__(self, bd, gd, nblocks=None, async=None, hermitian=None):+ def __init__(self, bd, gd, ksl, nblocks=None, async=None, hermitian=None): self.bd = bd self.gd = gd self.work1_xG = None@@ -39,7 +39,12 @@ self.async = async if hermitian is not None: self.hermitian = hermitian- self.bmd = BandMatrixDescriptor(bd, gd) #XXX prefix Blacs for 1D layout+ if ksl.using_blacs: #XXX prefix Blacs for 1D layout+ kwargs = ksl.get_keywords()+ kwargs.pop('timer', None)+ self.bmd = BlacsBandMatrixDescriptor(bd, gd, **kwargs)+ else:+ self.bmd = BandMatrixDescriptor(bd, gd) def allocate_work_arrays(self, dtype): """This is a little complicated, but let's look at the facts.Index: eigensolvers/eigensolver.py===================================================================--- eigensolvers/eigensolver.py (revision 6096)+++ eigensolvers/eigensolver.py (working copy)@@ -6,70 +6,12 @@ from gpaw.fd_operators import Laplace from gpaw.preconditioner import Preconditioner-from gpaw.utilities.lapack import diagonalize, sldiagonalize from gpaw.utilities.blas import axpy, r2k, gemm from gpaw.utilities.tools import apply_subspace_mask from gpaw.utilities import unpack-from gpaw.utilities import scalapack-from gpaw import sl_diagonalize from gpaw import debug -class BaseDiagonalizer:- def __init__(self, world, gd, bd, kpt_comm):- self.world = world- self.gd = gd- self.bd = bd- self.kpt_comm = kpt_comm-- def diagonalize(self, H_NN, eps_n):- nbands = self.bd.nbands- eps_N = np.empty(nbands)- info = self._diagonalize(H_NN, eps_N)- if info != 0:- raise RuntimeError('Failed to diagonalize: %d' % info)-- if self.gd.comm.rank == 0:- self.bd.distribute(eps_N, eps_n)- self.bd.comm.broadcast(H_NN, 0)-- self.gd.comm.broadcast(H_NN, 0)- self.gd.comm.broadcast(eps_n, 0)-- def _diagonalize(self, H_NN, eps_n):- raise NotImplementedError---class SLDiagonalizer(BaseDiagonalizer):- """Original ScaLAPACK diagonalizer using - redundantly distributed arrays."""- def __init__(self, world, gd, bd, kpt_comm, root=0):- BaseDiagonalizer.__init__(self, world, gd, bd, kpt_comm)- bcommsize = bd.comm.size- gcommsize = gd.comm.size- shiftks = kpt_comm.rank * bcommsize * gcommsize- block_ranks = shiftks + np.arange(bcommsize * gcommsize)- blockcomm = world.new_communicator(block_ranks)- self.blockcomm = blockcomm- self.root = root- # Keep buffers?-- def _diagonalize(self, H_NN, eps_N):- # Work is done on BLACS grid, but one processor still collects- # all eigenvectors. Only processors on the BLACS grid return- # meaningful values of info.- return sldiagonalize(H_NN, eps_N, self.blockcomm, root=self.root)---class LapackDiagonalizer(BaseDiagonalizer):- """Serial diagonalizer."""- def _diagonalize(self, H_NN, eps_N):- # Only one processor really does any work.- if self.gd.comm.rank == 0 and self.bd.comm.rank == 0:- return diagonalize(H_NN, eps_N)- else:- return 0- class Eigensolver: def __init__(self, keep_htpsit=True): self.keep_htpsit = keep_htpsit@@ -85,6 +27,7 @@ self.dtype = wfs.dtype self.bd = wfs.bd self.gd = wfs.gd+ self.diagonalizer = wfs.ksl self.nbands = wfs.nbands self.mynbands = wfs.mynbands @@ -108,13 +51,6 @@ if kpt.eps_n is None: kpt.eps_n = np.empty(self.mynbands) - if sl_diagonalize:- self.diagonalizer = SLDiagonalizer(self.world, self.gd, - self.bd, self.kpt_comm)- else:- self.diagonalizer = LapackDiagonalizer(self.world, self.gd,- self.bd, self.kpt_comm)- self.initialized = True def iterate(self, hamiltonian, wfs):@@ -234,10 +170,10 @@ raise NotImplementedError else: H_nn = hamiltonian.xc.xcfunc.exx.grr(wfs, kpt, Htpsit_xG,- hamiltonian)+ hamiltonian) self.timer.stop('calc_matrix') - diagonalizationstring = self.diagonalizer.__class__.__name__+ diagonalizationstring = repr(self.diagonalizer) wfs.timer.start(diagonalizationstring) self.diagonalizer.diagonalize(H_nn, kpt.eps_n) # The two lines below will go away soonIndex: wavefunctions.py===================================================================--- wavefunctions.py (revision 6096)+++ 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')@@ -478,8 +477,7 @@ if S_qMM is None: # XXX # First time: assert T_qMM is None- from gpaw.blacs import BlacsOrbitalLayouts- if isinstance(self.ksl, BlacsOrbitalLayouts): # XXX+ if self.ksl.using_blacs: # XXX self.tci.set_matrix_distribution(Mstart, mynao) S_qMM = np.empty((nq, mynao, nao), self.dtype)@@ -558,8 +556,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')@@ -965,7 +961,10 @@ lcaobd = BandDescriptor(lcaonbands, self.band_comm, self.bd.strided) assert lcaobd.mynbands == lcaomynbands #XXX - lcaowfs = LCAOWaveFunctions(self.gd, self.ksl, self.nspins,+ from gpaw.blacs import get_kohn_sham_layouts+ lcaoksl = get_kohn_sham_layouts('lcao', self.ksl.using_blacs, self.gd,+ lcaobd, nao=nao, **self.ksl.get_keywords())+ lcaowfs = LCAOWaveFunctions(self.gd, lcaoksl, self.nspins, self.nvalence, self.setups, lcaobd, self.dtype, self.world, self.kpt_comm, self.gamma, self.bzk_kc, self.ibzk_kc,@@ -977,11 +976,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)+ lcaoksl) # XXX when density matrix is properly distributed, be sure to # update the density here also eigensolver.iterate(hamiltonian, lcaowfs)Index: overlap.py===================================================================--- overlap.py (revision 6096)+++ overlap.py (working copy)@@ -32,7 +32,7 @@ def __init__(self, wfs): """Create the Overlap operator."""- self.operator = MatrixOperator(wfs.bd, wfs.gd)+ self.operator = MatrixOperator(wfs.bd, wfs.gd, wfs.ksl) self.timer = wfs.timer self.domain_comm = wfs.gd.comm self.band_comm = wfs.bd.comm@@ -82,7 +82,9 @@ self.timer.stop('calc_matrix') self.timer.start('inverse_cholesky') - if sl_inverse_cholesky:+ if wfs.ksl.using_blacs: #XXX TODO BAD BAD BAD UGLY HACK+ self.operator.bmd.nndescriptor.inverse_cholesky(S_nn, 'L')+ elif sl_inverse_cholesky: assert parallel and scalapack() if inverse_cholesky(S_nn, 0) != 0: raise RuntimeError('Orthogonalization failed! You may want to check your structure.')Index: test/parallel/ut_hsblacs.py===================================================================--- test/parallel/ut_hsblacs.py (revision 6096)+++ test/parallel/ut_hsblacs.py (working copy)@@ -64,6 +64,7 @@ def verify_blacs_stuff(self): # TODO do more here :) ksl = BlacsBandLayouts(self.gd, self.bd, self.mcpus, self.ncpus, 6)+ self.assertTrue(ksl.using_blacs) class UTBandParallelBlacsSetup_Blocked(UTBandParallelBlacsSetup):Index: lcao/eigensolver.py===================================================================--- lcao/eigensolver.py (revision 6096)+++ 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 6096)+++ paw.py (working copy)@@ -516,26 +516,36 @@ domain_comm, parsize) - # Construct orbital descriptor for LCAO mode or initialization:+ # Figure out layout of Kohn-Sham orbitals for LCAO or FD mode:+ kwargs = {'timer': self.timer}++ if par.mode == 'lcao':+ kwargs['nao'] = nao+ from gpaw import extra_parameters- use_blacs = (par.parallel['scalapack']- or extra_parameters.get('blacs'))- if use_blacs:+ if par.parallel['scalapack'] or extra_parameters.get('blacs'): sl_diagonalize = par.parallel['scalapack'] if sl_diagonalize is None: from gpaw import sl_diagonalize+ mcpus, ncpus, blocksize = sl_diagonalize[:3]+ kwargs.update(mcpus=mcpus, ncpus=ncpus, blocksize=blocksize)+ use_blacs = True+ """+ args += (mcpus, ncpus, blocksize) from gpaw.blacs import BlacsOrbitalLayouts- ncpus, mcpus, blocksize = sl_diagonalize[:3] ksl = BlacsOrbitalLayouts(gd, self.bd, nao,- ncpus, mcpus, blocksize,- self.timer)+ mcpus, ncpus, blocksize, self.timer)+ """ else:+ use_blacs = False+ """ # XXX This is actually a non-BLACS object. The class should be # defined somewhere else from gpaw.blacs import OrbitalLayouts ksl = OrbitalLayouts(gd, self.bd, setups.nao)-- + """+ from gpaw.blacs import get_kohn_sham_layouts+ ksl = get_kohn_sham_layouts(par.mode, use_blacs, gd, self.bd, **kwargs) # do k-point analysis here? XXX args = (gd, ksl, nspins, nvalence, setups, self.bd, dtype, world, kpt_comm,