update doc

This commit is contained in:
ykcheungab
2026-03-06 15:05:24 +08:00
parent c95e810e87
commit 44b14cded1
20 changed files with 359 additions and 178 deletions
+19 -6
View File
@@ -91,12 +91,18 @@ class nHHD(object):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.h
elif self.poisson_backend == 'cupy':
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
if hasattr(self, 'h') and self.h is not None:
return self.h
return cp.asnumpy(self._h_gpu)
def get_normal(self):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.d
elif self.poisson_backend == 'cupy':
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
if hasattr(self, 'd') and self.d is not None:
return self.d
return cp.asnumpy(self._d_gpu)
def _solve_poisson_backend(self, field, verbose=False):
@@ -249,12 +255,11 @@ class nHHD(object):
print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
with self.stream1:
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
with self.stream2:
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
self.stream1.synchronize()
self.stream2.synchronize()
# Use a single stream for both solves. Running concurrent FFT-based
# solves on separate streams can cause occasional long-run jitter.
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
cp.cuda.get_current_stream().synchronize()
if verbose > 0:
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
@@ -282,6 +287,9 @@ class nHHD(object):
if verbose > 0:
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
# Ensure all default-stream work (gradient, etc.) is complete before use
cp.cuda.get_current_stream().synchronize()
# Store GPU arrays
self._d_gpu = d_gpu
self._h_gpu = h_gpu
@@ -467,6 +475,11 @@ class nHHD(object):
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
# Keep stable CPU copies for get_normal()/get_shear() so callers get this frame's
# result instead of reading GPU memory that can change (avoids CuPy-only jumps)
self.d = cp.asnumpy(self._d_gpu)
self.h = cp.asnumpy(self._h_gpu)
flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)