reduce h2d trasfer

This commit is contained in:
bluepulsar
2026-02-22 13:46:05 +08:00
parent 10eb891d0b
commit 3fcd3531e8
+19 -24
View File
@@ -103,7 +103,7 @@ class nHHD(object):
Dispatch Poisson solve using the selected backend.
"""
if self.poisson_backend == 'opencl':
return self.psolver.solve_opencl(field, verbose, profile=False)
return self.psolver.solve_opencl(field, verbose, profile=True)
if self.poisson_backend == 'cupy':
return self.psolver.solve_gpu(field, verbose)
@@ -126,13 +126,17 @@ class nHHD(object):
# compute div curl
(self.div, self.curlw) = self.mesh.divcurl(vfield)
if self.poisson_backend == 'opencl':
self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32)
self.curlw = numpy.ascontiguousarray(self.curlw, dtype=numpy.float32)
if verbose > 0:
print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials
if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=False)
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=True)
else:
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
@@ -1651,12 +1655,12 @@ class PoissonSolver(object):
return full_conv[slices]
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=True):
"""
Batched OpenCL Poisson solver for two 2D fields.
Optimized: scatter H2D (~4x smaller transfer), fused multiply
(reads g_fft once), async D2H (one sync instead of two).
When profile=False, records and optionally prints per-stage timings.
When profile=True, records and optionally prints per-stage timings.
"""
if not OPENCL_AVAILABLE:
if verbose:
@@ -1734,7 +1738,7 @@ class PoissonSolver(object):
'src_plane': src_plane,
'dst_plane': dst_plane,
'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
'f_stage_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
'host_pad': numpy.zeros((2,) + full_shape, dtype=numpy.float32),
'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
@@ -1745,24 +1749,15 @@ class PoissonSolver(object):
fshape_c = c['fshape']
full_shape_c = c['full_shape']
# H2D: map pinned buffer, write f1/f2 directly (no f1_host/f2_host), unmap, then scatter
# H2D: write into pre-zeroed host pad buffer, single enqueue_copy (no map/unmap/scatter)
t0 = timeit.default_timer()
mapped_stage = c['f_stage_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.WRITE, is_blocking=True)
mapped_stage[0] = numpy.ascontiguousarray(f1, dtype=numpy.float32)
mapped_stage[1] = numpy.ascontiguousarray(f2, dtype=numpy.float32)
unmap_evt = mapped_stage.base.release(queue=queue)
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
scatter_batch(
queue, (fshape_c[1], fshape_c[0], 2), None,
c['f_stage_batch_dev'].data,
c['f_batch_dev'].data,
numpy.int32(fshape_c[1]),
numpy.int32(fshape_c[0]),
numpy.int32(full_shape_c[1]),
numpy.int32(c['src_plane']),
numpy.int32(c['dst_plane']),
wait_for=[unmap_evt]
)
host_pad = c['host_pad']
H, W = fshape_c
f1_f32 = numpy.ascontiguousarray(f1, dtype=numpy.float32)
f2_f32 = numpy.ascontiguousarray(f2, dtype=numpy.float32)
host_pad[0, :H, :W] = f1_f32
host_pad[1, :H, :W] = f2_f32
cl.enqueue_copy(queue, c['f_batch_dev'].data, host_pad, is_blocking=False)
if profile:
queue.finish()
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
@@ -1847,11 +1842,11 @@ class PoissonSolver(object):
)
return p1, p2
# def solve_opencl(self, f, verbose=False, profile=False):
# def solve_opencl(self, f, verbose=False, profile=True):
# """
# OpenCL Poisson solver using pyvkfft FFT convolution.
# Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids.
# When profile=False, records and optionally prints per-stage timings.
# When profile=True, records and optionally prints per-stage timings.
# """
# if not OPENCL_AVAILABLE:
# if verbose: