"""Speed of the GWN slicer (../../fcp_2810_gwn/scripts/gwn.py, defaults) on the GPU and on the CPU, by stage, against the package's encode on the same device. Writes ../report/gwn_speed.csv and ../report/gwn_speed.html. GPU: gwn.py as it is (CuPy kernels, direct sums). CPU: a twin made from gwn.py's own source with NumPy for CuPy, the surf slicer on its numba backend, and its summation kernels (wind, shepard, wind_grad, density) as numba loops of the same arithmetic (parallel over queries, float64 sums, float32 inputs). Every stage is timed synchronized; the summation kernels also count their work (queries x points). Clouds: clean area-uniform samples of the objects' meshes (seeded). env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/gwn_speed.py """ import csv import math import sys import time import types from pathlib import Path import numba as nb import numpy as np import trimesh sys.path.insert(0, str(Path(__file__).parent)) sys.path.insert(0, "/home/joao/code/fcp_devel/api") GWN_DIR = "/home/joao/code/data/meshes/fcp_2810_gwn/scripts" sys.path.insert(0, GWN_DIR) import cupy as cp # noqa: E402 import fcp # noqa: E402 import gwn as gwn_gpu # noqa: E402 from parts_report import CSS, OBJ, REPORT # noqa: E402 OBJECTS = ("mug", "dino", "open_box") SIZES = (2048, 10_000, 50_000) CPU_MAX = 10_000 # the CPU twin runs up to this many points (50k would take about 40 min per encode: estimated instead) REPEAT = 2 # GPU: timed runs after one warm-up, the minimum kept; CPU: one run (its kernels compiled once beforehand) OUT = REPORT / "gwn_speed.csv" # ------------------------------------------------------------------ numba twins of the summation kernels @nb.njit(parallel=True, fastmath=False, cache=True) def _wind_nb(X, V, E2, Q, out): m, q = len(X), len(Q) for t in nb.prange(q): qx, qy, qz = Q[t, 0], Q[t, 1], Q[t, 2] acc = 0.0 for j in range(m): dx, dy, dz = X[j, 0] - qx, X[j, 1] - qy, X[j, 2] - qz r2 = dx * dx + dy * dy + dz * dz + E2[j] inv = 1.0 / math.sqrt(r2) acc += (V[j, 0] * dx + V[j, 1] * dy + V[j, 2] * dz) * inv * inv * inv out[t] = acc * 0.07957747154594767 @nb.njit(parallel=True, fastmath=False, cache=True) def _shepard_nb(X, A, E2, Wv, Q, pw, out): m, q = len(X), len(Q) for t in nb.prange(q): qx, qy, qz = Q[t, 0], Q[t, 1], Q[t, 2] num, den = 0.0, 0.0 for j in range(m): dx, dy, dz = X[j, 0] - qx, X[j, 1] - qy, X[j, 2] - qz r2 = dx * dx + dy * dy + dz * dz + E2[j] g = A[j] / (r2 * r2) if pw == 4.0 else A[j] * r2 ** (-0.5 * pw) num += g * Wv[j] den += g out[t] = num / den @nb.njit(parallel=True, fastmath=False, cache=True) def _wind_grad_nb(X, V, E2, Q, out): m, q = len(X), len(Q) for t in nb.prange(q): qx, qy, qz = Q[t, 0], Q[t, 1], Q[t, 2] w = gx = gy = gz = 0.0 for j in range(m): dx, dy, dz = X[j, 0] - qx, X[j, 1] - qy, X[j, 2] - qz r2 = dx * dx + dy * dy + dz * dz + E2[j] inv = 1.0 / math.sqrt(r2) i3 = inv * inv * inv i5 = i3 * inv * inv vd = V[j, 0] * dx + V[j, 1] * dy + V[j, 2] * dz w += vd * i3 gx += -V[j, 0] * i3 + 3.0 * vd * dx * i5 gy += -V[j, 1] * i3 + 3.0 * vd * dy * i5 gz += -V[j, 2] * i3 + 3.0 * vd * dz * i5 c = 0.07957747154594767 out[t, 0], out[t, 1], out[t, 2], out[t, 3] = w * c, gx * c, gy * c, gz * c @nb.njit(parallel=True, fastmath=False, cache=True) def _density_nb(X, A, E2, Q, out): m, q = len(X), len(Q) for t in nb.prange(q): qx, qy, qz = Q[t, 0], Q[t, 1], Q[t, 2] acc = 0.0 for j in range(m): dx, dy, dz = X[j, 0] - qx, X[j, 1] - qy, X[j, 2] - qz d2 = dx * dx + dy * dy + dz * dz if d2 < 25.0 * E2[j]: acc += A[j] / (6.283185307 * E2[j]) * math.exp(-0.5 * d2 / E2[j]) out[t] = acc def _f32(a): return np.ascontiguousarray(np.asarray(a, np.float32)) def gpu_copy(tree=False): """A fresh copy of gwn.py (GPU), with the tree sums installed when tree=True (fastgwn.py).""" src = Path(GWN_DIR, "gwn.py").read_text() mod = types.ModuleType("gwn_gpu_copy") mod.__file__ = str(Path(GWN_DIR, "gwn.py")) exec(compile(src, mod.__file__, "exec"), mod.__dict__) if tree: import fastgwn fastgwn.install(mod, "cuda") return mod def cpu_twin(tree=False): """gwn.py's source with NumPy for CuPy, the CPU surf backend and the numba kernels: a module (with the tree sums installed when tree=True).""" src = Path(GWN_DIR, "gwn.py").read_text() for a, b in (("import cupy as cp", "import numpy as cp"), ('device="cuda"', 'device="cpu"'), ('_surf.backend("cuda")', '_surf.backend("cpu")')): assert a in src, a src = src.replace(a, b) mod = types.ModuleType("gwn_cpu") mod.__file__ = str(Path(GWN_DIR, "gwn.py")) exec(compile(src, mod.__file__, "exec"), mod.__dict__) def wind(X, V, E2, Q, chunk=None): out = np.empty(len(Q), np.float32) _wind_nb(_f32(X), _f32(V), _f32(E2), _f32(Q), out) return out def shepard(X, A, E2, Wv, Q, pw=4.0, chunk=None): out = np.empty(len(Q), np.float32) _shepard_nb(_f32(X), _f32(A), _f32(E2), _f32(Wv), _f32(Q), float(pw), out) return out def wind_grad(X, V, E2, Q, chunk=None): out = np.empty((len(Q), 4), np.float32) _wind_grad_nb(_f32(X), _f32(V), _f32(E2), _f32(Q), out) return out[:, 0], out[:, 1:] def density(X, A, E2, Q, chunk=None): out = np.empty(len(Q), np.float32) _density_nb(_f32(X), _f32(A), _f32(E2), _f32(Q), out) return out mod.wind, mod.shepard, mod.wind_grad, mod.density = wind, shepard, wind_grad, density if tree: import fastgwn fastgwn.install(mod, "cpu") return mod # ------------------------------------------------------------------ stage timing class Timer: """Wraps a module's kernels and stages: synchronized wall time and the work (queries x points) per name.""" def __init__(self, mod, gpu): self.mod, self.gpu, self.t, self.work = mod, gpu, {}, {} def sync(self): if self.gpu: cp.cuda.Device().synchronize() def wrap(self, owner, name, label, work=None): f = getattr(owner, name) def g(*a, **k): self.sync() t0 = time.perf_counter() r = f(*a, **k) self.sync() self.t[label] = self.t.get(label, 0.0) + time.perf_counter() - t0 if work is not None: self.work[label] = self.work.get(label, 0.0) + work(*a) return r setattr(owner, name, g) return f def run_gwn(mod, P, gpu): """One encode with the GWN slicer; per-stage times and work.""" T = Timer(mod, gpu) saved = {} # the kernels: wind's callers are the orientation, the points' own values and the field; split by caller below for nm, lab, wk in (("wind", "wind", lambda X, V, E2, Q, *r: len(Q) * len(X)), ("shepard", "shepard", lambda X, A, E2, Wv, Q, *r: len(Q) * len(X)), ("wind_grad", "wind_grad", lambda X, V, E2, Q, *r: len(Q) * len(X)), ("density", "density", lambda X, A, E2, Q, *r: len(Q) * len(X))): saved[nm] = T.wrap(mod, nm, lab, wk) saved["_orient"] = T.wrap(mod.GWN, "_orient", "orient (patch flips)") saved["__init__"] = T.wrap(mod.GWN, "__init__", "gwn init (eps, orient, splats, w at points)") saved["loops"] = T.wrap(mod.GWN, "loops", "field + contours (all layers)") saved["PointSlicer"] = T.wrap(mod, "PointSlicer", "surf (pose, k-NN, oriented directions)") try: T.sync() t0 = time.perf_counter() d = mod.encode_gwn(P, up=(0, 0, 1)) T.sync() total = time.perf_counter() - t0 finally: for nm in ("wind", "shepard", "wind_grad", "density", "PointSlicer"): setattr(mod, nm, saved[nm]) for nm in ("_orient", "__init__", "loops"): setattr(mod.GWN, nm, saved[nm]) return d, total, T.t, T.work def run_package(P, device): if device == "cuda": cp.cuda.Device().synchronize() t0 = time.perf_counter() d = fcp.encode(P, up=(0, 0, 1), device=device) if device == "cuda": cp.cuda.Device().synchronize() return d, time.perf_counter() - t0 def best(f, n=REPEAT, warm=True): if warm: f() # warm-up (kernel compilation, caches) runs = [f() for _ in range(n)] return min(runs, key=lambda r: r[1]) def compile_cpu(mod): """Compile the numba kernels once (tiny inputs) so the CPU runs need no warm-up.""" X = np.random.default_rng(0).random((8, 3)).astype(np.float32) E = np.ones(8, np.float32) mod.wind(X, X, E, X); mod.shepard(X, E, E, E, X); mod.wind_grad(X, X, E, X); mod.density(X, E, E, X) VARIANTS = (("cuda", "direct"), ("cuda", "tree"), ("cpu", "tree")) # the CPU direct sums: gwn_speed_direct.csv SIZES_BY = {("cuda", "direct"): (2048, 10_000, 50_000), ("cuda", "tree"): (2048, 10_000, 50_000, 100_000), ("cpu", "tree"): (2048, 10_000, 50_000, 100_000)} def main(): import fastgwn mods = {("cuda", "direct"): gpu_copy(False), ("cuda", "tree"): gpu_copy(True), ("cpu", "tree"): cpu_twin(True)} compile_cpu(cpu_twin(False)) # the numba direct kernels (the orientation's small sums stay direct) fastgwn.wind_tree(np.random.default_rng(0).random((64, 3)), np.ones((64, 3)), np.ones(64)) # numba tree build run_package(np.asarray(trimesh.load(OBJ["mug"], force="mesh").sample(2048, seed=5)), "cpu") # numba surf compiled rows = [] sizes = sorted({n for v in SIZES_BY.values() for n in v}) for name in OBJECTS: g = trimesh.load(OBJ[name], force="mesh") for n in sizes: P = np.asarray(g.sample(n, seed=5)) pk = {} for device, kind in VARIANTS: if n not in SIZES_BY[(device, kind)]: continue gpu = device == "cuda" mod = mods[(device, kind)] if device not in pk: pk[device] = best(lambda: run_package(P, device), n=REPEAT if gpu else 1, warm=gpu) dp, tp = pk[device] d, tg, st, wk = best(lambda: run_gwn(mod, P, gpu), n=REPEAT if gpu else 1, warm=True) inner = st.get("gwn init (eps, orient, splats, w at points)", 0.0) field = st.get("field + contours (all layers)", 0.0) rows.append(dict(object=name, points=n, device=device, sums=kind, package_s=round(tp, 3), gwn_s=round(tg, 3), surf_s=round(st.get("surf (pose, k-NN, oriented directions)", 0.0), 3), init_s=round(inner, 3), orient_s=round(st.get("orient (patch flips)", 0.0), 3), field_s=round(field, 3), wind_s=round(st.get("wind", 0.0), 3), shepard_s=round(st.get("shepard", 0.0), 3), parts_pkg=len(dp.parts), parts_gwn=len(d.parts))) print(rows[-1], flush=True) with OUT.open("w", newline="") as f: w = csv.DictWriter(f, rows[0].keys()) w.writeheader() w.writerows(rows) page(rows) def page(rows): keys = ("object", "points", "device", "sums", "package_s", "gwn_s", "surf_s", "init_s", "orient_s", "field_s", "wind_s", "shepard_s", "parts_pkg", "parts_gwn") head = ("objectpointsdeviceGWN sumspackage encode sGWN encode s" "surf sGWN init sof which orientation sfield + contours swind s" "shepard sparts packageparts GWN") old = "" if (REPORT / "gwn_speed_direct.csv").exists(): with (REPORT / "gwn_speed_direct.csv").open() as f: dr = [r for r in csv.DictReader(f) if r["device"] == "cpu"] old = ("

