integrated finger process pipeline on remapped 2d image
This commit is contained in:
ykcheungab
@@ -0,0 +1,5 @@
|
||||
from .capacitance import CapacitanceFFTSolver
|
||||
from .masked_sparse import MaskedSparseSolver
|
||||
from .rectangular import RectangularFFTSolver
|
||||
|
||||
__all__ = ["RectangularFFTSolver", "MaskedSparseSolver", "CapacitanceFFTSolver"]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Shared solver protocol and helpers."""
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,113 @@
|
||||
import numpy
|
||||
from scipy import linalg as dense_linalg
|
||||
|
||||
from ..config import MaskedPoissonConfig
|
||||
from ..domain import MaskedDomain, compute_boundary_mask, sample_boundary_coords
|
||||
|
||||
|
||||
class CapacitanceFFTSolver:
|
||||
"""Masked Poisson via rectangular FFT bulk solve + capacitance boundary correction."""
|
||||
|
||||
def __init__(self, config: MaskedPoissonConfig, bulk_solver):
|
||||
self.config = config
|
||||
self.bulk_solver = bulk_solver
|
||||
self._capfft_cache = None
|
||||
self._capfft_domain = None
|
||||
|
||||
def prepare(self, domain: MaskedDomain):
|
||||
"""Precompute capacitance matrix for boundary correction.
|
||||
|
||||
domain: MaskedDomain with built Laplacian
|
||||
returns: capacitance cache dict (also stored internally)
|
||||
"""
|
||||
self._capfft_domain = domain
|
||||
self._capfft_cache = self._build_capfft_cache(domain)
|
||||
return self._capfft_cache
|
||||
|
||||
def _build_capfft_cache(self, domain: MaskedDomain):
|
||||
active = domain.mask
|
||||
boundary_coords = numpy.argwhere(compute_boundary_mask(active))
|
||||
boundary_coords = sample_boundary_coords(boundary_coords, self.config)
|
||||
if boundary_coords.size == 0:
|
||||
pinned_rc = numpy.array(domain.active_coords[domain.pinned], dtype=numpy.int32)
|
||||
boundary_coords = pinned_rc.reshape(1, 2)
|
||||
|
||||
boundary_active_indices = domain.index_map[boundary_coords[:, 0], boundary_coords[:, 1]]
|
||||
boundary_count = int(len(boundary_active_indices))
|
||||
|
||||
cap_matrix = numpy.zeros((boundary_count, boundary_count), dtype=numpy.float64)
|
||||
zero_rhs = numpy.zeros(domain.num_active, dtype=numpy.float64)
|
||||
for col, (row, col_idx) in enumerate(boundary_coords):
|
||||
charge = numpy.zeros(active.shape, dtype=numpy.float32)
|
||||
charge[row, col_idx] = 1.0
|
||||
response = numpy.asarray(self.bulk_solver.solve(charge), dtype=numpy.float64)
|
||||
residual = domain.residual(response, zero_rhs)
|
||||
cap_matrix[:, col] = residual[boundary_active_indices]
|
||||
|
||||
reg_scale = self.config.reg
|
||||
if reg_scale <= 0:
|
||||
reg_scale = 1e-8
|
||||
reg = reg_scale * max(1.0, float(numpy.linalg.norm(cap_matrix, ord=numpy.inf)))
|
||||
cap_matrix_reg = cap_matrix + reg * numpy.eye(boundary_count, dtype=numpy.float64)
|
||||
lu, piv = dense_linalg.lu_factor(cap_matrix_reg)
|
||||
|
||||
return {
|
||||
"boundary_coords": boundary_coords,
|
||||
"boundary_active_indices": boundary_active_indices,
|
||||
"boundary_count": boundary_count,
|
||||
"cap_matrix": cap_matrix,
|
||||
"cap_matrix_reg": cap_matrix_reg,
|
||||
"cap_lu": lu,
|
||||
"cap_piv": piv,
|
||||
"reg": reg,
|
||||
}
|
||||
|
||||
def get_cache(self, domain: MaskedDomain):
|
||||
"""Return cached capacitance data, building if needed."""
|
||||
if self._capfft_cache is None or self._capfft_domain is not domain:
|
||||
self.prepare(domain)
|
||||
return self._capfft_cache
|
||||
|
||||
def solve(self, f, domain: MaskedDomain):
|
||||
"""Solve masked Poisson with capacitance-corrected rectangular FFT.
|
||||
|
||||
f: (H, W) scalar source field
|
||||
domain: MaskedDomain with prebuilt Laplacian
|
||||
returns: (H, W) scalar potential; zero outside mask
|
||||
"""
|
||||
rhs_active = domain.prepare_rhs(f)
|
||||
if rhs_active.size == 0:
|
||||
return numpy.zeros_like(f, dtype=numpy.float32)
|
||||
|
||||
rhs_full = numpy.zeros_like(f, dtype=numpy.float32)
|
||||
rhs_full[domain.mask] = rhs_active.astype(numpy.float32)
|
||||
|
||||
field = numpy.asarray(self.bulk_solver.solve(rhs_full), dtype=numpy.float64)
|
||||
cap_cache = self.get_cache(domain)
|
||||
if not cap_cache or cap_cache["boundary_count"] == 0:
|
||||
field_active = field[domain.mask]
|
||||
if field_active.size:
|
||||
field_active = field_active.copy()
|
||||
field_active -= field_active[domain.pinned]
|
||||
return domain.scatter(field_active, f)
|
||||
|
||||
boundary_rows = cap_cache["boundary_active_indices"]
|
||||
rhs_scale = max(1.0, float(numpy.linalg.norm(rhs_active)))
|
||||
for _ in range(self.config.max_iters):
|
||||
residual = domain.residual(field, rhs_active)
|
||||
boundary_residual = residual[boundary_rows]
|
||||
if float(numpy.linalg.norm(boundary_residual)) <= self.config.tol * rhs_scale:
|
||||
break
|
||||
|
||||
alpha = dense_linalg.lu_solve((cap_cache["cap_lu"], cap_cache["cap_piv"]), boundary_residual)
|
||||
charges = numpy.zeros_like(f, dtype=numpy.float32)
|
||||
coords = cap_cache["boundary_coords"]
|
||||
charges[coords[:, 0], coords[:, 1]] = alpha.astype(numpy.float32)
|
||||
correction = numpy.asarray(self.bulk_solver.solve(charges), dtype=numpy.float64)
|
||||
field -= correction
|
||||
|
||||
field_active = field[domain.mask]
|
||||
if field_active.size:
|
||||
field_active = field_active.copy()
|
||||
field_active -= field_active[domain.pinned]
|
||||
return domain.scatter(field_active, f)
|
||||
@@ -0,0 +1,21 @@
|
||||
import numpy
|
||||
|
||||
from ..domain import MaskedDomain
|
||||
|
||||
|
||||
class MaskedSparseSolver:
|
||||
"""Reference masked Poisson solver via factorized sparse Laplacian."""
|
||||
|
||||
def solve(self, f, domain: MaskedDomain):
|
||||
"""Solve masked Poisson on active domain pixels.
|
||||
|
||||
f: (H, W) scalar source field
|
||||
domain: MaskedDomain with prebuilt Laplacian
|
||||
returns: (H, W) scalar potential; zero outside mask
|
||||
"""
|
||||
rhs = domain.prepare_rhs(f)
|
||||
if rhs.size == 0:
|
||||
return numpy.zeros_like(f, dtype=numpy.float32)
|
||||
|
||||
sol_active = domain.factor(rhs) # (N,) reduced-system potential
|
||||
return domain.scatter(sol_active, f)
|
||||
@@ -0,0 +1,530 @@
|
||||
import numpy
|
||||
import timeit
|
||||
from scipy import signal
|
||||
|
||||
from ...backend_availability import (
|
||||
CUPY_AVAILABLE,
|
||||
OPENCL_AVAILABLE,
|
||||
cp,
|
||||
cp_signal,
|
||||
cl,
|
||||
cla,
|
||||
cl_rfftn,
|
||||
cl_irfftn,
|
||||
)
|
||||
from ..timing import Timer
|
||||
|
||||
|
||||
class RectangularFFTSolver:
|
||||
"""FFT Poisson solver on a full structured rectangular grid."""
|
||||
|
||||
def __init__(self, grid, spacings, poisson_backend="cpu", dtype=numpy.float32, stype="F"):
|
||||
self.gdims = grid
|
||||
self.gdx = spacings
|
||||
self.dim = len(self.gdims)
|
||||
self.dtype = dtype
|
||||
self.stype = stype
|
||||
self.poisson_backend = poisson_backend
|
||||
self.ready = False
|
||||
self.G = None
|
||||
|
||||
if self.poisson_backend == "cupy":
|
||||
self._gpu_cache = None
|
||||
if self.poisson_backend == "opencl":
|
||||
self._opencl_cache = None
|
||||
self._opencl_queue = None
|
||||
self._opencl_context = None
|
||||
self._opencl_crop_kernel = None
|
||||
self._opencl_scatter_kernel = None
|
||||
self._opencl_fused_mul_kernel = None
|
||||
self._opencl_scatter_batch_kernel = None
|
||||
self._opencl_crop_batch_kernel = None
|
||||
self._opencl_batch_mul_kernel = None
|
||||
self._opencl_last_timing = None
|
||||
self._opencv_cache = None
|
||||
|
||||
def _create_radialGrid(self, verbose):
|
||||
mtimer = Timer()
|
||||
if verbose:
|
||||
print(" - creating distance kernel:", end="")
|
||||
|
||||
rdims = tuple(2 * d - 1 for d in self.gdims)
|
||||
R = numpy.indices(rdims, dtype=self.dtype)
|
||||
for d in range(self.dim):
|
||||
R[d] += 1 - self.gdims[d]
|
||||
R[d] *= self.gdx[d]
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.absolute(R[0], R[0])
|
||||
R = R[0]
|
||||
elif self.dim == 2:
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
zval = 0.5 * (self.gdx[0] + self.gdx[1]) / 2.0
|
||||
R[self.gdims[0] - 1, self.gdims[1] - 1] = zval
|
||||
elif self.dim == 3:
|
||||
numpy.hypot(R[2], R[1], R[1])
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
zval = 0.5 * (self.gdx[0] + self.gdx[1] + self.gdx[2]) / 3.0
|
||||
R[self.gdims[0] - 1, self.gdims[1] - 1, self.gdims[2] - 1] = zval
|
||||
|
||||
if verbose:
|
||||
print(R.shape, R.min(), R.max())
|
||||
mtimer.end()
|
||||
return R
|
||||
|
||||
def prepare(self, verbose=False):
|
||||
"""Precompute Green's function for FFT convolution.
|
||||
|
||||
returns: None (sets self.G and self.ready)
|
||||
"""
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print("\nInitializing rectangular FFT Poisson solver, type =", self.stype)
|
||||
|
||||
self.G = None
|
||||
if self.stype == "F":
|
||||
self.G = self._create_radialGrid(verbose)
|
||||
elif self.stype == "S":
|
||||
raise SyntaxError("Spatial solver not suggested for structured grids")
|
||||
|
||||
ltimer = Timer()
|
||||
if verbose:
|
||||
print(" - computing the Green's function:")
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.multiply(self.G, 0.5, self.G)
|
||||
elif self.dim == 2:
|
||||
numpy.log(self.G, self.G)
|
||||
numpy.multiply(self.G, (0.5 / numpy.pi), self.G)
|
||||
elif self.dim == 3:
|
||||
numpy.reciprocal(self.G, self.G)
|
||||
numpy.multiply(self.G, (-0.25 / numpy.pi), self.G)
|
||||
|
||||
if verbose:
|
||||
print(self.G.shape, self.G.min(), self.G.max(), end="")
|
||||
ltimer.end()
|
||||
|
||||
self.ready = True
|
||||
if verbose:
|
||||
print("Rectangular FFT Poisson solver initialized", end="")
|
||||
gtimer.end()
|
||||
|
||||
def solve(self, f, verbose=False):
|
||||
"""Solve Poisson on the full rectangular grid via FFT convolution.
|
||||
|
||||
f: (H, W) or (Z, Y, X) scalar source field
|
||||
returns: (H, W) scalar potential, same shape as f
|
||||
"""
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print("\nSolving rectangular Poisson Eq.")
|
||||
|
||||
if self.stype == "F":
|
||||
p = signal.fftconvolve(f, self.G, mode="same")
|
||||
numpy.multiply(p, numpy.prod(self.gdx), p)
|
||||
else:
|
||||
raise SyntaxError("Spatial solver not suggested for structured grids")
|
||||
|
||||
if verbose:
|
||||
print("Rectangular Poisson solver finished")
|
||||
gtimer.end()
|
||||
return p
|
||||
|
||||
def solve_gpu(self, f, verbose=False, return_gpu=False):
|
||||
"""GPU FFT Poisson solve using CuPy.
|
||||
|
||||
f: (H, W) scalar source field on CPU
|
||||
return_gpu: if True, return CuPy array instead of numpy
|
||||
returns: (H, W) scalar potential
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print("CuPy not available, falling back to CPU solver")
|
||||
return self.solve(f, verbose)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print("\nSolving Poisson Eq. on GPU")
|
||||
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
if self.stype == "F":
|
||||
if self._gpu_cache is None:
|
||||
if verbose:
|
||||
print(" - Transferring Green's function to GPU...")
|
||||
self._gpu_cache = {
|
||||
"G_gpu": cp.asarray(self.G, dtype=cp.float32),
|
||||
"gdx_prod": numpy.prod(self.gdx),
|
||||
}
|
||||
|
||||
f_gpu = cp.asarray(f, dtype=cp.float32)
|
||||
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache["G_gpu"], mode="same")
|
||||
p_gpu *= self._gpu_cache["gdx_prod"]
|
||||
|
||||
if verbose:
|
||||
print("Poisson solver (GPU) finished")
|
||||
gtimer.end()
|
||||
|
||||
if return_gpu:
|
||||
return p_gpu
|
||||
return cp.asnumpy(p_gpu)
|
||||
|
||||
raise SyntaxError("Spatial solver not suggested for structured grids")
|
||||
|
||||
def _solve_gpu_internal(self, f_gpu):
|
||||
"""Internal GPU solve; f_gpu and return value are CuPy arrays."""
|
||||
if self.stype == "F":
|
||||
if self._gpu_cache is None:
|
||||
self._gpu_cache = {
|
||||
"G_gpu": cp.asarray(self.G, dtype=cp.float32),
|
||||
"gdx_prod": numpy.prod(self.gdx),
|
||||
}
|
||||
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache["G_gpu"], mode="same")
|
||||
p_gpu *= self._gpu_cache["gdx_prod"]
|
||||
return p_gpu
|
||||
raise SyntaxError("Spatial solver not suggested for structured grids")
|
||||
|
||||
def clear_gpu_cache(self):
|
||||
"""Release CuPy FFT cache buffers."""
|
||||
if self._gpu_cache is not None:
|
||||
self._gpu_cache = None
|
||||
if CUPY_AVAILABLE:
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
|
||||
def _get_opencl_queue(self):
|
||||
if self._opencl_queue is not None:
|
||||
return self._opencl_queue
|
||||
|
||||
if not OPENCL_AVAILABLE:
|
||||
raise RuntimeError("OpenCL/pyvkfft is not available")
|
||||
|
||||
ctx = None
|
||||
for platform in cl.get_platforms():
|
||||
for device in platform.get_devices():
|
||||
if device.type & cl.device_type.GPU == 0:
|
||||
continue
|
||||
ctx = cl.Context(devices=(device,))
|
||||
break
|
||||
if ctx is not None:
|
||||
break
|
||||
|
||||
if ctx is None:
|
||||
raise RuntimeError("No OpenCL GPU device found")
|
||||
|
||||
self._opencl_queue = cl.CommandQueue(ctx)
|
||||
self._opencl_context = ctx
|
||||
return self._opencl_queue
|
||||
|
||||
def _get_opencl_scatter_batch_kernel(self, queue):
|
||||
if self._opencl_scatter_batch_kernel is not None:
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void scatter_batch(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int src_height,
|
||||
const int dst_width,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < src_width && gid_y < src_height) {
|
||||
int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
|
||||
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx];
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_scatter_batch_kernel = program.scatter_batch
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
def _get_opencl_crop_batch_kernel(self, queue):
|
||||
if self._opencl_crop_batch_kernel is not None:
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void crop_batch(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int dst_width,
|
||||
const int dst_height,
|
||||
const int start_x,
|
||||
const int start_y,
|
||||
const float scale,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < dst_width && gid_y < dst_height) {
|
||||
int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
|
||||
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx] * scale;
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_crop_batch_kernel = program.crop_batch
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
def _get_opencl_batch_mul_kernel(self, queue):
|
||||
if self._opencl_batch_mul_kernel is not None:
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void batch_complex_mul(
|
||||
__global float2* fft_batch,
|
||||
__global const float2* g,
|
||||
const int n
|
||||
) {
|
||||
const int b = get_global_id(0);
|
||||
const int base = b * n;
|
||||
const int i0 = get_global_id(1) * 4;
|
||||
const int i_end = min(i0 + 4, n);
|
||||
for (int i = i0; i < i_end; i++) {
|
||||
float2 gv = g[i];
|
||||
float2 a = fft_batch[base + i];
|
||||
fft_batch[base + i] = (float2)(fma(a.x, gv.x, -(a.y * gv.y)),
|
||||
fma(a.x, gv.y, a.y * gv.x));
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_batch_mul_kernel = program.batch_complex_mul
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
def _get_opencl_alloc_host_ptr(self):
|
||||
ctx = self._opencl_context
|
||||
mf = cl.mem_flags
|
||||
|
||||
def allocator(nbytes):
|
||||
return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)
|
||||
|
||||
return allocator
|
||||
|
||||
@staticmethod
|
||||
def _center_crop_same(full_conv, target_shape):
|
||||
starts = [(full_conv.shape[i] - target_shape[i]) // 2 for i in range(len(target_shape))]
|
||||
slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
|
||||
return full_conv[slices]
|
||||
|
||||
def solve_batch_opencl(self, f1, f2, verbose=False, profile=False):
|
||||
"""Batched OpenCL FFT Poisson solve for two 2D fields.
|
||||
|
||||
f1, f2: (H, W) scalar source fields, same shape
|
||||
returns: (p1, p2) each (H, W) scalar potential
|
||||
"""
|
||||
if not OPENCL_AVAILABLE:
|
||||
if verbose:
|
||||
print("OpenCL/pyvkfft not available, falling back to CPU solver")
|
||||
return self.solve(f1, verbose), self.solve(f2, verbose)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f1.shape
|
||||
if fshape != f2.shape:
|
||||
raise ValueError("solve_batch_opencl requires inputs with identical shapes")
|
||||
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
timing = {"cache_prepare_ms": 0.0}
|
||||
t_total = timeit.default_timer()
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
queue = self._get_opencl_queue()
|
||||
timing["queue_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
full_shape = tuple(2 * s - 1 for s in fshape)
|
||||
timing["host_prepare_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
cache_valid = (
|
||||
self._opencl_cache is not None
|
||||
and self._opencl_cache["fshape"] == fshape
|
||||
and self._opencl_cache["full_shape"] == full_shape
|
||||
and "f_batch_dev" in self._opencl_cache
|
||||
)
|
||||
|
||||
if not cache_valid:
|
||||
t0 = timeit.default_timer()
|
||||
g_host = numpy.asarray(self.G, dtype=numpy.float32)
|
||||
g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
||||
g_slices = tuple(slice(0, s) for s in g_host.shape)
|
||||
g_pad[g_slices] = g_host
|
||||
g_dev = cla.to_device(queue, g_pad)
|
||||
g_fft_dev = cl_rfftn(g_dev)
|
||||
if profile:
|
||||
queue.finish()
|
||||
|
||||
same_starts = [fshape[i] - 1 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:
|
||||
fft_n_elements *= d
|
||||
src_plane = fshape[0] * fshape[1]
|
||||
dst_plane = full_shape[0] * full_shape[1]
|
||||
alloc_pinned = self._get_opencl_alloc_host_ptr()
|
||||
|
||||
self._opencl_cache = {
|
||||
"fshape": fshape,
|
||||
"full_shape": full_shape,
|
||||
"fft_shape": fft_shape,
|
||||
"fft_n_elements": fft_n_elements,
|
||||
"g_fft_dev": g_fft_dev,
|
||||
"same_starts": same_starts,
|
||||
"gdx_prod": numpy.prod(self.gdx),
|
||||
"src_plane": src_plane,
|
||||
"dst_plane": dst_plane,
|
||||
"f_batch_dev": cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
|
||||
"host_stage": numpy.empty((2,) + fshape, dtype=numpy.float32),
|
||||
"f_stage_dev": cla.empty(queue, (2,) + fshape, 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),
|
||||
"p1_host": numpy.empty(fshape, dtype=numpy.float32),
|
||||
"p2_host": numpy.empty(fshape, dtype=numpy.float32),
|
||||
}
|
||||
timing["cache_prepare_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
c = self._opencl_cache
|
||||
fshape_c = c["fshape"]
|
||||
full_shape_c = c["full_shape"]
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
H, W = fshape_c
|
||||
f1_f32 = f1 if (f1.flags["C_CONTIGUOUS"] and f1.dtype == numpy.float32) else numpy.ascontiguousarray(f1, dtype=numpy.float32)
|
||||
f2_f32 = f2 if (f2.flags["C_CONTIGUOUS"] and f2.dtype == numpy.float32) else numpy.ascontiguousarray(f2, dtype=numpy.float32)
|
||||
host_stage = c["host_stage"]
|
||||
numpy.copyto(host_stage[0], f1_f32)
|
||||
numpy.copyto(host_stage[1], f2_f32)
|
||||
f_stage_dev = c["f_stage_dev"]
|
||||
cl.enqueue_copy(queue, f_stage_dev.data, host_stage, is_blocking=False)
|
||||
c["f_batch_dev"].fill(0, queue=queue)
|
||||
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
|
||||
scatter_batch(
|
||||
queue,
|
||||
(W, H, 2),
|
||||
None,
|
||||
f_stage_dev.data,
|
||||
c["f_batch_dev"].data,
|
||||
numpy.int32(W),
|
||||
numpy.int32(H),
|
||||
numpy.int32(full_shape_c[1]),
|
||||
numpy.int32(H * W),
|
||||
numpy.int32(c["dst_plane"]),
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing["h2d_set_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
fft_batch = cl_rfftn(c["f_batch_dev"], c["f_fft_batch_dev"], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing["fft_f_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
n_fft = c["fft_n_elements"]
|
||||
global_n = max(1, (n_fft + 3) // 4)
|
||||
local_n = min(256, global_n)
|
||||
global_n = ((global_n + local_n - 1) // local_n) * local_n
|
||||
batch_mul = self._get_opencl_batch_mul_kernel(queue)
|
||||
batch_mul(
|
||||
queue,
|
||||
(2, global_n),
|
||||
(1, local_n),
|
||||
fft_batch.data,
|
||||
c["g_fft_dev"].data,
|
||||
numpy.int32(n_fft),
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing["mul_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
cl_irfftn(fft_batch, c["conv_batch_dev"], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing["ifft_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
crop_batch = self._get_opencl_crop_batch_kernel(queue)
|
||||
same_starts = c["same_starts"]
|
||||
crop_batch(
|
||||
queue,
|
||||
(fshape_c[1], fshape_c[0], 2),
|
||||
None,
|
||||
c["conv_batch_dev"].data,
|
||||
c["p_batch_dev"].data,
|
||||
numpy.int32(full_shape_c[1]),
|
||||
numpy.int32(fshape_c[1]),
|
||||
numpy.int32(fshape_c[0]),
|
||||
numpy.int32(same_starts[1]),
|
||||
numpy.int32(same_starts[0]),
|
||||
numpy.float32(c["gdx_prod"]),
|
||||
numpy.int32(c["dst_plane"]),
|
||||
numpy.int32(c["src_plane"]),
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing["crop_scale_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
p1 = c["p1_host"]
|
||||
p2 = c["p2_host"]
|
||||
mapped_p = c["p_batch_dev"].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
|
||||
numpy.copyto(p1, mapped_p[0])
|
||||
numpy.copyto(p2, mapped_p[1])
|
||||
mapped_p.base.release(queue=queue)
|
||||
timing["d2h_get_ms"] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
timing["total_ms"] = (timeit.default_timer() - t_total) * 1000.0
|
||||
self._opencl_last_timing = timing
|
||||
|
||||
if verbose:
|
||||
print("Poisson solver (OpenCL FFT batch2) finished")
|
||||
if profile:
|
||||
print(
|
||||
"OpenCL batch2 timing [ms]: "
|
||||
f"total={timing['total_ms']:.3f}, "
|
||||
f"queue={timing['queue_ms']:.3f}, "
|
||||
f"cache={timing['cache_prepare_ms']:.3f}, "
|
||||
f"host={timing['host_prepare_ms']:.3f}, "
|
||||
f"h2d={timing['h2d_set_ms']:.3f}, "
|
||||
f"fft={timing['fft_f_ms']:.3f}, "
|
||||
f"mul={timing['mul_ms']:.3f}, "
|
||||
f"ifft={timing['ifft_ms']:.3f}, "
|
||||
f"d2h={timing['d2h_get_ms']:.3f}, "
|
||||
f"crop={timing['crop_scale_ms']:.3f}"
|
||||
)
|
||||
return p1, p2
|
||||
|
||||
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
|
||||
"""Alias for solve_batch_opencl (backward compatible name)."""
|
||||
return self.solve_batch_opencl(f1, f2, verbose=verbose, profile=profile)
|
||||
Reference in New Issue
Block a user