orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dq6a reach 18 ms per frame

Browse Source
This commit is contained in:
radxa-d6a
2026-04-03 12:52:48 +08:00
parent 424a4e37ce
commit 0a281c982c
4 changed files with 211 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/HHD.py
+127 -12
View File
@@ -26,12 +26,107 @@ import timeit
import sys
# import cv2
from scipy import signal, spatial
from numba import njit
from .backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
cp, cp_signal, _torque_kernel,
cl, cla, cl_rfftn, cl_irfftn
)
# -- Numba-accelerated finite-difference kernels for 2D HHD ------------------
# Matches numpy.gradient's second-order boundary stencil exactly.
@njit(cache=True)
def _divcurl_2d_nb(vfield, dy, dx, div_out, curl_out):
"""Fused divergence + curl of a 2D vector field via finite differences."""
H = vfield.shape[0]
W = vfield.shape[1]
idy2 = 1.0 / (2.0 * dy)
idx2 = 1.0 / (2.0 * dx)
idy_edge = 1.0 / (2.0 * dy)
idx_edge = 1.0 / (2.0 * dx)
for y in range(H):
for x in range(W):
# du/dx
if x == 0:
dudx = (-3.0*vfield[y,0,0] + 4.0*vfield[y,1,0] - vfield[y,2,0]) * idx_edge
elif x == W - 1:
dudx = (vfield[y,W-3,0] - 4.0*vfield[y,W-2,0] + 3.0*vfield[y,W-1,0]) * idx_edge
else:
dudx = (vfield[y,x+1,0] - vfield[y,x-1,0]) * idx2
# dv/dy
if y == 0:
dvdy = (-3.0*vfield[0,x,1] + 4.0*vfield[1,x,1] - vfield[2,x,1]) * idy_edge
elif y == H - 1:
dvdy = (vfield[H-3,x,1] - 4.0*vfield[H-2,x,1] + 3.0*vfield[H-1,x,1]) * idy_edge
else:
dvdy = (vfield[y+1,x,1] - vfield[y-1,x,1]) * idy2
# dv/dx
if x == 0:
dvdx = (-3.0*vfield[y,0,1] + 4.0*vfield[y,1,1] - vfield[y,2,1]) * idx_edge
elif x == W - 1:
dvdx = (vfield[y,W-3,1] - 4.0*vfield[y,W-2,1] + 3.0*vfield[y,W-1,1]) * idx_edge
else:
dvdx = (vfield[y,x+1,1] - vfield[y,x-1,1]) * idx2
# du/dy
if y == 0:
dudy = (-3.0*vfield[0,x,0] + 4.0*vfield[1,x,0] - vfield[2,x,0]) * idy_edge
elif y == H - 1:
dudy = (vfield[H-3,x,0] - 4.0*vfield[H-2,x,0] + 3.0*vfield[H-1,x,0]) * idy_edge
else:
dudy = (vfield[y+1,x,0] - vfield[y-1,x,0]) * idy2
div_out[y, x] = dudx + dvdy
curl_out[y, x] = dvdx - dudy
@njit(cache=True)
def _grad_rotgrad_h_2d_nb(nD, nRu, vfield, dy, dx, d_out, r_out, h_out):
"""Fused gradient(nD) + rotated_gradient(nRu) + h = vfield - d - r."""
H = nD.shape[0]
W = nD.shape[1]
idy2 = 1.0 / (2.0 * dy)
idx2 = 1.0 / (2.0 * dx)
idy_edge = 1.0 / (2.0 * dy)
idx_edge = 1.0 / (2.0 * dx)
for y in range(H):
for x in range(W):
# gradient(nD) → d = (ddx_nD, ddy_nD)
if x == 0:
ddx_nD = (-3.0*nD[y,0] + 4.0*nD[y,1] - nD[y,2]) * idx_edge
ddx_nRu = (-3.0*nRu[y,0] + 4.0*nRu[y,1] - nRu[y,2]) * idx_edge
elif x == W - 1:
ddx_nD = (nD[y,W-3] - 4.0*nD[y,W-2] + 3.0*nD[y,W-1]) * idx_edge
ddx_nRu = (nRu[y,W-3] - 4.0*nRu[y,W-2] + 3.0*nRu[y,W-1]) * idx_edge
else:
ddx_nD = (nD[y,x+1] - nD[y,x-1]) * idx2
ddx_nRu = (nRu[y,x+1] - nRu[y,x-1]) * idx2
if y == 0:
ddy_nD = (-3.0*nD[0,x] + 4.0*nD[1,x] - nD[2,x]) * idy_edge
ddy_nRu = (-3.0*nRu[0,x] + 4.0*nRu[1,x] - nRu[2,x]) * idy_edge
elif y == H - 1:
ddy_nD = (nD[H-3,x] - 4.0*nD[H-2,x] + 3.0*nD[H-1,x]) * idy_edge
ddy_nRu = (nRu[H-3,x] - 4.0*nRu[H-2,x] + 3.0*nRu[H-1,x]) * idy_edge
else:
ddy_nD = (nD[y+1,x] - nD[y-1,x]) * idy2
ddy_nRu = (nRu[y+1,x] - nRu[y-1,x]) * idy2
d0 = ddx_nD
d1 = ddy_nD
r0 = -ddy_nRu
r1 = ddx_nRu
d_out[y, x, 0] = d0
d_out[y, x, 1] = d1
r_out[y, x, 0] = r0
r_out[y, x, 1] = r1
h_out[y, x, 0] = vfield[y, x, 0] - d0 - r0
h_out[y, x, 1] = vfield[y, x, 1] - d1 - r1
# ---------------------------------------------------------------------------
class nHHD(object):
@@ -79,8 +174,23 @@ class nHHD(object):
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
self._force_cpu_cache = None
self._decompose_bufs = None
self.psolver.prepare(False)
def _ensure_decompose_bufs(self, shape2d):
"""Lazily allocate reusable buffers for the 2D decompose fast path."""
if self._decompose_bufs is not None and self._decompose_bufs['shape'] == shape2d:
return self._decompose_bufs
self._decompose_bufs = {
'shape': shape2d,
'div': numpy.empty(shape2d, dtype=numpy.float32),
'curl': numpy.empty(shape2d, dtype=numpy.float32),
'd': numpy.empty(shape2d + (2,), dtype=numpy.float32),
'r': numpy.empty(shape2d + (2,), dtype=numpy.float32),
'h': numpy.empty(shape2d + (2,), dtype=numpy.float32),
}
return self._decompose_bufs
def get_force(self, vfield, verbose=False):
if self.poisson_backend == 'cupy':
return self.get_force_cuda(vfield, verbose)
@@ -129,19 +239,19 @@ class nHHD(object):
print( '\n ==> nhhd.decompose', vfield.shape)
if self.dim == 2:
shape2d = (vfield.shape[0], vfield.shape[1])
bufs = self._ensure_decompose_bufs(shape2d)
dy, dx = self.mesh.dx[0], self.mesh.dx[1]
# 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)
vf32 = vfield if vfield.dtype == numpy.float32 else vfield.astype(numpy.float32)
_divcurl_2d_nb(vf32, dy, dx, bufs['div'], bufs['curl'])
self.div = bufs['div']
self.curlw = bufs['curl']
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=self.debug)
else:
@@ -151,14 +261,19 @@ class nHHD(object):
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
# compute fields as gradients of potentials
self.d = self.mesh.gradient(self.nD)
nD_f32 = self.nD if self.nD.dtype == numpy.float32 else self.nD.astype(numpy.float32)
nRu_f32 = self.nRu if self.nRu.dtype == numpy.float32 else self.nRu.astype(numpy.float32)
_grad_rotgrad_h_2d_nb(nD_f32, nRu_f32, vf32, dy, dx,
bufs['d'], bufs['r'], bufs['h'])
self.d = bufs['d']
self.r = bufs['r']
self.h = bufs['h']
if verbose > 0:
print( 'd =', self.d.shape, self.d.min(), self.d.max())
self.r = self.mesh.rotated_gradient(self.nRu)
if verbose > 0:
print( 'r =', self.r.shape, self.r.min(), self.r.max())
print( 'h =', self.h.shape, self.h.min(), self.h.max())
return
else:
# compute div curl