The CPU with direct sums (the first benchmark)

" "" + "".join( f"" f"" for r in dr) + "
objectpointspackage sGWN sfield s
{r['object']}{r['points']}{r['package_s']}{r['gwn_s']}{r['field_s']}

Direct sums on the CPU: 1.65 × 10⁹ point × " "query evaluations per second; 50k estimated at 14–25 min, 100k at 50–81 min.

") (REPORT / "gwn_speed.html").write_text(f""" GWN slicer speed

← report · synthesis

GWN slicer: speed on the GPU and the CPU, direct and tree sums

GPU: GTX 1650. "direct": gwn.py's CuPy kernels (a sum over all points per query). "tree": the tree codes of scripts/fastgwn.py (Barnes-Hut with a second-order dipole expansion for the winding number, θ = 0.5; monopole for the local iso-value, θ = 0.35; numba tree build; numba traversal on the CPU, float32 CUDA traversal on the GPU). CPU: gwn.py's twin (NumPy, the surf slicer's numba backend; {nb.config.NUMBA_NUM_THREADS} threads). GPU times: best of {REPEAT} runs after a warm-up; CPU: one run after compiling. Synchronized. The package's encode on the same device for comparison. Clean area-uniform clouds (seed 5), 128 layers, 1024 atoms. Generated {time.strftime('%Y-%m-%d %H:%M')} by scripts/gwn_speed.py.

{head}{''.join("" + "".join(f"" for k in keys) + "" for r in rows)}
{r[k]}
{old}

Accuracy of the tree sums: report/ablation.csv, configuration "gwn-tree" against "gwn" (the synthesis).

""") if __name__ == "__main__": main()