fix incorrect fft result from USE_HOST_PTR

This commit is contained in:
ykcheungab-dragonq6a
2026-02-17 19:49:26 +08:00
parent b84ca53bbc
commit 2479b1fa9a
+9 -43
View File
@@ -1655,13 +1655,13 @@ class PoissonSolver(object):
g_fft_dev = cl_rfftn(g_dev)
if profile:
queue.finish()
same_starts = [(full_shape[i] - f_host.shape[i]) // 2 for i in range(self.dim)]
linear_shape = tuple(2 * s - 1 for s in f_host.shape)
same_starts = [(linear_shape[i] - f_host.shape[i]) // 2 for i in range(self.dim)]
same_slices = tuple(
slice(same_starts[i], same_starts[i] + f_host.shape[i])
for i in range(self.dim)
)
fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
p_host = numpy.empty(f_host.shape, dtype=numpy.float32)
self._opencl_cache = {
'fshape': f_host.shape,
'full_shape': full_shape,
@@ -1676,21 +1676,9 @@ class PoissonSolver(object):
'f_stage_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64),
'conv_full_dev': cla.empty(queue, full_shape, dtype=numpy.float32),
'p_host': p_host,
'p_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
'p_host': numpy.empty(f_host.shape, dtype=numpy.float32),
}
# USE_HOST_PTR: crop kernel writes directly into host numpy array
try:
p_buf = cl.Buffer(queue.context,
cl.mem_flags.READ_WRITE | cl.mem_flags.USE_HOST_PTR,
hostbuf=p_host)
self._opencl_cache['p_dev'] = cla.Array(
queue, f_host.shape, dtype=numpy.float32, data=p_buf)
self._opencl_cache['_p_buf'] = p_buf
self._opencl_cache['_use_host_ptr'] = True
except Exception:
self._opencl_cache['p_dev'] = cla.empty(
queue, f_host.shape, dtype=numpy.float32)
self._opencl_cache['_use_host_ptr'] = False
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
# H2D: upload only input-sized data, scatter into padded device buffer
@@ -1762,10 +1750,6 @@ class PoissonSolver(object):
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
t0 = timeit.default_timer()
if self._opencl_cache.get('_use_host_ptr'):
queue.finish()
p = self._opencl_cache['p_host']
else:
p = p_dev.get(ary=self._opencl_cache['p_host'])
timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
else:
@@ -1870,7 +1854,8 @@ class PoissonSolver(object):
if profile:
queue.finish()
same_starts = [(full_shape[i] - f1_host.shape[i]) // 2 for i in range(self.dim)]
linear_shape = tuple(2 * s - 1 for s in f1_host.shape)
same_starts = [(linear_shape[i] - f1_host.shape[i]) // 2 for i in range(self.dim)]
fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
fft_n_elements = 1
for d in fft_shape:
@@ -1878,9 +1863,6 @@ class PoissonSolver(object):
src_plane = f1_host.shape[0] * f1_host.shape[1]
dst_plane = full_shape[0] * full_shape[1]
p_shape = (2,) + f1_host.shape
p_batch_host = numpy.empty(p_shape, dtype=numpy.float32)
self._opencl_cache = {
'fshape': f1_host.shape,
'full_shape': full_shape,
@@ -1896,23 +1878,10 @@ class PoissonSolver(object):
'f_stage_batch_dev': cla.empty(queue, (2,) + f1_host.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_host': p_batch_host,
'p_batch_dev': cla.empty(queue, (2,) + f1_host.shape, dtype=numpy.float32),
'p_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
'f_stage_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
}
# USE_HOST_PTR: crop kernel writes directly into host numpy array
# On unified memory (Qualcomm), this avoids the explicit D2H memcpy
try:
p_buf = cl.Buffer(queue.context,
cl.mem_flags.READ_WRITE | cl.mem_flags.USE_HOST_PTR,
hostbuf=p_batch_host)
self._opencl_cache['p_batch_dev'] = cla.Array(
queue, p_shape, dtype=numpy.float32, data=p_buf)
self._opencl_cache['_p_buf'] = p_buf
self._opencl_cache['_use_host_ptr'] = True
except Exception:
self._opencl_cache['p_batch_dev'] = cla.empty(
queue, p_shape, dtype=numpy.float32)
self._opencl_cache['_use_host_ptr'] = False
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
c = self._opencl_cache
@@ -1987,11 +1956,8 @@ class PoissonSolver(object):
queue.finish()
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
# D2H: zero-copy if USE_HOST_PTR, else explicit transfer
# D2H: single transfer for both results
t0 = timeit.default_timer()
if c.get('_use_host_ptr'):
queue.finish()
else:
cl.enqueue_copy(queue, c['p_batch_host'], c['p_batch_dev'].data, is_blocking=True)
p1 = c['p_batch_host'][0]
p2 = c['p_batch_host'][1]