Skip to content

Set per-device kernel shared memory on every GPU #806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tvogels wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Set per-device kernel shared memory on every GPU #806

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion gpu4pyscf/grad/rhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@
# dd_cache_size = nf * min(THREADS, _nearest_power2(SHM_SIZE//(g_size*3*8)))
DD_CACHE_MAX = 101250 * (SHM_SIZE//(45*1024))

libvhf_rys.RYS_build_vjk_ip1_init(ctypes.c_int(SHM_SIZE))
libvhf_rys.RYS_build_vjk_ip1_init.restype = ctypes.c_int

def _jk_energy_per_atom(vhfopt, dm, j_factor=1., k_factor=1., verbose=None):
'''
Expand Down Expand Up @@ -97,6 +97,9 @@ def proc():
device_id = cp.cuda.device.get_device_id()
log = logger.new_logger(mol, verbose)
cput0 = log.init_timer()
err = libvhf_rys.RYS_build_vjk_ip1_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('RYS build_vjk_ip1 CUDA kernel initialization failed')

timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
Expand Down
10 changes: 8 additions & 2 deletions gpu4pyscf/hessian/rhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@
ALIGNED = 4
DD_CACHE_MAX = rhf_grad.DD_CACHE_MAX

libvhf_rys.RYS_build_vjk_ip1_init(ctypes.c_int(SHM_SIZE))
libvhf_rys.RYS_build_ejk_ip2_init(ctypes.c_int(SHM_SIZE))
libvhf_rys.RYS_build_vjk_ip1_init.restype = ctypes.c_int
libvhf_rys.RYS_build_ejk_ip2_init.restype = ctypes.c_int

def hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
mo1=None, mo_e1=None, h1mo=None,
Expand Down Expand Up @@ -209,6 +209,9 @@ def proc():

timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libvhf_rys.RYS_build_ejk_ip2_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('RYS build_ejk_ip2 CUDA kernel initialization failed')
kern1 = libvhf_rys.RYS_per_atom_jk_ip2_type12
kern2 = libvhf_rys.RYS_per_atom_jk_ip2_type3

Expand Down Expand Up @@ -426,6 +429,9 @@ def proc():
timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
kern = libvhf_rys.RYS_build_jk_ip1
err = libvhf_rys.RYS_build_vjk_ip1_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('RYS build_vjk_ip1 CUDA kernel initialization failed')

_dms = cp.asarray(dms)

Expand Down
4 changes: 4 additions & 0 deletions gpu4pyscf/pbc/scf/j_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
]

libvhf_md.PBC_build_j.restype = ctypes.c_int
libvhf_md.init_mdj_constant.restype = ctypes.c_int

def get_j(cell, dm, hermi=0, kpts=None, kpts_band=None, vhfopt=None,
verbose=None):
Expand Down Expand Up @@ -242,6 +243,9 @@ def proc(dm_xyz):

timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libvhf_md.init_mdj_constant(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError(f'MD-J kernel init failed on Device {device_id}')
kern = libvhf_md.PBC_build_j
rys_envs = self.rys_envs

Expand Down
18 changes: 15 additions & 3 deletions gpu4pyscf/pbc/scf/rsjk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,8 +55,9 @@
]

libpbc.PBC_build_k.restype = ctypes.c_int
libpbc.PBC_build_k_init(ctypes.c_int(SHM_SIZE))
libpbc.PBC_build_jk_ip1_init(ctypes.c_int(SHM_SIZE))
libpbc.PBC_build_k_init.restype = ctypes.c_int
libpbc.PBC_build_jk_ip1_init.restype = ctypes.c_int
libpbc.PBC_build_j_init.restype = ctypes.c_int
libpbc.PBC_per_atom_jk_ip1.restype = ctypes.c_int
libpbc.PBC_jk_strain_deriv.restype = ctypes.c_int

Expand Down Expand Up @@ -349,6 +350,9 @@ def proc(dms, dm_cond):

timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libpbc.PBC_build_k_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError(f'PBC build_k kernel init failed on Device {device_id}')
kern = libpbc.PBC_build_k
rys_envs = self.rys_envs
rsjk_omega = -self.omega
Expand Down Expand Up @@ -711,7 +715,6 @@ def _get_j_sr(self, dm, hermi, kpts=None, kpts_band=None):

diffuse_exps, diffuse_ctr_coef = extract_pgto_params(supmol, 'diffuse')

libpbc.PBC_build_j_init(ctypes.c_int(SHM_SIZE))
libpbc.PBC_build_j.restype = ctypes.c_int

uniq_l_ctr = cell.uniq_l_ctr
Expand Down Expand Up @@ -747,6 +750,9 @@ def proc(dms, dm_cond):

timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libpbc.PBC_build_j_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError(f'PBC build_j kernel init failed on Device {device_id}')
kern = libpbc.PBC_build_j
rys_envs = self.rys_envs
rsjk_omega = -self.omega
Expand Down Expand Up @@ -1003,6 +1009,9 @@ def proc(dms, dm_cond):
t1 = log.timer_debug1(f'ejk_sr initialization on Device {device_id}', *t0)
timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libpbc.PBC_build_jk_ip1_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError(f'PBC build_jk_ip1 kernel init failed on Device {device_id}')
kern = libpbc.PBC_per_atom_jk_ip1
rys_envs = self.rys_envs
omega = -self.omega
Expand Down Expand Up @@ -1433,6 +1442,9 @@ def proc(dms, dm_cond):
t1 = log.timer_debug1(f'ejk_sr_strain_deriv initialization on Device {device_id}', *t0)
timing_collection = _TimingCollector(log.timer_debug1)
kern_counts = 0
err = libpbc.PBC_build_jk_ip1_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError(f'PBC build_jk_ip1 kernel init failed on Device {device_id}')
kern = libpbc.PBC_jk_strain_deriv
rys_envs = self.rys_envs
omega = -self.omega
Expand Down
5 changes: 4 additions & 1 deletion gpu4pyscf/scf/j_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@

libvhf_md = load_library('libgvhf_md')
libvhf_md.MD_build_j.restype = ctypes.c_int
libvhf_md.init_mdj_constant(ctypes.c_int(SHM_SIZE))
libvhf_md.init_mdj_constant.restype = ctypes.c_int

def get_j(mol, dm, hermi=1, vhfopt=None, verbose=None):
'''Compute J matrix
Expand Down Expand Up @@ -161,6 +161,9 @@ def proc(dm_xyz):
stream = cp.cuda.stream.get_current_stream()
log = logger.new_logger(self.mol, verbose)
t1 = log.init_timer()
err = libvhf_md.init_mdj_constant(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('MD-J CUDA kernel initialization failed')
dm_xyz = asarray(dm_xyz) # transfer to current device
vj_xyz = cp.zeros_like(dm_xyz)

Expand Down
10 changes: 8 additions & 2 deletions gpu4pyscf/scf/jk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,8 @@
GROUP_SIZE = 256
Q_COND_MARGIN = 4.

libvhf_rys.RYS_build_k_init(ctypes.c_int(SHM_SIZE))
libvhf_rys.RYS_build_jk_init(ctypes.c_int(SHM_SIZE))
libvhf_rys.RYS_build_k_init.restype = ctypes.c_int
libvhf_rys.RYS_build_jk_init.restype = ctypes.c_int

def get_jk(mol, dm, hermi=0, vhfopt=None, with_j=True, with_k=True, verbose=None):
'''Compute J, K matrices
Expand Down Expand Up @@ -444,6 +444,9 @@ def proc(dms, dm_cond):
stream = cp.cuda.get_current_stream()
log = logger.new_logger(mol, verbose)
t0 = log.init_timer()
err = libvhf_rys.RYS_build_jk_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('RYS build_jk CUDA kernel initialization failed')
dms = cp.asarray(dms) # transfer to current device
dm_cond = cp.asarray(dm_cond)

Expand Down Expand Up @@ -739,6 +742,9 @@ def proc(dms, dm_cond):
stream = cp.cuda.stream.get_current_stream()
log = logger.new_logger(mol, verbose)
t0 = log.init_timer()
err = libvhf_rys.RYS_build_k_init(ctypes.c_int(SHM_SIZE))
if err != 0:
raise RuntimeError('RYS build_k CUDA kernel initialization failed')
dms = cp.asarray(dms) # transfer to current device
dm_cond = cp.asarray(dm_cond)

Expand Down
Loading