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. Dispatch Poisson solve using the selected backend.
""" """
if self.poisson_backend == 'opencl': 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': if self.poisson_backend == 'cupy':
return self.psolver.solve_gpu(field, verbose) return self.psolver.solve_gpu(field, verbose)
@@ -126,13 +126,17 @@ class nHHD(object):
# compute div curl # compute div curl
(self.div, self.curlw) = self.mesh.divcurl(vfield) (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: if verbose > 0:
print( 'div =', self.div.shape, self.div.min(), self.div.max()) print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials # compute potentials
if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE): 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: else:
self.nD = self._solve_poisson_backend(self.div, verbose > 0) self.nD = self._solve_poisson_backend(self.div, verbose > 0)
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0) self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
@@ -1651,12 +1655,12 @@ class PoissonSolver(object):
return full_conv[slices] 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. Batched OpenCL Poisson solver for two 2D fields.
Optimized: scatter H2D (~4x smaller transfer), fused multiply Optimized: scatter H2D (~4x smaller transfer), fused multiply
(reads g_fft once), async D2H (one sync instead of two). (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 not OPENCL_AVAILABLE:
if verbose: if verbose:
@@ -1734,7 +1738,7 @@ class PoissonSolver(object):
'src_plane': src_plane, 'src_plane': src_plane,
'dst_plane': dst_plane, 'dst_plane': dst_plane,
'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32), '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), '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), '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), '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'] fshape_c = c['fshape']
full_shape_c = c['full_shape'] 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() t0 = timeit.default_timer()
mapped_stage = c['f_stage_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.WRITE, is_blocking=True) host_pad = c['host_pad']
mapped_stage[0] = numpy.ascontiguousarray(f1, dtype=numpy.float32) H, W = fshape_c
mapped_stage[1] = numpy.ascontiguousarray(f2, dtype=numpy.float32) f1_f32 = numpy.ascontiguousarray(f1, dtype=numpy.float32)
unmap_evt = mapped_stage.base.release(queue=queue) f2_f32 = numpy.ascontiguousarray(f2, dtype=numpy.float32)
scatter_batch = self._get_opencl_scatter_batch_kernel(queue) host_pad[0, :H, :W] = f1_f32
scatter_batch( host_pad[1, :H, :W] = f2_f32
queue, (fshape_c[1], fshape_c[0], 2), None, cl.enqueue_copy(queue, c['f_batch_dev'].data, host_pad, is_blocking=False)
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]
)
if profile: if profile:
queue.finish() queue.finish()
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0 timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
@@ -1847,11 +1842,11 @@ class PoissonSolver(object):
) )
return p1, p2 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. # OpenCL Poisson solver using pyvkfft FFT convolution.
# Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids. # 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 not OPENCL_AVAILABLE:
# if verbose: # if verbose: