diff --git a/README.md b/README.md index 8f4eed99a..37dc7cfe5 100644 --- a/README.md +++ b/README.md @@ -100,7 +100,7 @@ Each benchmark falls into a single category. While such classification is not ac extrema, fft, lombscargle, sosfil, zmddft ### Simulation - ace, adv, amgmk, axhelm, bh, bspline-vgh, burger, cooling, ccsd-trpdrv, che, chemv, chi2, clenergy, cmp, cobahh, d2q9_bgk, d3q19_bgk, damage, ddbp, dslash, easyWave, eikonal, fdtd3d, feynman-kac, fhd, fluidSim, gibbs, goulash, gpp, grrt, haccmk, halo-finder, heartwall, heat, heat2d, henry, hexicton, hotspot, hotspot3D, hpl, hwt1d, hypterm, ising, iso2dfd, laplace, laplace3d, lavaMD, lid-driven-cavity, logic-resim, logic-rewrite, loopback, lsqt, lulesh, mcmd, md, mdh, metropolis, miniFE, minimod, minisweep, miniWeather, multimaterial, mxfp4, myocte, nbody, particle-diffusion, particlefilter, particles, pathfinder, pns, projectile, pso, qem, rainflow, rayleighBenardConvection, reaction, rsbench, rtm8, rushlarsen, s3d, su3, sundials, sheath, simplemoc, slit, sparkler, sph, sw4ck, tensorT, testSNAP, tissue, tpacf, tqs, tridiagonal, tsa, vanGenuchten, vmc, wenofv, wlcpow, wsm5, xlqc, xsbench + ace, adv, amgmk, axhelm, bh, bspline-vgh, burger, cooling, ccsd-trpdrv, che, chemv, chi2, clenergy, cmp, cobahh, d2q9_bgk, d3q19_bgk, damage, ddbp, dslash, easyWave, eikonal, fdtd3d, feynman-kac, fhd, fluidSim, gibbs, goulash, gpp, grrt, gsplat4d, haccmk, halo-finder, heartwall, heat, heat2d, henry, hexicton, hotspot, hotspot3D, hpl, hwt1d, hypterm, ising, iso2dfd, laplace, laplace3d, lavaMD, lid-driven-cavity, logic-resim, logic-rewrite, loopback, lsqt, lulesh, mcmd, md, mdh, metropolis, miniFE, minimod, minisweep, miniWeather, multimaterial, mxfp4, myocte, nbody, particle-diffusion, particlefilter, particles, pathfinder, pns, projectile, pso, qem, rainflow, rayleighBenardConvection, reaction, rsbench, rtm8, rushlarsen, s3d, su3, sundials, sheath, simplemoc, slit, sparkler, sph, sw4ck, tensorT, testSNAP, tissue, tpacf, tqs, tridiagonal, tsa, vanGenuchten, vmc, wenofv, wlcpow, wsm5, xlqc, xsbench ### Sorting bitonic-sort, hybridsort, is, merge, quicksort, radixsort, segsort, sort, sortKV, split, topk, warpsort @@ -848,6 +848,9 @@ Early results are shown [here](results/README.md) ### gru2 (cuda) Forward operations of a multi-layer gated recurrent unit (https://pytorch.org/) +### gsplat4d (cuda) + Physics-Aware 4D Dynamic Endoscopic Scene Simulations via MLLM-Guided Gaussian Splatting (https://arxiv.org/abs/2605.16022) + ### haccmk (c) The HACC microkernel (https://asc.llnl.gov/CORAL-benchmarks/#haccmk) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 374c7e71a..4e0f2b167 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -184,6 +184,7 @@ set(HECBENCH_POC_BENCHMARKS grrt gru gru2 + gsplat4d haccmk hausdorff haversine diff --git a/src/gsplat4d-cuda/CMakeLists.txt b/src/gsplat4d-cuda/CMakeLists.txt new file mode 100644 index 000000000..52b4506e5 --- /dev/null +++ b/src/gsplat4d-cuda/CMakeLists.txt @@ -0,0 +1,11 @@ +# gsplat4d-cuda/CMakeLists.txt + +add_hecbench_benchmark( + NAME gsplat4d + MODEL cuda + SOURCES main.cu + CATEGORIES simulation + TEST_ARGS 200000 1280 720 100 + TEST_REGEX "(?s)MPM step: PASS.*4D preprocess \([0-9]+ of [0-9]+ gaussians visible\): PASS.*Rasterizer: PASS" + TEST_TIMEOUT 300 +) diff --git a/src/gsplat4d-cuda/Makefile b/src/gsplat4d-cuda/Makefile new file mode 100644 index 000000000..46ab3d4fc --- /dev/null +++ b/src/gsplat4d-cuda/Makefile @@ -0,0 +1,56 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = nvcc +OPTIMIZE = yes +DEBUG = no +ARCH = sm_60 +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cu + +obj = $(source:.cu=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Xcompiler -Wall -arch=$(ARCH) + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cu reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-cuda/main.cu b/src/gsplat4d-cuda/main.cu new file mode 100644 index 000000000..69e5e94ce --- /dev/null +++ b/src/gsplat4d-cuda/main.cu @@ -0,0 +1,905 @@ +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#define CHECK(call) \ + do { \ + const cudaError_t err = (call); \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA error %s:%d '%s': %s\n", __FILE__, __LINE__, \ + #call, cudaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +#ifndef P2G_BLOCK +#define P2G_BLOCK 128 +#endif +#ifndef GRID_BLOCK +#define GRID_BLOCK 256 +#endif +#ifndef PREPROCESS_BLOCK +#define PREPROCESS_BLOCK 128 +#endif + +// --------------------------------------------------------------------------- +// stage 1: MLS-MPM +// +// The per particle 3x3 tensors are stored component major, so that the nine +// loads of a warp are nine fully coalesced 128 byte transactions instead of +// nine strided ones. The particles are pre-sorted by cell (see +// generate_scene), which keeps the scattered atomics of a warp inside a +// handful of cache lines. +// --------------------------------------------------------------------------- + +__device__ __forceinline__ void quad_weights(float fx, float& w0, float& w1, + float& w2) +{ + const float a = 1.5f - fx; + const float b = fx - 1.0f; + const float c = fx - 0.5f; + w0 = 0.5f * a * a; + w1 = 0.75f - b * b; + w2 = 0.5f * c * c; +} + +// One thread block scatters one chunk of particles, all of which belong to the +// same MPM_BLOCK^3 cell block. Their entire stencil footprint therefore fits a +// MPM_TILE^3 tile held in shared memory, so the scatter costs shared atomics +// plus one flush of the tile, rather than 108 global atomics per particle. +// A particle that has drifted out of its block since the binning still lands +// correctly through the global fallback. +__global__ void __launch_bounds__(P2G_BLOCK) +mpm_p2g_kernel(int n, MpmParams p, + const int* __restrict__ chunk_start, + const int* __restrict__ chunk_block, + const float4* __restrict__ mean, + const float4* __restrict__ velocity, + const float* __restrict__ affine_in, + const float* __restrict__ defgrad, + float* __restrict__ grid) +{ + __shared__ float tile[MPM_TILE_CELLS * 4]; + + for (int k = threadIdx.x; k < MPM_TILE_CELLS * 4; k += P2G_BLOCK) tile[k] = 0.0f; + + const int chunk = blockIdx.x; + const int begin = chunk_start[chunk]; + const int end = chunk_start[chunk + 1]; + + const int gb = chunk_block[chunk]; + const int obz = (gb % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int oby = ((gb / MPM_BLOCKS_PER_DIM) % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int obx = (gb / (MPM_BLOCKS_PER_DIM * MPM_BLOCKS_PER_DIM)) * MPM_BLOCK - 1; + + __syncthreads(); + + for (int i = begin + threadIdx.x; i < end; i += P2G_BLOCK) { + const float4 x = mean[i]; + const float4 v = velocity[i]; + + float C[9], F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) C[k] = affine_in[(size_t)k * n + i]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + // F <- (I + dt C) F + float F[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = fmaf(p.dt * C[a * 3 + k], F0[k * 3 + b], acc); + F[a * 3 + b] = acc; + } + + // neo-Hookean stress, folded into the MLS-MPM affine momentum + const float det = + F[0] * (F[4] * F[8] - F[5] * F[7]) - + F[1] * (F[3] * F[8] - F[5] * F[6]) + + F[2] * (F[3] * F[7] - F[4] * F[6]); + const float safe_J = (fabsf(det) < 1e-6f) ? ((det < 0.0f) ? -1e-6f : 1e-6f) : det; + const float inv_J = 1.0f / safe_J; + + const float cof[9] = { + (F[4] * F[8] - F[5] * F[7]), -(F[3] * F[8] - F[5] * F[6]), (F[3] * F[7] - F[4] * F[6]), + -(F[1] * F[8] - F[2] * F[7]), (F[0] * F[8] - F[2] * F[6]), -(F[0] * F[7] - F[1] * F[6]), + (F[1] * F[5] - F[2] * F[4]), -(F[0] * F[5] - F[2] * F[3]), (F[0] * F[4] - F[1] * F[3]) }; + + const float coeff = (p.lambda * __logf(fabsf(safe_J)) - p.mu) * inv_J; + float P[9]; + #pragma unroll + for (int k = 0; k < 9; k++) P[k] = fmaf(coeff, cof[k], p.mu * F[k]); + + const float s = -p.dt * p.particle_volume * 4.0f * p.inv_dx * p.inv_dx; + float affine[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 3; k++) acc = fmaf(P[a * 3 + k], F[b * 3 + k], acc); + affine[a * 3 + b] = fmaf(s, acc, p.particle_mass * C[a * 3 + b]); + } + + // F is only used for the stress above; it is deliberately not persisted + // here. The single per-step deformation-gradient update is applied once, in + // mpm_g2p, from the original F0 (this matches the host reference; writing F + // back here would advance the gradient twice per step). + + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + const float mv[3] = { p.particle_mass * v.x, p.particle_mass * v.y, + p.particle_mass * v.z }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float dpz = ((float)c - fz) * p.dx; + const float w = wxy * wz[c]; + + const int lx = ix - obx, ly = iy - oby, lz = iz - obz; + const bool local = (unsigned)lx < MPM_TILE && (unsigned)ly < MPM_TILE && + (unsigned)lz < MPM_TILE; + float* cell = local + ? tile + 4 * ((lx * MPM_TILE + ly) * MPM_TILE + lz) + : grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + + #pragma unroll + for (int k = 0; k < 3; k++) { + float impulse = affine[k * 3 + 0] * dpx; + impulse = fmaf(affine[k * 3 + 1], dpy, impulse); + impulse = fmaf(affine[k * 3 + 2], dpz, impulse); + atomicAdd(cell + k, w * (mv[k] + impulse)); + } + atomicAdd(cell + 3, w * p.particle_mass); + } + } + } + } + + __syncthreads(); + + // flush the tile into the global grid + for (int c = threadIdx.x; c < MPM_TILE_CELLS; c += P2G_BLOCK) { + const int lz = c % MPM_TILE; + const int ly = (c / MPM_TILE) % MPM_TILE; + const int lx = c / (MPM_TILE * MPM_TILE); + const int ix = obx + lx, iy = oby + ly, iz = obz + lz; + if ((unsigned)ix >= MPM_GRID || (unsigned)iy >= MPM_GRID || + (unsigned)iz >= MPM_GRID) + continue; + + const float4 acc = *(const float4*)(tile + 4 * c); + if (acc.x == 0.0f && acc.y == 0.0f && acc.z == 0.0f && acc.w == 0.0f) continue; + + float* dst = grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + atomicAdd(dst + 0, acc.x); + atomicAdd(dst + 1, acc.y); + atomicAdd(dst + 2, acc.z); + atomicAdd(dst + 3, acc.w); + } +} + +__global__ void __launch_bounds__(GRID_BLOCK) +mpm_grid_kernel(MpmParams p, float4* __restrict__ grid) +{ + const int cell = blockIdx.x * blockDim.x + threadIdx.x; + if (cell >= MPM_CELLS) return; + + float4 g = grid[cell]; + if (g.w <= 0.0f) { + grid[cell] = make_float4(0.0f, 0.0f, 0.0f, g.w); + return; + } + + const float inv_mass = 1.0f / g.w; + float vx = g.x * inv_mass; + float vy = fmaf(p.dt, p.gravity, g.y * inv_mass); + float vz = g.z * inv_mass; + + const int iz = cell % MPM_GRID; + const int iy = (cell / MPM_GRID) % MPM_GRID; + const int ix = cell / (MPM_GRID * MPM_GRID); + + if (ix < p.boundary && vx < 0.0f) vx = 0.0f; + if (ix >= MPM_GRID - p.boundary && vx > 0.0f) vx = 0.0f; + if (iy < p.boundary && vy < 0.0f) vy = 0.0f; + if (iy >= MPM_GRID - p.boundary && vy > 0.0f) vy = 0.0f; + if (iz < p.boundary && vz < 0.0f) vz = 0.0f; + if (iz >= MPM_GRID - p.boundary && vz > 0.0f) vz = 0.0f; + + grid[cell] = make_float4(vx, vy, vz, g.w); +} + +__global__ void __launch_bounds__(P2G_BLOCK) +mpm_g2p_kernel(int n, MpmParams p, + float4* __restrict__ mean, + float4* __restrict__ velocity, + float* __restrict__ affine_out, + float* __restrict__ defgrad, + const float4* __restrict__ grid) +{ + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + const float4 x = mean[i]; + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + float nv[3] = { 0.0f, 0.0f, 0.0f }; + float nC[9] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float w = wxy * wz[c]; + const float4 g = grid[(size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz]; + const float gv[3] = { g.x, g.y, g.z }; + const float dpos[3] = { dpx, dpy, ((float)c - fz) * p.dx }; + + #pragma unroll + for (int k = 0; k < 3; k++) { + nv[k] = fmaf(w, gv[k], nv[k]); + const float wg = 4.0f * p.inv_dx * p.inv_dx * w * gv[k]; + #pragma unroll + for (int l = 0; l < 3; l++) nC[k * 3 + l] = fmaf(wg, dpos[l], nC[k * 3 + l]); + } + } + } + } + + float F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = fmaf(p.dt * nC[a * 3 + k], F0[k * 3 + b], acc); + defgrad[(size_t)(a * 3 + b) * n + i] = acc; + } + + #pragma unroll + for (int k = 0; k < 9; k++) affine_out[(size_t)k * n + i] = nC[k]; + + velocity[i] = make_float4(nv[0], nv[1], nv[2], 0.0f); + mean[i] = make_float4(fmaf(p.dt, nv[0], x.x), fmaf(p.dt, nv[1], x.y), + fmaf(p.dt, nv[2], x.z), x.w); +} + +// --------------------------------------------------------------------------- +// stage 2: condition the 4D Gaussian on t, project, shade +// +// The frustum and opacity tests come before the 48 spherical harmonic loads, +// so a culled Gaussian costs no bandwidth beyond its pose. +// --------------------------------------------------------------------------- + +__global__ void __launch_bounds__(PREPROCESS_BLOCK) +preprocess_kernel(int n, float time, Camera cam, + const float4* __restrict__ mean4, + const float4* __restrict__ scale4, + const float4* __restrict__ quat_l, + const float4* __restrict__ quat_r, + const float* __restrict__ opacity_in, + const float* __restrict__ sh, + float2* __restrict__ mean2d, + float4* __restrict__ conic_opacity, + float4* __restrict__ color_depth, + int* __restrict__ radii) +{ + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + radii[i] = 0; + + const float4 mu = mean4[i]; + const float4 sc = scale4[i]; + const float4 ql = quat_l[i]; + const float4 qr = quat_r[i]; + + // M = L(ql) R(qr), the 4D rotation as a pair of isoclinic rotations + const float lw = ql.x, lx = ql.y, ly = ql.z, lz = ql.w; + const float rw = qr.x, rx = qr.y, ry = qr.z, rz = qr.w; + + const float L[16] = { + lw, -lx, -ly, -lz, + lx, lw, -lz, ly, + ly, lz, lw, -lx, + lz, -ly, lx, lw }; + const float R[16] = { + rw, -rx, -ry, -rz, + rx, rw, rz, -ry, + ry, -rz, rw, rx, + rz, ry, -rx, rw }; + + float M[16]; + #pragma unroll + for (int a = 0; a < 4; a++) + #pragma unroll + for (int b = 0; b < 4; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) acc = fmaf(L[a * 4 + k], R[k * 4 + b], acc); + M[a * 4 + b] = acc; + } + + // the quaternion basis is ordered (t, x, y, z) + const float s2[4] = { sc.w * sc.w, sc.x * sc.x, sc.y * sc.y, sc.z * sc.z }; + + // only the 3x3 spatial block, the spatio-temporal column and Sigma_tt are + // needed, so the remaining entries of the 4x4 product are never formed + float st[3], sxyz[6]; + #pragma unroll + for (int a = 0; a < 3; a++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) acc = fmaf(M[(a + 1) * 4 + k] * s2[k], M[k], acc); + st[a] = acc; + } + float sigma_tt = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) sigma_tt = fmaf(M[k] * s2[k], M[k], sigma_tt); + + int idx = 0; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = a; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) + acc = fmaf(M[(a + 1) * 4 + k] * s2[k], M[(b + 1) * 4 + k], acc); + sxyz[idx++] = acc; + } + + const float inv_tt = 1.0f / sigma_tt; + const float dt = time - mu.w; + + const float alpha = opacity_in[i] * __expf(-0.5f * dt * dt * inv_tt); + if (alpha < MIN_OPACITY) return; + + const float mean3[3] = { fmaf(dt * inv_tt, st[0], mu.x), + fmaf(dt * inv_tt, st[1], mu.y), + fmaf(dt * inv_tt, st[2], mu.z) }; + const float cov3[6] = { sxyz[0] - st[0] * st[0] * inv_tt, + sxyz[1] - st[0] * st[1] * inv_tt, + sxyz[2] - st[0] * st[2] * inv_tt, + sxyz[3] - st[1] * st[1] * inv_tt, + sxyz[4] - st[1] * st[2] * inv_tt, + sxyz[5] - st[2] * st[2] * inv_tt }; + + // world -> camera + const float px = mean3[0] - cam.cam_pos[0]; + const float py = mean3[1] - cam.cam_pos[1]; + const float pz = mean3[2] - cam.cam_pos[2]; + float t[3]; + #pragma unroll + for (int a = 0; a < 3; a++) + t[a] = fmaf(cam.view[a * 3 + 0], px, + fmaf(cam.view[a * 3 + 1], py, cam.view[a * 3 + 2] * pz)); + + if (t[2] < cam.near_plane) return; + + const float inv_z = 1.0f / t[2]; + const float inv_z2 = inv_z * inv_z; + + const float lim_x = 1.3f * (0.5f * cam.width) / cam.focal_x; + const float lim_y = 1.3f * (0.5f * cam.height) / cam.focal_y; + const float tx = t[2] * fminf(lim_x, fmaxf(-lim_x, t[0] * inv_z)); + const float ty = t[2] * fminf(lim_y, fmaxf(-lim_y, t[1] * inv_z)); + + const float J[6] = { cam.focal_x * inv_z, 0.0f, -cam.focal_x * tx * inv_z2, + 0.0f, cam.focal_y * inv_z, -cam.focal_y * ty * inv_z2 }; + + float T[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + T[a * 3 + b] = fmaf(J[a * 3 + 0], cam.view[0 * 3 + b], + fmaf(J[a * 3 + 1], cam.view[1 * 3 + b], + J[a * 3 + 2] * cam.view[2 * 3 + b])); + + const float S[9] = { cov3[0], cov3[1], cov3[2], + cov3[1], cov3[3], cov3[4], + cov3[2], cov3[4], cov3[5] }; + float TS[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + TS[a * 3 + b] = fmaf(T[a * 3 + 0], S[0 * 3 + b], + fmaf(T[a * 3 + 1], S[1 * 3 + b], + T[a * 3 + 2] * S[2 * 3 + b])); + + float ca = 0.3f, cb = 0.0f, cc = 0.3f; + #pragma unroll + for (int k = 0; k < 3; k++) { + ca = fmaf(TS[0 * 3 + k], T[0 * 3 + k], ca); + cb = fmaf(TS[0 * 3 + k], T[1 * 3 + k], cb); + cc = fmaf(TS[1 * 3 + k], T[1 * 3 + k], cc); + } + + const float det = fmaf(ca, cc, -cb * cb); + if (det <= 0.0f) return; + const float inv_det = 1.0f / det; + + const float mid = 0.5f * (ca + cc); + const float disc = sqrtf(fmaxf(0.1f, fmaf(mid, mid, -det))); + const float radius = ceilf(3.0f * sqrtf(fmaxf(mid + disc, mid - disc))); + + const float mx = fmaf(cam.focal_x * inv_z, t[0], cam.center_x); + const float my = fmaf(cam.focal_y * inv_z, t[1], cam.center_y); + + if (mx + radius < 0.0f || mx - radius > (float)cam.width || + my + radius < 0.0f || my - radius > (float)cam.height) + return; + + // spherical harmonics, evaluated only for the Gaussians that survive + const float inv_len = rsqrtf(fmaf(px, px, fmaf(py, py, pz * pz))); + const float dx = px * inv_len, dy = py * inv_len, dz = pz * inv_len; + + float basis[SH_COEFFS]; + basis[0] = SH_C0; + basis[1] = -SH_C1 * dy; + basis[2] = SH_C1 * dz; + basis[3] = -SH_C1 * dx; + + const float xx = dx * dx, yy = dy * dy, zz = dz * dz; + const float xy = dx * dy, yz = dy * dz, xz = dx * dz; + basis[4] = SH_C2_0 * xy; + basis[5] = SH_C2_1 * yz; + basis[6] = SH_C2_2 * (2.0f * zz - xx - yy); + basis[7] = SH_C2_3 * xz; + basis[8] = SH_C2_4 * (xx - yy); + basis[9] = SH_C3_0 * dy * (3.0f * xx - yy); + basis[10] = SH_C3_1 * xy * dz; + basis[11] = SH_C3_2 * dy * (4.0f * zz - xx - yy); + basis[12] = SH_C3_3 * dz * (2.0f * zz - 3.0f * xx - 3.0f * yy); + basis[13] = SH_C3_4 * dx * (4.0f * zz - xx - yy); + basis[14] = SH_C3_5 * dz * (xx - yy); + basis[15] = SH_C3_6 * dx * (xx - 3.0f * yy); + + float rgb[3] = { 0.0f, 0.0f, 0.0f }; + #pragma unroll + for (int c = 0; c < SH_COEFFS; c++) { + const float w = basis[c]; + #pragma unroll + for (int ch = 0; ch < 3; ch++) + rgb[ch] = fmaf(w, sh[((size_t)c * 3 + ch) * n + i], rgb[ch]); + } + + radii[i] = (int)radius; + mean2d[i] = make_float2(mx, my); + conic_opacity[i] = make_float4(cc * inv_det, -cb * inv_det, ca * inv_det, alpha); + color_depth[i] = make_float4(fmaxf(rgb[0] + 0.5f, 0.0f), + fmaxf(rgb[1] + 0.5f, 0.0f), + fmaxf(rgb[2] + 0.5f, 0.0f), t[2]); +} + +// --------------------------------------------------------------------------- +// stage 3: tile rasterizer +// +// One block owns one 16x16 tile. The block cooperatively stages a batch of +// BLOCK_SIZE Gaussians in shared memory, so each Gaussian is read from global +// memory once per tile rather than once per pixel, and the whole block leaves +// as soon as all of its pixels are saturated. +// --------------------------------------------------------------------------- + +__global__ void __launch_bounds__(BLOCK_SIZE) +render_kernel(Camera cam, + const float2* __restrict__ mean2d, + const float4* __restrict__ conic_opacity, + const float4* __restrict__ color_depth, + const int* __restrict__ tile_offsets, + const int* __restrict__ tile_list, + float4* __restrict__ image) +{ + __shared__ float2 s_xy[BLOCK_SIZE]; + __shared__ float4 s_co[BLOCK_SIZE]; + __shared__ float4 s_color[BLOCK_SIZE]; + + const int tile = blockIdx.y * cam.tiles_x + blockIdx.x; + const int x = blockIdx.x * BLOCK_X + threadIdx.x; + const int y = blockIdx.y * BLOCK_Y + threadIdx.y; + const int lane = threadIdx.y * BLOCK_X + threadIdx.x; + + const bool inside = (x < cam.width) && (y < cam.height); + bool done = !inside; + + const float pixf_x = (float)x + 0.5f; + const float pixf_y = (float)y + 0.5f; + + const int begin = tile_offsets[tile]; + const int end = tile_offsets[tile + 1]; + const int rounds = (end - begin + BLOCK_SIZE - 1) / BLOCK_SIZE; + + float transmittance = 1.0f; + float r = 0.0f, g = 0.0f, b = 0.0f; + + int todo = end - begin; + for (int round = 0; round < rounds; round++, todo -= BLOCK_SIZE) { + // every pixel of the tile is saturated, so no further Gaussian can + // contribute to this block + if (__syncthreads_count(done) == BLOCK_SIZE) break; + + const int fetch = begin + round * BLOCK_SIZE + lane; + if (fetch < end) { + const int gid = tile_list[fetch]; + s_xy[lane] = mean2d[gid]; + s_co[lane] = conic_opacity[gid]; + s_color[lane] = color_depth[gid]; + } + __syncthreads(); + + const int count = min(BLOCK_SIZE, todo); + for (int j = 0; j < count && !done; j++) { + const float2 xy = s_xy[j]; + const float4 co = s_co[j]; + const float dx = xy.x - pixf_x; + const float dy = xy.y - pixf_y; + + const float power = -0.5f * fmaf(co.x, dx * dx, co.z * dy * dy) - co.y * dx * dy; + if (power > 0.0f) continue; + + const float alpha = fminf(0.99f, co.w * __expf(power)); + if (alpha < MIN_OPACITY) continue; + + const float weight = alpha * transmittance; + const float4 c = s_color[j]; + r = fmaf(c.x, weight, r); + g = fmaf(c.y, weight, g); + b = fmaf(c.z, weight, b); + + transmittance *= 1.0f - alpha; + if (transmittance < MIN_TRANSMITTANCE) done = true; + } + __syncthreads(); + } + + if (inside) + image[(size_t)y * cam.width + x] = make_float4(r, g, b, 1.0f - transmittance); +} + +// --------------------------------------------------------------------------- + +int main(int argc, char* argv[]) +{ + if (argc != 5) { + printf("Usage: %s \n", + argv[0]); + return 1; + } + + const int n = atoi(argv[1]); + const int width = atoi(argv[2]); + const int height = atoi(argv[3]); + const int repeat = atoi(argv[4]); + + if (n <= 0 || width <= 0 || height <= 0 || repeat <= 0) { + printf("Error: number of gaussians, image width, image height, and repeat " + "must all be positive integers (got n=%d, width=%d, height=%d, " + "repeat=%d)\n", n, width, height, repeat); + return 1; + } + + Camera cam; + setup_camera(width, height, cam); + MpmParams mpm; + setup_mpm(mpm); + const float time = 0.5f; // the instant of the 4D scene that is rendered + + Scene scene; + std::vector h_image, h_ref_image; + std::vector h_mean2d, h_conic, h_color; + std::vector h_radii; + std::vector tile_offsets, tile_list; + try { + generate_scene(n, scene); + h_mean2d.resize((size_t)2 * n); + h_conic.resize((size_t)4 * n); + h_color.resize((size_t)4 * n); + h_radii.resize(n); + h_image.resize(4 * (size_t)width * height); + h_ref_image.resize(4 * (size_t)width * height); + } catch (const std::bad_alloc&) { + printf("Failed to allocate the host buffers for %d gaussians and a %d x %d " + "image\n", n, width, height); + return 1; + } + + const int num_tiles = cam.tiles_x * cam.tiles_y; + printf("Gaussians: %d, image: %d x %d (%d tiles), MPM grid: %d^3\n", + n, width, height, num_tiles, MPM_GRID); + + // the host reference evolves its own copy of the scene + Scene ref_scene = scene; + + float4 *d_mean, *d_scale, *d_quat_l, *d_quat_r, *d_velocity; + float *d_opacity, *d_sh, *d_affine, *d_defgrad; + float4 *d_grid; + CHECK(cudaMalloc((void**)&d_mean, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_scale, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_quat_l, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_quat_r, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_velocity, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_opacity, sizeof(float) * n)); + CHECK(cudaMalloc((void**)&d_sh, sizeof(float) * SH_COEFFS * 3 * (size_t)n)); + CHECK(cudaMalloc((void**)&d_affine, sizeof(float) * 9 * (size_t)n)); + CHECK(cudaMalloc((void**)&d_defgrad, sizeof(float) * 9 * (size_t)n)); + CHECK(cudaMalloc((void**)&d_grid, sizeof(float4) * MPM_CELLS)); + + const int num_chunks = (int)scene.chunk_block.size(); + int *d_chunk_start, *d_chunk_block; + CHECK(cudaMalloc((void**)&d_chunk_start, sizeof(int) * (num_chunks + 1))); + CHECK(cudaMalloc((void**)&d_chunk_block, sizeof(int) * num_chunks)); + CHECK(cudaMemcpy(d_chunk_start, scene.chunk_start.data(), + sizeof(int) * (num_chunks + 1), cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(d_chunk_block, scene.chunk_block.data(), + sizeof(int) * num_chunks, cudaMemcpyHostToDevice)); + + // the host keeps xyzw and xyz layouts, the device wants 16 byte vectors + std::vector pack(4 * (size_t)n); + auto upload_vec4 = [&](const std::vector& src, int comps, float4* dst) { + for (int i = 0; i < n; i++) { + pack[4 * (size_t)i + 0] = src[(size_t)comps * i + 0]; + pack[4 * (size_t)i + 1] = src[(size_t)comps * i + 1]; + pack[4 * (size_t)i + 2] = src[(size_t)comps * i + 2]; + pack[4 * (size_t)i + 3] = (comps == 4) ? src[(size_t)comps * i + 3] : 0.0f; + } + CHECK(cudaMemcpy(dst, pack.data(), sizeof(float4) * n, cudaMemcpyHostToDevice)); + }; + + // the 3x3 tensors are stored component major on the device + std::vector pack9(9 * (size_t)n); + auto upload_tensor = [&](const std::vector& src, float* dst) { + for (int k = 0; k < 9; k++) + for (int i = 0; i < n; i++) pack9[(size_t)k * n + i] = src[9 * (size_t)i + k]; + CHECK(cudaMemcpy(dst, pack9.data(), sizeof(float) * 9 * (size_t)n, + cudaMemcpyHostToDevice)); + }; + + upload_vec4(scene.mean, 4, d_mean); + upload_vec4(scene.scale, 4, d_scale); + upload_vec4(scene.quat_l, 4, d_quat_l); + upload_vec4(scene.quat_r, 4, d_quat_r); + upload_vec4(scene.velocity, 3, d_velocity); + upload_tensor(scene.affine, d_affine); + upload_tensor(scene.defgrad, d_defgrad); + CHECK(cudaMemcpy(d_opacity, scene.opacity.data(), sizeof(float) * n, + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(d_sh, scene.sh.data(), sizeof(float) * SH_COEFFS * 3 * (size_t)n, + cudaMemcpyHostToDevice)); + + float2 *d_mean2d; + float4 *d_conic, *d_color, *d_image; + int *d_radii, *d_tile_offsets, *d_tile_list; + CHECK(cudaMalloc((void**)&d_mean2d, sizeof(float2) * n)); + CHECK(cudaMalloc((void**)&d_conic, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_color, sizeof(float4) * n)); + CHECK(cudaMalloc((void**)&d_radii, sizeof(int) * n)); + CHECK(cudaMalloc((void**)&d_image, sizeof(float4) * (size_t)width * height)); + CHECK(cudaMalloc((void**)&d_tile_offsets, sizeof(int) * (num_tiles + 1))); + + const int particle_blocks = (n + P2G_BLOCK - 1) / P2G_BLOCK; + const int cell_blocks = (MPM_CELLS + GRID_BLOCK - 1) / GRID_BLOCK; + const int preprocess_blocks = (n + PREPROCESS_BLOCK - 1) / PREPROCESS_BLOCK; + + auto mpm_step = [&]() { + CHECK(cudaMemset(d_grid, 0, sizeof(float4) * MPM_CELLS)); + mpm_p2g_kernel<<>>( + n, mpm, d_chunk_start, d_chunk_block, d_mean, d_velocity, d_affine, + d_defgrad, (float*)d_grid); + mpm_grid_kernel<<>>(mpm, d_grid); + mpm_g2p_kernel<<>>( + n, mpm, d_mean, d_velocity, d_affine, d_defgrad, d_grid); + }; + + // --- stage 1: one MPM step, verified against the reference --------------- + mpm_step(); + CHECK(cudaGetLastError()); + + { + std::vector ref_grid(4 * (size_t)MPM_CELLS); + reference_p2g(ref_scene, mpm, ref_grid.data()); + reference_grid_update(mpm, ref_grid.data()); + reference_g2p(ref_scene, mpm, ref_grid.data()); + + std::vector got(4 * (size_t)n); + CHECK(cudaMemcpy(got.data(), d_mean, sizeof(float4) * n, cudaMemcpyDeviceToHost)); + int mpm_errors = 0; + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(got[4 * (size_t)i + k], ref_scene.mean[4 * (size_t)i + k], 1e-4f)) + mpm_errors++; + CHECK(cudaMemcpy(got.data(), d_velocity, sizeof(float4) * n, cudaMemcpyDeviceToHost)); + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(got[4 * (size_t)i + k], ref_scene.velocity[3 * (size_t)i + k], 1e-3f)) + mpm_errors++; + printf("MPM step: %s\n", mpm_errors == 0 ? "PASS" : "FAIL"); + } + + CHECK(cudaDeviceSynchronize()); + auto start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) mpm_step(); + CHECK(cudaDeviceSynchronize()); + auto end = std::chrono::steady_clock::now(); + auto time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the MPM step (p2g, grid, g2p): %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 2: the 4D preprocess, from the undeformed scene --------------- + upload_vec4(scene.mean, 4, d_mean); + + preprocess_kernel<<>>( + n, time, cam, d_mean, d_scale, d_quat_l, d_quat_r, d_opacity, d_sh, + d_mean2d, d_conic, d_color, d_radii); + + std::vector ref_mean2d(2 * (size_t)n), ref_conic(4 * (size_t)n), + ref_color(4 * (size_t)n); + { + std::vector ref_radii(n); + reference_preprocess(scene, cam, time, ref_mean2d.data(), ref_conic.data(), + ref_color.data(), ref_radii.data()); + + CHECK(cudaMemcpy(h_mean2d.data(), d_mean2d, sizeof(float2) * n, cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(h_conic.data(), d_conic, sizeof(float4) * n, cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(h_color.data(), d_color, sizeof(float4) * n, cudaMemcpyDeviceToHost)); + CHECK(cudaMemcpy(h_radii.data(), d_radii, sizeof(int) * n, cudaMemcpyDeviceToHost)); + + int pre_errors = 0; + int visible = 0; + for (int i = 0; i < n && pre_errors == 0; i++) { + if (ref_radii[i] > 0) visible++; + // the extent is the ceiling of a float, so it may land either side of + // an integer when the host and the device contract differently + if (abs(h_radii[i] - ref_radii[i]) > 1) pre_errors++; + if (ref_radii[i] == 0 || h_radii[i] == 0) continue; + for (int k = 0; k < 2; k++) + if (!close_enough(h_mean2d[2 * (size_t)i + k], ref_mean2d[2 * (size_t)i + k], 1e-3f)) + pre_errors++; + for (int k = 0; k < 4; k++) { + if (!close_enough(h_conic[4 * (size_t)i + k], ref_conic[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + if (!close_enough(h_color[4 * (size_t)i + k], ref_color[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + } + } + printf("4D preprocess (%d of %d gaussians visible): %s\n", visible, n, + pre_errors == 0 ? "PASS" : "FAIL"); + + // The tile lists are host side setup, shared by the reference and the + // device rasterizer. Stage 3 is fed the reference splats, so that it is + // verified on its own rather than against the rounding of stage 2. + build_tile_lists(cam, n, ref_mean2d.data(), ref_color.data(), + ref_radii.data(), tile_offsets, tile_list); + reference_render(cam, ref_mean2d.data(), ref_conic.data(), ref_color.data(), + tile_offsets.data(), tile_list.data(), h_ref_image.data()); + } + + CHECK(cudaDeviceSynchronize()); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) { + preprocess_kernel<<>>( + n, time, cam, d_mean, d_scale, d_quat_l, d_quat_r, d_opacity, d_sh, + d_mean2d, d_conic, d_color, d_radii); + } + CHECK(cudaDeviceSynchronize()); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the 4D preprocess kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 3: the tile rasterizer --------------------------------------- + CHECK(cudaMemcpy(d_mean2d, ref_mean2d.data(), sizeof(float2) * n, + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(d_conic, ref_conic.data(), sizeof(float4) * n, + cudaMemcpyHostToDevice)); + CHECK(cudaMemcpy(d_color, ref_color.data(), sizeof(float4) * n, + cudaMemcpyHostToDevice)); + + const size_t list_size = tile_list.empty() ? 1 : tile_list.size(); + CHECK(cudaMalloc((void**)&d_tile_list, sizeof(int) * list_size)); + CHECK(cudaMemcpy(d_tile_offsets, tile_offsets.data(), sizeof(int) * (num_tiles + 1), + cudaMemcpyHostToDevice)); + if (!tile_list.empty()) + CHECK(cudaMemcpy(d_tile_list, tile_list.data(), sizeof(int) * tile_list.size(), + cudaMemcpyHostToDevice)); + + printf("Gaussian instances after tiling: %zu (%.1f per tile)\n", + tile_list.size(), (double)tile_list.size() / num_tiles); + + const dim3 render_block(BLOCK_X, BLOCK_Y); + const dim3 render_grid(cam.tiles_x, cam.tiles_y); + + CHECK(cudaMemset(d_image, 0, sizeof(float4) * (size_t)width * height)); + render_kernel<<>>( + cam, d_mean2d, d_conic, d_color, d_tile_offsets, d_tile_list, d_image); + CHECK(cudaGetLastError()); + CHECK(cudaMemcpy(h_image.data(), d_image, sizeof(float4) * (size_t)width * height, + cudaMemcpyDeviceToHost)); + + int render_errors = 0; + for (size_t k = 0; k < h_image.size() && render_errors == 0; k++) + if (!close_enough(h_image[k], h_ref_image[k], 1e-3f)) render_errors++; + printf("Rasterizer: %s\n", render_errors == 0 ? "PASS" : "FAIL"); + + CHECK(cudaDeviceSynchronize()); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) { + render_kernel<<>>( + cam, d_mean2d, d_conic, d_color, d_tile_offsets, d_tile_list, d_image); + } + CHECK(cudaDeviceSynchronize()); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the rasterizer kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + CHECK(cudaFree(d_mean)); CHECK(cudaFree(d_scale)); + CHECK(cudaFree(d_quat_l)); CHECK(cudaFree(d_quat_r)); + CHECK(cudaFree(d_velocity)); CHECK(cudaFree(d_opacity)); + CHECK(cudaFree(d_sh)); CHECK(cudaFree(d_affine)); CHECK(cudaFree(d_defgrad)); + CHECK(cudaFree(d_grid)); CHECK(cudaFree(d_mean2d)); CHECK(cudaFree(d_conic)); + CHECK(cudaFree(d_color)); CHECK(cudaFree(d_radii)); CHECK(cudaFree(d_image)); + CHECK(cudaFree(d_tile_offsets)); CHECK(cudaFree(d_tile_list)); + + return 0; +} diff --git a/src/gsplat4d-cuda/reference.h b/src/gsplat4d-cuda/reference.h new file mode 100644 index 000000000..58c4077e1 --- /dev/null +++ b/src/gsplat4d-cuda/reference.h @@ -0,0 +1,917 @@ +#ifndef REFERENCE_H +#define REFERENCE_H + +#include +#include +#include +#include +#include +#include +#include + +// 4D Gaussian splatting of a deformable scene, the kind of workload found in +// physics-aware endoscopic simulators built on 4DGS and the material point +// method (inspired by EndoGSim, https://arxiv.org/abs/2605.16022). This is an +// independently written synthetic benchmark: EndoGSim's own code and data are +// not publicly released, so nothing here is ported from it. The scene is +// procedurally generated (see generate_scene) and every stage is verified +// against the host reference in this file. +// +// Three stages are benchmarked as independent microbenchmarks. They are each +// timed and validated on their own from a well-defined input; they are NOT run +// as one fused pipeline (the MPM step does not feed its deformed positions into +// the preprocess, and the rasterizer consumes the reference splats), so the +// three timings must not be read as a single end-to-end frame time: +// +// 1. an MLS-MPM step (Hu et al., https://arxiv.org/abs/1806.01923) that +// advects the Gaussians as material particles: particle-to-grid +// scatter, grid update, grid-to-particle gather; +// 2. the 4D preprocess, which conditions every 4D Gaussian on the current +// time, projects the conditional 3D covariance to a 2D conic with the +// EWA Jacobian (Zwicker et al., 2001) and evaluates the view-dependent +// colour from spherical harmonics; +// 3. the tile-based rasterizer, which alpha-blends the Gaussians of each +// 16x16 tile front to back (Kerbl et al., https://arxiv.org/abs/2308.04079). +// The per-tile binning and depth sort are host-side setup shared by every +// backend (build_tile_lists) and are deliberately excluded from the timed +// region, which measures list-consumption rasterization only. +// +// The 4D Gaussian follows the native formulation of Yang et al. +// (https://arxiv.org/abs/2310.10642): a 4x4 covariance built from a pair of +// isoclinic quaternions, conditioned on t to give a 3D Gaussian whose opacity +// decays with the temporal distance. + +#define BLOCK_X 16 +#define BLOCK_Y 16 +#define BLOCK_SIZE (BLOCK_X * BLOCK_Y) + +// spherical harmonics up to degree 3 (16 coefficients per colour channel) +#define SH_DEGREE 3 +#define SH_COEFFS 16 + +// MLS-MPM background grid +#define MPM_GRID 64 +#define MPM_CELLS (MPM_GRID * MPM_GRID * MPM_GRID) + +// The particles are binned into MPM_BLOCK^3 cell blocks so that the +// particle-to-grid scatter can accumulate in a block local tile instead of +// hammering the global grid with atomics. A particle in a block writes to +// cells [start - 1, start + MPM_BLOCK + 1], hence the tile is three cells +// wider than the block. +#define MPM_BLOCK 4 +#define MPM_TILE (MPM_BLOCK + 3) +#define MPM_TILE_CELLS (MPM_TILE * MPM_TILE * MPM_TILE) +#define MPM_BLOCKS_PER_DIM (MPM_GRID / MPM_BLOCK) +// the largest number of particles one thread block scatters in one pass +#define MPM_CHUNK 512 + +// the transmittance below which a pixel stops accumulating +#define MIN_TRANSMITTANCE 1e-4f +// Gaussians fainter than this after the temporal decay are dropped +#define MIN_OPACITY (1.0f / 255.0f) + +struct Camera { + float view[9]; // world -> camera rotation, row major + float cam_pos[3]; // camera centre in world space + float focal_x, focal_y; + float center_x, center_y; + float near_plane; + int width, height; + int tiles_x, tiles_y; +}; + +struct MpmParams { + float dt; + float dx, inv_dx; + float particle_volume; + float particle_mass; + float mu, lambda; // Lame parameters of the neo-Hookean model + float gravity; + int boundary; // number of cells clamped at each wall +}; + +// --------------------------------------------------------------------------- +// deterministic input generation (identical on the host and every backend) +// --------------------------------------------------------------------------- + +// xorshift32: reproducible without depending on the host's rand() +static inline unsigned rng_next(unsigned& state) +{ + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + return state; +} + +static inline float rng_uniform(unsigned& state) +{ + return (rng_next(state) >> 8) * (1.0f / 16777216.0f); +} + +static inline float rng_signed(unsigned& state) +{ + return 2.0f * rng_uniform(state) - 1.0f; +} + +static inline void normalize4(float* q) +{ + const float n = sqrtf(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]); + const float inv = (n > 0.0f) ? 1.0f / n : 1.0f; + for (int i = 0; i < 4; i++) q[i] *= inv; +} + +// The scene is a slab of tissue in the unit cube, so that it is both a valid +// MPM domain and visible from the camera placed in front of it. +struct Scene { + int num_gaussians; + + std::vector mean; // 4 * N: xyz in the unit cube, t in [0,1] + std::vector scale; // 4 * N: spatial extents and the temporal one + std::vector quat_l; // 4 * N + std::vector quat_r; // 4 * N + std::vector opacity; // N + std::vector sh; // SH_COEFFS * 3 * N, coefficient major + + std::vector velocity; // 3 * N + std::vector affine; // 9 * N: the APIC velocity gradient C + std::vector defgrad; // 9 * N: the deformation gradient F + + // particle chunks, each holding up to MPM_CHUNK particles of one cell block + std::vector chunk_start; // num_chunks + 1 + std::vector chunk_block; // num_chunks, the linear cell block index +}; + +static void reorder_particles(Scene& s, const std::vector& order) +{ + const int n = s.num_gaussians; + const Scene t = s; + for (int i = 0; i < n; i++) { + const int o = order[i]; + for (int k = 0; k < 4; k++) { + s.mean[4 * (size_t)i + k] = t.mean[4 * (size_t)o + k]; + s.scale[4 * (size_t)i + k] = t.scale[4 * (size_t)o + k]; + s.quat_l[4 * (size_t)i + k] = t.quat_l[4 * (size_t)o + k]; + s.quat_r[4 * (size_t)i + k] = t.quat_r[4 * (size_t)o + k]; + } + s.opacity[i] = t.opacity[o]; + for (int k = 0; k < 3; k++) + s.velocity[3 * (size_t)i + k] = t.velocity[3 * (size_t)o + k]; + for (int k = 0; k < 9; k++) { + s.affine[9 * (size_t)i + k] = t.affine[9 * (size_t)o + k]; + s.defgrad[9 * (size_t)i + k] = t.defgrad[9 * (size_t)o + k]; + } + for (int c = 0; c < SH_COEFFS * 3; c++) + s.sh[(size_t)c * n + i] = t.sh[(size_t)c * n + o]; + } +} + +static void generate_scene(int n, Scene& s) +{ + s.num_gaussians = n; + s.mean.resize((size_t)4 * n); + s.scale.resize((size_t)4 * n); + s.quat_l.resize((size_t)4 * n); + s.quat_r.resize((size_t)4 * n); + s.opacity.resize(n); + s.sh.resize((size_t)SH_COEFFS * 3 * n); + s.velocity.resize((size_t)3 * n); + s.affine.resize((size_t)9 * n); + s.defgrad.resize((size_t)9 * n); + + unsigned state = 123456789u; + + for (int i = 0; i < n; i++) { + // particles fill the middle of the grid, away from the walls + s.mean[4 * (size_t)i + 0] = 0.2f + 0.6f * rng_uniform(state); + s.mean[4 * (size_t)i + 1] = 0.2f + 0.4f * rng_uniform(state); + s.mean[4 * (size_t)i + 2] = 0.2f + 0.6f * rng_uniform(state); + s.mean[4 * (size_t)i + 3] = rng_uniform(state); // time of birth + + // small anisotropic splats, and a temporal extent that keeps a useful + // fraction of the scene alive at any one time + s.scale[4 * (size_t)i + 0] = 0.004f + 0.010f * rng_uniform(state); + s.scale[4 * (size_t)i + 1] = 0.004f + 0.010f * rng_uniform(state); + s.scale[4 * (size_t)i + 2] = 0.004f + 0.010f * rng_uniform(state); + s.scale[4 * (size_t)i + 3] = 0.03f + 0.15f * rng_uniform(state); + + float ql[4], qr[4]; + for (int k = 0; k < 4; k++) ql[k] = rng_signed(state); + normalize4(ql); + // qr near conj(ql) keeps the rotation close to a spatial one; the + // perturbation is the space-time coupling that moves the Gaussian + qr[0] = ql[0] + 0.25f * rng_signed(state); + for (int k = 1; k < 4; k++) qr[k] = -ql[k] + 0.25f * rng_signed(state); + normalize4(qr); + for (int k = 0; k < 4; k++) s.quat_l[4 * (size_t)i + k] = ql[k]; + for (int k = 0; k < 4; k++) s.quat_r[4 * (size_t)i + k] = qr[k]; + + s.opacity[i] = 0.15f + 0.75f * rng_uniform(state); + + // the band 0 coefficient carries the base colour, the higher bands the + // view dependent sheen of wet tissue + for (int c = 0; c < SH_COEFFS; c++) { + const float amp = (c == 0) ? 0.5f : 0.15f / (float)(1 + c); + for (int ch = 0; ch < 3; ch++) + s.sh[((size_t)c * 3 + ch) * n + i] = amp * rng_signed(state); + } + + // a slow swirl, so that the first step already deforms the material + const float rx = s.mean[4 * (size_t)i + 0] - 0.5f; + const float rz = s.mean[4 * (size_t)i + 2] - 0.5f; + s.velocity[3 * (size_t)i + 0] = -2.0f * rz; + s.velocity[3 * (size_t)i + 1] = 0.0f; + s.velocity[3 * (size_t)i + 2] = 2.0f * rx; + + for (int k = 0; k < 9; k++) s.affine[9 * (size_t)i + k] = 0.0f; + for (int k = 0; k < 9; k++) + s.defgrad[9 * (size_t)i + k] = (k % 4 == 0) ? 1.0f : 0.0f; + } + + // Bin the particles by MPM_BLOCK^3 cell block, then by cell. Production MPM + // codes keep the particles binned for exactly this reason: it bounds the + // grid footprint of a thread block, which is what lets the scatter stay in + // local memory, and it keeps the gather reading a few cache lines. + std::vector order(n); + std::vector key(n); + std::vector block(n); + for (int i = 0; i < n; i++) { + order[i] = i; + const int ix = std::min(MPM_GRID - 1, std::max(0, (int)(s.mean[4 * (size_t)i + 0] * MPM_GRID))); + const int iy = std::min(MPM_GRID - 1, std::max(0, (int)(s.mean[4 * (size_t)i + 1] * MPM_GRID))); + const int iz = std::min(MPM_GRID - 1, std::max(0, (int)(s.mean[4 * (size_t)i + 2] * MPM_GRID))); + block[i] = ((ix / MPM_BLOCK) * MPM_BLOCKS_PER_DIM + (iy / MPM_BLOCK)) * + MPM_BLOCKS_PER_DIM + (iz / MPM_BLOCK); + key[i] = (ix * MPM_GRID + iy) * MPM_GRID + iz; + } + std::stable_sort(order.begin(), order.end(), [&](int a, int b) { + return (block[a] != block[b]) ? (block[a] < block[b]) : (key[a] < key[b]); + }); + + reorder_particles(s, order); + + // cut the sorted particles into chunks that hold one cell block each + s.chunk_start.clear(); + s.chunk_block.clear(); + int begin = 0; + while (begin < n) { + const int b = block[order[begin]]; + int end = begin; + while (end < n && block[order[end]] == b && end - begin < MPM_CHUNK) end++; + s.chunk_start.push_back(begin); + s.chunk_block.push_back(b); + begin = end; + } + s.chunk_start.push_back(n); + + // Within a chunk the particles are still sorted by cell, so neighbouring + // lanes would scatter into the same cell and serialize on it. Interleaving + // the chunk hands each lane a particle from a different cell while keeping + // the loads of a warp contiguous. + std::vector shuffled(n); + const int lanes = 32; + for (size_t c = 0; c + 1 < s.chunk_start.size(); c++) { + const int cb = s.chunk_start[c], ce = s.chunk_start[c + 1]; + const int len = ce - cb; + const int rows = (len + lanes - 1) / lanes; + int out = cb; + for (int r = 0; r < rows; r++) + for (int l = 0; l < lanes; l++) { + const int src = l * rows + r; + if (src < len) shuffled[out++] = cb + src; + } + } + reorder_particles(s, shuffled); +} + +static void setup_camera(int width, int height, Camera& cam) +{ + // Look at the centre of the tissue block from a yawed and slightly raised + // position, so that the view rotation exercises all nine entries. + const float target[3] = { 0.5f, 0.4f, 0.5f }; + const float yaw = 0.35f, pitch = 0.20f, dist = 1.6f; + + // f points from the camera to the target + const float f[3] = { sinf(yaw) * cosf(pitch), -sinf(pitch), + cosf(yaw) * cosf(pitch) }; + for (int k = 0; k < 3; k++) cam.cam_pos[k] = target[k] - dist * f[k]; + + // r = up x f, u = f x r + const float up[3] = { 0.0f, 1.0f, 0.0f }; + float r[3] = { up[1] * f[2] - up[2] * f[1], + up[2] * f[0] - up[0] * f[2], + up[0] * f[1] - up[1] * f[0] }; + const float rn = sqrtf(r[0] * r[0] + r[1] * r[1] + r[2] * r[2]); + for (int k = 0; k < 3; k++) r[k] /= rn; + const float u[3] = { f[1] * r[2] - f[2] * r[1], + f[2] * r[0] - f[0] * r[2], + f[0] * r[1] - f[1] * r[0] }; + + for (int k = 0; k < 3; k++) { + cam.view[0 * 3 + k] = r[k]; + cam.view[1 * 3 + k] = u[k]; + cam.view[2 * 3 + k] = f[k]; + } + + cam.focal_x = 0.8f * width; + cam.focal_y = 0.8f * width; + cam.center_x = 0.5f * width; + cam.center_y = 0.5f * height; + cam.near_plane = 0.2f; + cam.width = width; + cam.height = height; + cam.tiles_x = (width + BLOCK_X - 1) / BLOCK_X; + cam.tiles_y = (height + BLOCK_Y - 1) / BLOCK_Y; +} + +static void setup_mpm(MpmParams& p) +{ + p.dx = 1.0f / MPM_GRID; + p.inv_dx = (float)MPM_GRID; + p.dt = 5e-5f; + p.particle_volume = p.dx * p.dx * p.dx * 0.25f; + p.particle_mass = p.particle_volume * 1000.0f; + // soft tissue: a low Young's modulus and a nearly incompressible Poisson + const float E = 5.0e3f, nu = 0.4f; + p.mu = E / (2.0f * (1.0f + nu)); + p.lambda = E * nu / ((1.0f + nu) * (1.0f - 2.0f * nu)); + p.gravity = -9.8f; + p.boundary = 3; + return; +} + +// --------------------------------------------------------------------------- +// shared math, written so the device kernels can use exactly the same formulas +// --------------------------------------------------------------------------- + +// Build the 4x4 covariance of a 4D Gaussian from the scale and the pair of +// isoclinic quaternions, and condition it on the time t. Returns the 6 unique +// entries of the conditional 3D covariance, the conditional mean and the +// opacity attenuation. Yang et al., https://arxiv.org/abs/2310.10642 +static inline void condition_4d_gaussian( + const float* mean4, const float* scale4, + const float* ql, const float* qr, float time, + float* mean3, float* cov3, float* opacity_scale) +{ + // Left and right isoclinic rotations. The quaternion basis is ordered + // (t, x, y, z): with qr = conj(ql) the rotation is purely spatial, and the + // deviation of qr from conj(ql) is what couples space and time, that is, + // what makes the Gaussian travel as t advances. + const float lw = ql[0], lx = ql[1], ly = ql[2], lz = ql[3]; + const float rw = qr[0], rx = qr[1], ry = qr[2], rz = qr[3]; + + const float L[16] = { + lw, -lx, -ly, -lz, + lx, lw, -lz, ly, + ly, lz, lw, -lx, + lz, -ly, lx, lw }; + const float R[16] = { + rw, -rx, -ry, -rz, + rx, rw, rz, -ry, + ry, -rz, rw, rx, + rz, ry, -rx, rw }; + + // M = L * R is the 4D rotation + float M[16]; + for (int i = 0; i < 4; i++) + for (int j = 0; j < 4; j++) { + float acc = 0.0f; + for (int k = 0; k < 4; k++) acc += L[i * 4 + k] * R[k * 4 + j]; + M[i * 4 + j] = acc; + } + + // Sigma = M diag(s^2) M^T, in the (t, x, y, z) basis + const float s2[4] = { scale4[3] * scale4[3], scale4[0] * scale4[0], + scale4[1] * scale4[1], scale4[2] * scale4[2] }; + float sigma[16]; + for (int i = 0; i < 4; i++) + for (int j = 0; j < 4; j++) { + float acc = 0.0f; + for (int k = 0; k < 4; k++) acc += M[i * 4 + k] * s2[k] * M[j * 4 + k]; + sigma[i * 4 + j] = acc; + } + + // condition the spatial block on the temporal one + const float sigma_tt = sigma[0]; + const float inv_tt = 1.0f / sigma_tt; + const float st[3] = { sigma[4], sigma[8], sigma[12] }; // Sigma_{xyz,t} + const float dt = time - mean4[3]; + + for (int i = 0; i < 3; i++) mean3[i] = mean4[i] + dt * st[i] * inv_tt; + + cov3[0] = sigma[5] - st[0] * st[0] * inv_tt; // xx + cov3[1] = sigma[6] - st[0] * st[1] * inv_tt; // xy + cov3[2] = sigma[7] - st[0] * st[2] * inv_tt; // xz + cov3[3] = sigma[10] - st[1] * st[1] * inv_tt; // yy + cov3[4] = sigma[11] - st[1] * st[2] * inv_tt; // yz + cov3[5] = sigma[15] - st[2] * st[2] * inv_tt; // zz + + *opacity_scale = expf(-0.5f * dt * dt * inv_tt); +} + +// Project the 3D covariance to the 2D conic with the EWA Jacobian. +// Returns false if the Gaussian is behind the near plane or degenerate. +static inline bool project_gaussian( + const float* mean3, const float* cov3, const Camera& cam, + float* mean2d, float* conic, float* radius, float* depth) +{ + // world -> camera + const float p[3] = { mean3[0] - cam.cam_pos[0], + mean3[1] - cam.cam_pos[1], + mean3[2] - cam.cam_pos[2] }; + float t[3]; + for (int i = 0; i < 3; i++) + t[i] = cam.view[i * 3 + 0] * p[0] + cam.view[i * 3 + 1] * p[1] + + cam.view[i * 3 + 2] * p[2]; + + if (t[2] < cam.near_plane) return false; + + const float inv_z = 1.0f / t[2]; + const float inv_z2 = inv_z * inv_z; + + // J is the Jacobian of the perspective projection, clamped as in EWA + // splatting so that Gaussians far off axis stay well conditioned + const float lim_x = 1.3f * (0.5f * cam.width) / cam.focal_x; + const float lim_y = 1.3f * (0.5f * cam.height) / cam.focal_y; + const float tx = t[2] * std::min(lim_x, std::max(-lim_x, t[0] * inv_z)); + const float ty = t[2] * std::min(lim_y, std::max(-lim_y, t[1] * inv_z)); + + const float J[6] = { cam.focal_x * inv_z, 0.0f, -cam.focal_x * tx * inv_z2, + 0.0f, cam.focal_y * inv_z, -cam.focal_y * ty * inv_z2 }; + + // T = J * W (2x3) + float T[6]; + for (int i = 0; i < 2; i++) + for (int j = 0; j < 3; j++) + T[i * 3 + j] = J[i * 3 + 0] * cam.view[0 * 3 + j] + + J[i * 3 + 1] * cam.view[1 * 3 + j] + + J[i * 3 + 2] * cam.view[2 * 3 + j]; + + const float S[9] = { cov3[0], cov3[1], cov3[2], + cov3[1], cov3[3], cov3[4], + cov3[2], cov3[4], cov3[5] }; + + // cov2d = T S T^T, with the low pass filter of one third of a pixel + float TS[6]; + for (int i = 0; i < 2; i++) + for (int j = 0; j < 3; j++) + TS[i * 3 + j] = T[i * 3 + 0] * S[0 * 3 + j] + T[i * 3 + 1] * S[1 * 3 + j] + + T[i * 3 + 2] * S[2 * 3 + j]; + + float a = 0.0f, b = 0.0f, c = 0.0f; + for (int k = 0; k < 3; k++) { + a += TS[0 * 3 + k] * T[0 * 3 + k]; + b += TS[0 * 3 + k] * T[1 * 3 + k]; + c += TS[1 * 3 + k] * T[1 * 3 + k]; + } + a += 0.3f; + c += 0.3f; + + const float det = a * c - b * b; + if (det <= 0.0f) return false; + const float inv_det = 1.0f / det; + + conic[0] = c * inv_det; + conic[1] = -b * inv_det; + conic[2] = a * inv_det; + + // three sigma of the larger principal axis + const float mid = 0.5f * (a + c); + const float disc = sqrtf(std::max(0.1f, mid * mid - det)); + *radius = ceilf(3.0f * sqrtf(std::max(mid + disc, mid - disc))); + + mean2d[0] = cam.focal_x * t[0] * inv_z + cam.center_x; + mean2d[1] = cam.focal_y * t[1] * inv_z + cam.center_y; + *depth = t[2]; + return true; +} + +// Real spherical harmonics up to degree 3, the basis used by every Gaussian +// splatting implementation. +// spelled out as macros so that the device kernels and the host reference +// share one definition +#define SH_C0 0.28209479177387814f +#define SH_C1 0.4886025119029199f +#define SH_C2_0 1.0925484305920792f +#define SH_C2_1 -1.0925484305920792f +#define SH_C2_2 0.31539156525252005f +#define SH_C2_3 -1.0925484305920792f +#define SH_C2_4 0.5462742152960396f +#define SH_C3_0 -0.5900435899266435f +#define SH_C3_1 2.890611442640554f +#define SH_C3_2 -0.4570457994644658f +#define SH_C3_3 0.3731763325901154f +#define SH_C3_4 -0.4570457994644658f +#define SH_C3_5 1.445305721320277f +#define SH_C3_6 -0.5900435899266435f + +// `sh` is coefficient major: sh[(c * 3 + channel) * n + index] +static inline void eval_sh(const float* sh, int n, int index, + const float* dir, float* rgb) +{ + const float x = dir[0], y = dir[1], z = dir[2]; + float basis[SH_COEFFS]; + + basis[0] = SH_C0; + basis[1] = -SH_C1 * y; + basis[2] = SH_C1 * z; + basis[3] = -SH_C1 * x; + + const float xx = x * x, yy = y * y, zz = z * z; + const float xy = x * y, yz = y * z, xz = x * z; + basis[4] = SH_C2_0 * xy; + basis[5] = SH_C2_1 * yz; + basis[6] = SH_C2_2 * (2.0f * zz - xx - yy); + basis[7] = SH_C2_3 * xz; + basis[8] = SH_C2_4 * (xx - yy); + + basis[9] = SH_C3_0 * y * (3.0f * xx - yy); + basis[10] = SH_C3_1 * xy * z; + basis[11] = SH_C3_2 * y * (4.0f * zz - xx - yy); + basis[12] = SH_C3_3 * z * (2.0f * zz - 3.0f * xx - 3.0f * yy); + basis[13] = SH_C3_4 * x * (4.0f * zz - xx - yy); + basis[14] = SH_C3_5 * z * (xx - yy); + basis[15] = SH_C3_6 * x * (xx - 3.0f * yy); + + for (int ch = 0; ch < 3; ch++) { + float acc = 0.0f; + for (int c = 0; c < SH_COEFFS; c++) + acc += basis[c] * sh[((size_t)c * 3 + ch) * n + index]; + rgb[ch] = std::max(acc + 0.5f, 0.0f); + } +} + +// --------------------------------------------------------------------------- +// host reference +// --------------------------------------------------------------------------- + +// Stage 2: condition on time, project, shade. +static void reference_preprocess( + const Scene& s, const Camera& cam, float time, + float* mean2d, float* conic_opacity, float* color_depth, int* radii) +{ + const int n = s.num_gaussians; + + for (int i = 0; i < n; i++) { + radii[i] = 0; + mean2d[2 * (size_t)i + 0] = 0.0f; + mean2d[2 * (size_t)i + 1] = 0.0f; + for (int k = 0; k < 4; k++) { + conic_opacity[4 * (size_t)i + k] = 0.0f; + color_depth[4 * (size_t)i + k] = 0.0f; + } + + float mean3[3], cov3[6], opacity_scale; + condition_4d_gaussian(&s.mean[4 * (size_t)i], &s.scale[4 * (size_t)i], + &s.quat_l[4 * (size_t)i], &s.quat_r[4 * (size_t)i], + time, mean3, cov3, &opacity_scale); + + const float alpha = s.opacity[i] * opacity_scale; + if (alpha < MIN_OPACITY) continue; + + float m2d[2], conic[3], radius, depth; + if (!project_gaussian(mean3, cov3, cam, m2d, conic, &radius, &depth)) + continue; + + if (m2d[0] + radius < 0.0f || m2d[0] - radius > (float)cam.width || + m2d[1] + radius < 0.0f || m2d[1] - radius > (float)cam.height) + continue; + + float dir[3] = { mean3[0] - cam.cam_pos[0], + mean3[1] - cam.cam_pos[1], + mean3[2] - cam.cam_pos[2] }; + const float inv_len = + 1.0f / sqrtf(dir[0] * dir[0] + dir[1] * dir[1] + dir[2] * dir[2]); + for (int k = 0; k < 3; k++) dir[k] *= inv_len; + + float rgb[3]; + eval_sh(s.sh.data(), n, i, dir, rgb); + + radii[i] = (int)radius; + mean2d[2 * (size_t)i + 0] = m2d[0]; + mean2d[2 * (size_t)i + 1] = m2d[1]; + conic_opacity[4 * (size_t)i + 0] = conic[0]; + conic_opacity[4 * (size_t)i + 1] = conic[1]; + conic_opacity[4 * (size_t)i + 2] = conic[2]; + conic_opacity[4 * (size_t)i + 3] = alpha; + color_depth[4 * (size_t)i + 0] = rgb[0]; + color_depth[4 * (size_t)i + 1] = rgb[1]; + color_depth[4 * (size_t)i + 2] = rgb[2]; + color_depth[4 * (size_t)i + 3] = depth; + } +} + +// Build the per tile lists, sorted front to back. Production implementations +// do this with a radix sort over (tile, depth) keys; here it is host side +// setup shared by every backend, so that the rasterizer sees identical input. +static void build_tile_lists( + const Camera& cam, int n, const float* mean2d, const float* color_depth, + const int* radii, std::vector& tile_offsets, std::vector& tile_list) +{ + const int num_tiles = cam.tiles_x * cam.tiles_y; + tile_offsets.assign(num_tiles + 1, 0); + + std::vector counts(num_tiles, 0); + for (int i = 0; i < n; i++) { + if (radii[i] <= 0) continue; + const float r = (float)radii[i]; + const int x0 = std::max(0, (int)floorf((mean2d[2 * (size_t)i + 0] - r) / BLOCK_X)); + const int x1 = std::min(cam.tiles_x - 1, (int)floorf((mean2d[2 * (size_t)i + 0] + r) / BLOCK_X)); + const int y0 = std::max(0, (int)floorf((mean2d[2 * (size_t)i + 1] - r) / BLOCK_Y)); + const int y1 = std::min(cam.tiles_y - 1, (int)floorf((mean2d[2 * (size_t)i + 1] + r) / BLOCK_Y)); + for (int ty = y0; ty <= y1; ty++) + for (int tx = x0; tx <= x1; tx++) + counts[ty * cam.tiles_x + tx]++; + } + + // The prefix sums and per-tile cursors below are int, and the device tile + // buffers are int-indexed, so guard against a total gaussian/tile intersection + // count that would overflow int before it corrupts memory. + size_t total = 0; + for (int t = 0; t < num_tiles; t++) total += (size_t)counts[t]; + if (total > (size_t)INT_MAX) { + fprintf(stderr, + "Too many gaussian/tile intersections (%zu); reduce the number of " + "gaussians or the image size\n", total); + exit(EXIT_FAILURE); + } + + for (int t = 0; t < num_tiles; t++) tile_offsets[t + 1] = tile_offsets[t] + counts[t]; + tile_list.assign(tile_offsets[num_tiles], 0); + + std::vector cursor(tile_offsets.begin(), tile_offsets.end() - 1); + for (int i = 0; i < n; i++) { + if (radii[i] <= 0) continue; + const float r = (float)radii[i]; + const int x0 = std::max(0, (int)floorf((mean2d[2 * (size_t)i + 0] - r) / BLOCK_X)); + const int x1 = std::min(cam.tiles_x - 1, (int)floorf((mean2d[2 * (size_t)i + 0] + r) / BLOCK_X)); + const int y0 = std::max(0, (int)floorf((mean2d[2 * (size_t)i + 1] - r) / BLOCK_Y)); + const int y1 = std::min(cam.tiles_y - 1, (int)floorf((mean2d[2 * (size_t)i + 1] + r) / BLOCK_Y)); + for (int ty = y0; ty <= y1; ty++) + for (int tx = x0; tx <= x1; tx++) + tile_list[cursor[ty * cam.tiles_x + tx]++] = i; + } + + for (int t = 0; t < num_tiles; t++) { + std::sort(tile_list.begin() + tile_offsets[t], + tile_list.begin() + tile_offsets[t + 1], + [&](int a, int b) { + const float da = color_depth[4 * (size_t)a + 3]; + const float db = color_depth[4 * (size_t)b + 3]; + return (da != db) ? (da < db) : (a < b); + }); + } +} + +// Stage 3: front to back alpha blending of one tile list per tile. +static void reference_render( + const Camera& cam, const float* mean2d, const float* conic_opacity, + const float* color_depth, const int* tile_offsets, const int* tile_list, + float* image) +{ + for (int ty = 0; ty < cam.tiles_y; ty++) { + for (int tx = 0; tx < cam.tiles_x; tx++) { + const int tile = ty * cam.tiles_x + tx; + const int begin = tile_offsets[tile], end = tile_offsets[tile + 1]; + + for (int py = 0; py < BLOCK_Y; py++) { + const int y = ty * BLOCK_Y + py; + if (y >= cam.height) continue; + for (int px = 0; px < BLOCK_X; px++) { + const int x = tx * BLOCK_X + px; + if (x >= cam.width) continue; + + float transmittance = 1.0f; + float rgb[3] = { 0.0f, 0.0f, 0.0f }; + + for (int k = begin; k < end; k++) { + const int g = tile_list[k]; + const float dx = mean2d[2 * (size_t)g + 0] - (float)x - 0.5f; + const float dy = mean2d[2 * (size_t)g + 1] - (float)y - 0.5f; + const float* co = &conic_opacity[4 * (size_t)g]; + const float power = + -0.5f * (co[0] * dx * dx + co[2] * dy * dy) - co[1] * dx * dy; + if (power > 0.0f) continue; + + const float alpha = std::min(0.99f, co[3] * expf(power)); + if (alpha < MIN_OPACITY) continue; + + const float weight = alpha * transmittance; + for (int ch = 0; ch < 3; ch++) + rgb[ch] += color_depth[4 * (size_t)g + ch] * weight; + + transmittance *= 1.0f - alpha; + if (transmittance < MIN_TRANSMITTANCE) break; + } + + const size_t o = 4 * ((size_t)y * cam.width + x); + image[o + 0] = rgb[0]; + image[o + 1] = rgb[1]; + image[o + 2] = rgb[2]; + image[o + 3] = 1.0f - transmittance; + } + } + } + } +} + +// --------------------------------------------------------------------------- +// MLS-MPM reference (stage 1) +// --------------------------------------------------------------------------- + +static inline void quadratic_weights(float fx, float w[3]) +{ + w[0] = 0.5f * (1.5f - fx) * (1.5f - fx); + w[1] = 0.75f - (fx - 1.0f) * (fx - 1.0f); + w[2] = 0.5f * (fx - 0.5f) * (fx - 0.5f); +} + +// The neo-Hookean first Piola-Kirchhoff stress, mapped to the MLS-MPM +// affine momentum contribution: -dt * volume * 4 * inv_dx^2 * P F^T +static inline void neo_hookean_stress(const float* F, const MpmParams& p, + float* stress) +{ + const float J = + F[0] * (F[4] * F[8] - F[5] * F[7]) - + F[1] * (F[3] * F[8] - F[5] * F[6]) + + F[2] * (F[3] * F[7] - F[4] * F[6]); + const float safe_J = (fabsf(J) < 1e-6f) ? ((J < 0.0f) ? -1e-6f : 1e-6f) : J; + const float inv_J = 1.0f / safe_J; + + // cofactor matrix = J * F^{-T} + const float C[9] = { + (F[4] * F[8] - F[5] * F[7]), -(F[3] * F[8] - F[5] * F[6]), (F[3] * F[7] - F[4] * F[6]), + -(F[1] * F[8] - F[2] * F[7]), (F[0] * F[8] - F[2] * F[6]), -(F[0] * F[7] - F[1] * F[6]), + (F[1] * F[5] - F[2] * F[4]), -(F[0] * F[5] - F[2] * F[3]), (F[0] * F[4] - F[1] * F[3]) }; + + // P = mu (F - F^{-T}) + lambda log(J) F^{-T} + const float coeff = (p.lambda * logf(fabsf(safe_J)) - p.mu) * inv_J; + float P[9]; + for (int k = 0; k < 9; k++) P[k] = p.mu * F[k] + coeff * C[k]; + + // stress = -dt * volume * 4 * inv_dx^2 * P F^T + const float s = -p.dt * p.particle_volume * 4.0f * p.inv_dx * p.inv_dx; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) { + float acc = 0.0f; + for (int k = 0; k < 3; k++) acc += P[i * 3 + k] * F[j * 3 + k]; + stress[i * 3 + j] = s * acc; + } +} + +static void reference_p2g(const Scene& s, const MpmParams& p, float* grid) +{ + const int n = s.num_gaussians; + memset(grid, 0, sizeof(float) * 4 * (size_t)MPM_CELLS); + + for (int i = 0; i < n; i++) { + const float* x = &s.mean[4 * (size_t)i]; + const float* v = &s.velocity[3 * (size_t)i]; + const float* C = &s.affine[9 * (size_t)i]; + const float* F0 = &s.defgrad[9 * (size_t)i]; + + // F <- (I + dt C) F + float F[9]; + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + for (int k = 0; k < 3; k++) + acc += p.dt * C[a * 3 + k] * F0[k * 3 + b]; + F[a * 3 + b] = acc; + } + + float stress[9]; + neo_hookean_stress(F, p, stress); + + // affine = stress + mass * C + float affine[9]; + for (int k = 0; k < 9; k++) affine[k] = stress[k] + p.particle_mass * C[k]; + + const float gx = x[0] * p.inv_dx, gy = x[1] * p.inv_dx, gz = x[2] * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + float wx[3], wy[3], wz[3]; + quadratic_weights(gx - bx, wx); + quadratic_weights(gy - by, wy); + quadratic_weights(gz - bz, wz); + + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) + for (int c = 0; c < 3; c++) { + const int ix = bx + a, iy = by + b, iz = bz + c; + if (ix < 0 || ix >= MPM_GRID || iy < 0 || iy >= MPM_GRID || + iz < 0 || iz >= MPM_GRID) + continue; + const float w = wx[a] * wy[b] * wz[c]; + const float dpos[3] = { ((float)a - (gx - bx)) * p.dx, + ((float)b - (gy - by)) * p.dx, + ((float)c - (gz - bz)) * p.dx }; + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + for (int k = 0; k < 3; k++) { + const float impulse = affine[k * 3 + 0] * dpos[0] + + affine[k * 3 + 1] * dpos[1] + + affine[k * 3 + 2] * dpos[2]; + grid[o + k] += w * (p.particle_mass * v[k] + impulse); + } + grid[o + 3] += w * p.particle_mass; + } + } +} + +static void reference_grid_update(const MpmParams& p, float* grid) +{ + for (int ix = 0; ix < MPM_GRID; ix++) + for (int iy = 0; iy < MPM_GRID; iy++) + for (int iz = 0; iz < MPM_GRID; iz++) { + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + const float mass = grid[o + 3]; + if (mass <= 0.0f) { + grid[o + 0] = grid[o + 1] = grid[o + 2] = 0.0f; + continue; + } + const float inv_mass = 1.0f / mass; + float v[3] = { grid[o + 0] * inv_mass, grid[o + 1] * inv_mass, + grid[o + 2] * inv_mass }; + v[1] += p.dt * p.gravity; + + // sticky walls + if (ix < p.boundary && v[0] < 0.0f) v[0] = 0.0f; + if (ix >= MPM_GRID - p.boundary && v[0] > 0.0f) v[0] = 0.0f; + if (iy < p.boundary && v[1] < 0.0f) v[1] = 0.0f; + if (iy >= MPM_GRID - p.boundary && v[1] > 0.0f) v[1] = 0.0f; + if (iz < p.boundary && v[2] < 0.0f) v[2] = 0.0f; + if (iz >= MPM_GRID - p.boundary && v[2] > 0.0f) v[2] = 0.0f; + + for (int k = 0; k < 3; k++) grid[o + k] = v[k]; + } +} + +static void reference_g2p(Scene& s, const MpmParams& p, const float* grid) +{ + const int n = s.num_gaussians; + + for (int i = 0; i < n; i++) { + float* x = &s.mean[4 * (size_t)i]; + float* v = &s.velocity[3 * (size_t)i]; + float* C = &s.affine[9 * (size_t)i]; + float* F0 = &s.defgrad[9 * (size_t)i]; + + const float gx = x[0] * p.inv_dx, gy = x[1] * p.inv_dx, gz = x[2] * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + float wx[3], wy[3], wz[3]; + quadratic_weights(gx - bx, wx); + quadratic_weights(gy - by, wy); + quadratic_weights(gz - bz, wz); + + float new_v[3] = { 0.0f, 0.0f, 0.0f }; + float new_C[9] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) + for (int c = 0; c < 3; c++) { + const int ix = bx + a, iy = by + b, iz = bz + c; + if (ix < 0 || ix >= MPM_GRID || iy < 0 || iy >= MPM_GRID || + iz < 0 || iz >= MPM_GRID) + continue; + const float w = wx[a] * wy[b] * wz[c]; + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + const float gv[3] = { grid[o + 0], grid[o + 1], grid[o + 2] }; + const float dpos[3] = { ((float)a - (gx - bx)) * p.dx, + ((float)b - (gy - by)) * p.dx, + ((float)c - (gz - bz)) * p.dx }; + for (int k = 0; k < 3; k++) { + new_v[k] += w * gv[k]; + for (int l = 0; l < 3; l++) + new_C[k * 3 + l] += 4.0f * p.inv_dx * p.inv_dx * w * gv[k] * dpos[l]; + } + } + + // F <- (I + dt C) F, with the updated C + float F[9]; + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + for (int k = 0; k < 3; k++) + acc += p.dt * new_C[a * 3 + k] * F0[k * 3 + b]; + F[a * 3 + b] = acc; + } + + for (int k = 0; k < 3; k++) { + v[k] = new_v[k]; + x[k] += p.dt * new_v[k]; + } + for (int k = 0; k < 9; k++) { + C[k] = new_C[k]; + F0[k] = F[k]; + } + } +} + +static bool close_enough(float a, float b, float tol) +{ + return fabsf(a - b) <= tol * (1.0f + fabsf(b)); +} + +#endif diff --git a/src/gsplat4d-hip/CMakeLists.txt b/src/gsplat4d-hip/CMakeLists.txt new file mode 100644 index 000000000..760bdbb14 --- /dev/null +++ b/src/gsplat4d-hip/CMakeLists.txt @@ -0,0 +1,11 @@ +# gsplat4d-hip/CMakeLists.txt + +add_hecbench_benchmark( + NAME gsplat4d + MODEL hip + SOURCES main.cu + CATEGORIES simulation + TEST_ARGS 200000 1280 720 100 + TEST_REGEX "(?s)MPM step: PASS.*4D preprocess \([0-9]+ of [0-9]+ gaussians visible\): PASS.*Rasterizer: PASS" + TEST_TIMEOUT 300 +) diff --git a/src/gsplat4d-hip/Makefile b/src/gsplat4d-hip/Makefile new file mode 100644 index 000000000..dd6fdb9d3 --- /dev/null +++ b/src/gsplat4d-hip/Makefile @@ -0,0 +1,55 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = hipcc +OPTIMIZE = yes +DEBUG = no +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cu + +obj = $(source:.cu=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../gsplat4d-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cu ../gsplat4d-cuda/reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-hip/main.cu b/src/gsplat4d-hip/main.cu new file mode 100644 index 000000000..6be7a9b96 --- /dev/null +++ b/src/gsplat4d-hip/main.cu @@ -0,0 +1,906 @@ +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#define CHECK(call) \ + do { \ + const hipError_t err = (call); \ + if (err != hipSuccess) { \ + fprintf(stderr, "HIP error %s:%d '%s': %s\n", __FILE__, __LINE__, \ + #call, hipGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +#ifndef P2G_BLOCK +#define P2G_BLOCK 128 +#endif +#ifndef GRID_BLOCK +#define GRID_BLOCK 256 +#endif +#ifndef PREPROCESS_BLOCK +#define PREPROCESS_BLOCK 128 +#endif + +// --------------------------------------------------------------------------- +// stage 1: MLS-MPM +// +// The per particle 3x3 tensors are stored component major, so that the nine +// loads of a warp are nine fully coalesced 128 byte transactions instead of +// nine strided ones. The particles are pre-sorted by cell (see +// generate_scene), which keeps the scattered atomics of a warp inside a +// handful of cache lines. +// --------------------------------------------------------------------------- + +__device__ __forceinline__ void quad_weights(float fx, float& w0, float& w1, + float& w2) +{ + const float a = 1.5f - fx; + const float b = fx - 1.0f; + const float c = fx - 0.5f; + w0 = 0.5f * a * a; + w1 = 0.75f - b * b; + w2 = 0.5f * c * c; +} + +// One thread block scatters one chunk of particles, all of which belong to the +// same MPM_BLOCK^3 cell block. Their entire stencil footprint therefore fits a +// MPM_TILE^3 tile held in shared memory, so the scatter costs shared atomics +// plus one flush of the tile, rather than 108 global atomics per particle. +// A particle that has drifted out of its block since the binning still lands +// correctly through the global fallback. +__global__ void __launch_bounds__(P2G_BLOCK) +mpm_p2g_kernel(int n, MpmParams p, + const int* __restrict__ chunk_start, + const int* __restrict__ chunk_block, + const float4* __restrict__ mean, + const float4* __restrict__ velocity, + const float* __restrict__ affine_in, + const float* __restrict__ defgrad, + float* __restrict__ grid) +{ + __shared__ float tile[MPM_TILE_CELLS * 4]; + + for (int k = threadIdx.x; k < MPM_TILE_CELLS * 4; k += P2G_BLOCK) tile[k] = 0.0f; + + const int chunk = blockIdx.x; + const int begin = chunk_start[chunk]; + const int end = chunk_start[chunk + 1]; + + const int gb = chunk_block[chunk]; + const int obz = (gb % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int oby = ((gb / MPM_BLOCKS_PER_DIM) % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int obx = (gb / (MPM_BLOCKS_PER_DIM * MPM_BLOCKS_PER_DIM)) * MPM_BLOCK - 1; + + __syncthreads(); + + for (int i = begin + threadIdx.x; i < end; i += P2G_BLOCK) { + const float4 x = mean[i]; + const float4 v = velocity[i]; + + float C[9], F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) C[k] = affine_in[(size_t)k * n + i]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + // F <- (I + dt C) F + float F[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = fmaf(p.dt * C[a * 3 + k], F0[k * 3 + b], acc); + F[a * 3 + b] = acc; + } + + // neo-Hookean stress, folded into the MLS-MPM affine momentum + const float det = + F[0] * (F[4] * F[8] - F[5] * F[7]) - + F[1] * (F[3] * F[8] - F[5] * F[6]) + + F[2] * (F[3] * F[7] - F[4] * F[6]); + const float safe_J = (fabsf(det) < 1e-6f) ? ((det < 0.0f) ? -1e-6f : 1e-6f) : det; + const float inv_J = 1.0f / safe_J; + + const float cof[9] = { + (F[4] * F[8] - F[5] * F[7]), -(F[3] * F[8] - F[5] * F[6]), (F[3] * F[7] - F[4] * F[6]), + -(F[1] * F[8] - F[2] * F[7]), (F[0] * F[8] - F[2] * F[6]), -(F[0] * F[7] - F[1] * F[6]), + (F[1] * F[5] - F[2] * F[4]), -(F[0] * F[5] - F[2] * F[3]), (F[0] * F[4] - F[1] * F[3]) }; + + const float coeff = (p.lambda * __logf(fabsf(safe_J)) - p.mu) * inv_J; + float P[9]; + #pragma unroll + for (int k = 0; k < 9; k++) P[k] = fmaf(coeff, cof[k], p.mu * F[k]); + + const float s = -p.dt * p.particle_volume * 4.0f * p.inv_dx * p.inv_dx; + float affine[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 3; k++) acc = fmaf(P[a * 3 + k], F[b * 3 + k], acc); + affine[a * 3 + b] = fmaf(s, acc, p.particle_mass * C[a * 3 + b]); + } + + // F is only used for the stress above; it is deliberately not persisted + // here. The single per-step deformation-gradient update is applied once, in + // mpm_g2p, from the original F0 (this matches the host reference; writing F + // back here would advance the gradient twice per step). + + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + const float mv[3] = { p.particle_mass * v.x, p.particle_mass * v.y, + p.particle_mass * v.z }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float dpz = ((float)c - fz) * p.dx; + const float w = wxy * wz[c]; + + const int lx = ix - obx, ly = iy - oby, lz = iz - obz; + const bool local = (unsigned)lx < MPM_TILE && (unsigned)ly < MPM_TILE && + (unsigned)lz < MPM_TILE; + float* cell = local + ? tile + 4 * ((lx * MPM_TILE + ly) * MPM_TILE + lz) + : grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + + #pragma unroll + for (int k = 0; k < 3; k++) { + float impulse = affine[k * 3 + 0] * dpx; + impulse = fmaf(affine[k * 3 + 1], dpy, impulse); + impulse = fmaf(affine[k * 3 + 2], dpz, impulse); + atomicAdd(cell + k, w * (mv[k] + impulse)); + } + atomicAdd(cell + 3, w * p.particle_mass); + } + } + } + } + + __syncthreads(); + + // flush the tile into the global grid + for (int c = threadIdx.x; c < MPM_TILE_CELLS; c += P2G_BLOCK) { + const int lz = c % MPM_TILE; + const int ly = (c / MPM_TILE) % MPM_TILE; + const int lx = c / (MPM_TILE * MPM_TILE); + const int ix = obx + lx, iy = oby + ly, iz = obz + lz; + if ((unsigned)ix >= MPM_GRID || (unsigned)iy >= MPM_GRID || + (unsigned)iz >= MPM_GRID) + continue; + + const float4 acc = *(const float4*)(tile + 4 * c); + if (acc.x == 0.0f && acc.y == 0.0f && acc.z == 0.0f && acc.w == 0.0f) continue; + + float* dst = grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + atomicAdd(dst + 0, acc.x); + atomicAdd(dst + 1, acc.y); + atomicAdd(dst + 2, acc.z); + atomicAdd(dst + 3, acc.w); + } +} + +__global__ void __launch_bounds__(GRID_BLOCK) +mpm_grid_kernel(MpmParams p, float4* __restrict__ grid) +{ + const int cell = blockIdx.x * blockDim.x + threadIdx.x; + if (cell >= MPM_CELLS) return; + + float4 g = grid[cell]; + if (g.w <= 0.0f) { + grid[cell] = make_float4(0.0f, 0.0f, 0.0f, g.w); + return; + } + + const float inv_mass = 1.0f / g.w; + float vx = g.x * inv_mass; + float vy = fmaf(p.dt, p.gravity, g.y * inv_mass); + float vz = g.z * inv_mass; + + const int iz = cell % MPM_GRID; + const int iy = (cell / MPM_GRID) % MPM_GRID; + const int ix = cell / (MPM_GRID * MPM_GRID); + + if (ix < p.boundary && vx < 0.0f) vx = 0.0f; + if (ix >= MPM_GRID - p.boundary && vx > 0.0f) vx = 0.0f; + if (iy < p.boundary && vy < 0.0f) vy = 0.0f; + if (iy >= MPM_GRID - p.boundary && vy > 0.0f) vy = 0.0f; + if (iz < p.boundary && vz < 0.0f) vz = 0.0f; + if (iz >= MPM_GRID - p.boundary && vz > 0.0f) vz = 0.0f; + + grid[cell] = make_float4(vx, vy, vz, g.w); +} + +__global__ void __launch_bounds__(P2G_BLOCK) +mpm_g2p_kernel(int n, MpmParams p, + float4* __restrict__ mean, + float4* __restrict__ velocity, + float* __restrict__ affine_out, + float* __restrict__ defgrad, + const float4* __restrict__ grid) +{ + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + const float4 x = mean[i]; + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + float nv[3] = { 0.0f, 0.0f, 0.0f }; + float nC[9] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float w = wxy * wz[c]; + const float4 g = grid[(size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz]; + const float gv[3] = { g.x, g.y, g.z }; + const float dpos[3] = { dpx, dpy, ((float)c - fz) * p.dx }; + + #pragma unroll + for (int k = 0; k < 3; k++) { + nv[k] = fmaf(w, gv[k], nv[k]); + const float wg = 4.0f * p.inv_dx * p.inv_dx * w * gv[k]; + #pragma unroll + for (int l = 0; l < 3; l++) nC[k * 3 + l] = fmaf(wg, dpos[l], nC[k * 3 + l]); + } + } + } + } + + float F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = fmaf(p.dt * nC[a * 3 + k], F0[k * 3 + b], acc); + defgrad[(size_t)(a * 3 + b) * n + i] = acc; + } + + #pragma unroll + for (int k = 0; k < 9; k++) affine_out[(size_t)k * n + i] = nC[k]; + + velocity[i] = make_float4(nv[0], nv[1], nv[2], 0.0f); + mean[i] = make_float4(fmaf(p.dt, nv[0], x.x), fmaf(p.dt, nv[1], x.y), + fmaf(p.dt, nv[2], x.z), x.w); +} + +// --------------------------------------------------------------------------- +// stage 2: condition the 4D Gaussian on t, project, shade +// +// The frustum and opacity tests come before the 48 spherical harmonic loads, +// so a culled Gaussian costs no bandwidth beyond its pose. +// --------------------------------------------------------------------------- + +__global__ void __launch_bounds__(PREPROCESS_BLOCK) +preprocess_kernel(int n, float time, Camera cam, + const float4* __restrict__ mean4, + const float4* __restrict__ scale4, + const float4* __restrict__ quat_l, + const float4* __restrict__ quat_r, + const float* __restrict__ opacity_in, + const float* __restrict__ sh, + float2* __restrict__ mean2d, + float4* __restrict__ conic_opacity, + float4* __restrict__ color_depth, + int* __restrict__ radii) +{ + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + radii[i] = 0; + + const float4 mu = mean4[i]; + const float4 sc = scale4[i]; + const float4 ql = quat_l[i]; + const float4 qr = quat_r[i]; + + // M = L(ql) R(qr), the 4D rotation as a pair of isoclinic rotations + const float lw = ql.x, lx = ql.y, ly = ql.z, lz = ql.w; + const float rw = qr.x, rx = qr.y, ry = qr.z, rz = qr.w; + + const float L[16] = { + lw, -lx, -ly, -lz, + lx, lw, -lz, ly, + ly, lz, lw, -lx, + lz, -ly, lx, lw }; + const float R[16] = { + rw, -rx, -ry, -rz, + rx, rw, rz, -ry, + ry, -rz, rw, rx, + rz, ry, -rx, rw }; + + float M[16]; + #pragma unroll + for (int a = 0; a < 4; a++) + #pragma unroll + for (int b = 0; b < 4; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) acc = fmaf(L[a * 4 + k], R[k * 4 + b], acc); + M[a * 4 + b] = acc; + } + + // the quaternion basis is ordered (t, x, y, z) + const float s2[4] = { sc.w * sc.w, sc.x * sc.x, sc.y * sc.y, sc.z * sc.z }; + + // only the 3x3 spatial block, the spatio-temporal column and Sigma_tt are + // needed, so the remaining entries of the 4x4 product are never formed + float st[3], sxyz[6]; + #pragma unroll + for (int a = 0; a < 3; a++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) acc = fmaf(M[(a + 1) * 4 + k] * s2[k], M[k], acc); + st[a] = acc; + } + float sigma_tt = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) sigma_tt = fmaf(M[k] * s2[k], M[k], sigma_tt); + + int idx = 0; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = a; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) + acc = fmaf(M[(a + 1) * 4 + k] * s2[k], M[(b + 1) * 4 + k], acc); + sxyz[idx++] = acc; + } + + const float inv_tt = 1.0f / sigma_tt; + const float dt = time - mu.w; + + const float alpha = opacity_in[i] * __expf(-0.5f * dt * dt * inv_tt); + if (alpha < MIN_OPACITY) return; + + const float mean3[3] = { fmaf(dt * inv_tt, st[0], mu.x), + fmaf(dt * inv_tt, st[1], mu.y), + fmaf(dt * inv_tt, st[2], mu.z) }; + const float cov3[6] = { sxyz[0] - st[0] * st[0] * inv_tt, + sxyz[1] - st[0] * st[1] * inv_tt, + sxyz[2] - st[0] * st[2] * inv_tt, + sxyz[3] - st[1] * st[1] * inv_tt, + sxyz[4] - st[1] * st[2] * inv_tt, + sxyz[5] - st[2] * st[2] * inv_tt }; + + // world -> camera + const float px = mean3[0] - cam.cam_pos[0]; + const float py = mean3[1] - cam.cam_pos[1]; + const float pz = mean3[2] - cam.cam_pos[2]; + float t[3]; + #pragma unroll + for (int a = 0; a < 3; a++) + t[a] = fmaf(cam.view[a * 3 + 0], px, + fmaf(cam.view[a * 3 + 1], py, cam.view[a * 3 + 2] * pz)); + + if (t[2] < cam.near_plane) return; + + const float inv_z = 1.0f / t[2]; + const float inv_z2 = inv_z * inv_z; + + const float lim_x = 1.3f * (0.5f * cam.width) / cam.focal_x; + const float lim_y = 1.3f * (0.5f * cam.height) / cam.focal_y; + const float tx = t[2] * fminf(lim_x, fmaxf(-lim_x, t[0] * inv_z)); + const float ty = t[2] * fminf(lim_y, fmaxf(-lim_y, t[1] * inv_z)); + + const float J[6] = { cam.focal_x * inv_z, 0.0f, -cam.focal_x * tx * inv_z2, + 0.0f, cam.focal_y * inv_z, -cam.focal_y * ty * inv_z2 }; + + float T[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + T[a * 3 + b] = fmaf(J[a * 3 + 0], cam.view[0 * 3 + b], + fmaf(J[a * 3 + 1], cam.view[1 * 3 + b], + J[a * 3 + 2] * cam.view[2 * 3 + b])); + + const float S[9] = { cov3[0], cov3[1], cov3[2], + cov3[1], cov3[3], cov3[4], + cov3[2], cov3[4], cov3[5] }; + float TS[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + TS[a * 3 + b] = fmaf(T[a * 3 + 0], S[0 * 3 + b], + fmaf(T[a * 3 + 1], S[1 * 3 + b], + T[a * 3 + 2] * S[2 * 3 + b])); + + float ca = 0.3f, cb = 0.0f, cc = 0.3f; + #pragma unroll + for (int k = 0; k < 3; k++) { + ca = fmaf(TS[0 * 3 + k], T[0 * 3 + k], ca); + cb = fmaf(TS[0 * 3 + k], T[1 * 3 + k], cb); + cc = fmaf(TS[1 * 3 + k], T[1 * 3 + k], cc); + } + + const float det = fmaf(ca, cc, -cb * cb); + if (det <= 0.0f) return; + const float inv_det = 1.0f / det; + + const float mid = 0.5f * (ca + cc); + const float disc = sqrtf(fmaxf(0.1f, fmaf(mid, mid, -det))); + const float radius = ceilf(3.0f * sqrtf(fmaxf(mid + disc, mid - disc))); + + const float mx = fmaf(cam.focal_x * inv_z, t[0], cam.center_x); + const float my = fmaf(cam.focal_y * inv_z, t[1], cam.center_y); + + if (mx + radius < 0.0f || mx - radius > (float)cam.width || + my + radius < 0.0f || my - radius > (float)cam.height) + return; + + // spherical harmonics, evaluated only for the Gaussians that survive + const float inv_len = rsqrtf(fmaf(px, px, fmaf(py, py, pz * pz))); + const float dx = px * inv_len, dy = py * inv_len, dz = pz * inv_len; + + float basis[SH_COEFFS]; + basis[0] = SH_C0; + basis[1] = -SH_C1 * dy; + basis[2] = SH_C1 * dz; + basis[3] = -SH_C1 * dx; + + const float xx = dx * dx, yy = dy * dy, zz = dz * dz; + const float xy = dx * dy, yz = dy * dz, xz = dx * dz; + basis[4] = SH_C2_0 * xy; + basis[5] = SH_C2_1 * yz; + basis[6] = SH_C2_2 * (2.0f * zz - xx - yy); + basis[7] = SH_C2_3 * xz; + basis[8] = SH_C2_4 * (xx - yy); + basis[9] = SH_C3_0 * dy * (3.0f * xx - yy); + basis[10] = SH_C3_1 * xy * dz; + basis[11] = SH_C3_2 * dy * (4.0f * zz - xx - yy); + basis[12] = SH_C3_3 * dz * (2.0f * zz - 3.0f * xx - 3.0f * yy); + basis[13] = SH_C3_4 * dx * (4.0f * zz - xx - yy); + basis[14] = SH_C3_5 * dz * (xx - yy); + basis[15] = SH_C3_6 * dx * (xx - 3.0f * yy); + + float rgb[3] = { 0.0f, 0.0f, 0.0f }; + #pragma unroll + for (int c = 0; c < SH_COEFFS; c++) { + const float w = basis[c]; + #pragma unroll + for (int ch = 0; ch < 3; ch++) + rgb[ch] = fmaf(w, sh[((size_t)c * 3 + ch) * n + i], rgb[ch]); + } + + radii[i] = (int)radius; + mean2d[i] = make_float2(mx, my); + conic_opacity[i] = make_float4(cc * inv_det, -cb * inv_det, ca * inv_det, alpha); + color_depth[i] = make_float4(fmaxf(rgb[0] + 0.5f, 0.0f), + fmaxf(rgb[1] + 0.5f, 0.0f), + fmaxf(rgb[2] + 0.5f, 0.0f), t[2]); +} + +// --------------------------------------------------------------------------- +// stage 3: tile rasterizer +// +// One block owns one 16x16 tile. The block cooperatively stages a batch of +// BLOCK_SIZE Gaussians in shared memory, so each Gaussian is read from global +// memory once per tile rather than once per pixel, and the whole block leaves +// as soon as all of its pixels are saturated. +// --------------------------------------------------------------------------- + +__global__ void __launch_bounds__(BLOCK_SIZE) +render_kernel(Camera cam, + const float2* __restrict__ mean2d, + const float4* __restrict__ conic_opacity, + const float4* __restrict__ color_depth, + const int* __restrict__ tile_offsets, + const int* __restrict__ tile_list, + float4* __restrict__ image) +{ + __shared__ float2 s_xy[BLOCK_SIZE]; + __shared__ float4 s_co[BLOCK_SIZE]; + __shared__ float4 s_color[BLOCK_SIZE]; + + const int tile = blockIdx.y * cam.tiles_x + blockIdx.x; + const int x = blockIdx.x * BLOCK_X + threadIdx.x; + const int y = blockIdx.y * BLOCK_Y + threadIdx.y; + const int lane = threadIdx.y * BLOCK_X + threadIdx.x; + + const bool inside = (x < cam.width) && (y < cam.height); + bool done = !inside; + + const float pixf_x = (float)x + 0.5f; + const float pixf_y = (float)y + 0.5f; + + const int begin = tile_offsets[tile]; + const int end = tile_offsets[tile + 1]; + const int rounds = (end - begin + BLOCK_SIZE - 1) / BLOCK_SIZE; + + float transmittance = 1.0f; + float r = 0.0f, g = 0.0f, b = 0.0f; + + int todo = end - begin; + for (int round = 0; round < rounds; round++, todo -= BLOCK_SIZE) { + // every pixel of the tile is saturated, so no further Gaussian can + // contribute to this block + if (__syncthreads_count(done) == BLOCK_SIZE) break; + + const int fetch = begin + round * BLOCK_SIZE + lane; + if (fetch < end) { + const int gid = tile_list[fetch]; + s_xy[lane] = mean2d[gid]; + s_co[lane] = conic_opacity[gid]; + s_color[lane] = color_depth[gid]; + } + __syncthreads(); + + const int count = min(BLOCK_SIZE, todo); + for (int j = 0; j < count && !done; j++) { + const float2 xy = s_xy[j]; + const float4 co = s_co[j]; + const float dx = xy.x - pixf_x; + const float dy = xy.y - pixf_y; + + const float power = -0.5f * fmaf(co.x, dx * dx, co.z * dy * dy) - co.y * dx * dy; + if (power > 0.0f) continue; + + const float alpha = fminf(0.99f, co.w * __expf(power)); + if (alpha < MIN_OPACITY) continue; + + const float weight = alpha * transmittance; + const float4 c = s_color[j]; + r = fmaf(c.x, weight, r); + g = fmaf(c.y, weight, g); + b = fmaf(c.z, weight, b); + + transmittance *= 1.0f - alpha; + if (transmittance < MIN_TRANSMITTANCE) done = true; + } + __syncthreads(); + } + + if (inside) + image[(size_t)y * cam.width + x] = make_float4(r, g, b, 1.0f - transmittance); +} + +// --------------------------------------------------------------------------- + +int main(int argc, char* argv[]) +{ + if (argc != 5) { + printf("Usage: %s \n", + argv[0]); + return 1; + } + + const int n = atoi(argv[1]); + const int width = atoi(argv[2]); + const int height = atoi(argv[3]); + const int repeat = atoi(argv[4]); + + if (n <= 0 || width <= 0 || height <= 0 || repeat <= 0) { + printf("Error: number of gaussians, image width, image height, and repeat " + "must all be positive integers (got n=%d, width=%d, height=%d, " + "repeat=%d)\n", n, width, height, repeat); + return 1; + } + + Camera cam; + setup_camera(width, height, cam); + MpmParams mpm; + setup_mpm(mpm); + const float time = 0.5f; // the instant of the 4D scene that is rendered + + Scene scene; + std::vector h_image, h_ref_image; + std::vector h_mean2d, h_conic, h_color; + std::vector h_radii; + std::vector tile_offsets, tile_list; + try { + generate_scene(n, scene); + h_mean2d.resize((size_t)2 * n); + h_conic.resize((size_t)4 * n); + h_color.resize((size_t)4 * n); + h_radii.resize(n); + h_image.resize(4 * (size_t)width * height); + h_ref_image.resize(4 * (size_t)width * height); + } catch (const std::bad_alloc&) { + printf("Failed to allocate the host buffers for %d gaussians and a %d x %d " + "image\n", n, width, height); + return 1; + } + + const int num_tiles = cam.tiles_x * cam.tiles_y; + printf("Gaussians: %d, image: %d x %d (%d tiles), MPM grid: %d^3\n", + n, width, height, num_tiles, MPM_GRID); + + // the host reference evolves its own copy of the scene + Scene ref_scene = scene; + + float4 *d_mean, *d_scale, *d_quat_l, *d_quat_r, *d_velocity; + float *d_opacity, *d_sh, *d_affine, *d_defgrad; + float4 *d_grid; + CHECK(hipMalloc((void**)&d_mean, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_scale, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_quat_l, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_quat_r, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_velocity, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_opacity, sizeof(float) * n)); + CHECK(hipMalloc((void**)&d_sh, sizeof(float) * SH_COEFFS * 3 * (size_t)n)); + CHECK(hipMalloc((void**)&d_affine, sizeof(float) * 9 * (size_t)n)); + CHECK(hipMalloc((void**)&d_defgrad, sizeof(float) * 9 * (size_t)n)); + CHECK(hipMalloc((void**)&d_grid, sizeof(float4) * MPM_CELLS)); + + const int num_chunks = (int)scene.chunk_block.size(); + int *d_chunk_start, *d_chunk_block; + CHECK(hipMalloc((void**)&d_chunk_start, sizeof(int) * (num_chunks + 1))); + CHECK(hipMalloc((void**)&d_chunk_block, sizeof(int) * num_chunks)); + CHECK(hipMemcpy(d_chunk_start, scene.chunk_start.data(), + sizeof(int) * (num_chunks + 1), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(d_chunk_block, scene.chunk_block.data(), + sizeof(int) * num_chunks, hipMemcpyHostToDevice)); + + // the host keeps xyzw and xyz layouts, the device wants 16 byte vectors + std::vector pack(4 * (size_t)n); + auto upload_vec4 = [&](const std::vector& src, int comps, float4* dst) { + for (int i = 0; i < n; i++) { + pack[4 * (size_t)i + 0] = src[(size_t)comps * i + 0]; + pack[4 * (size_t)i + 1] = src[(size_t)comps * i + 1]; + pack[4 * (size_t)i + 2] = src[(size_t)comps * i + 2]; + pack[4 * (size_t)i + 3] = (comps == 4) ? src[(size_t)comps * i + 3] : 0.0f; + } + CHECK(hipMemcpy(dst, pack.data(), sizeof(float4) * n, hipMemcpyHostToDevice)); + }; + + // the 3x3 tensors are stored component major on the device + std::vector pack9(9 * (size_t)n); + auto upload_tensor = [&](const std::vector& src, float* dst) { + for (int k = 0; k < 9; k++) + for (int i = 0; i < n; i++) pack9[(size_t)k * n + i] = src[9 * (size_t)i + k]; + CHECK(hipMemcpy(dst, pack9.data(), sizeof(float) * 9 * (size_t)n, + hipMemcpyHostToDevice)); + }; + + upload_vec4(scene.mean, 4, d_mean); + upload_vec4(scene.scale, 4, d_scale); + upload_vec4(scene.quat_l, 4, d_quat_l); + upload_vec4(scene.quat_r, 4, d_quat_r); + upload_vec4(scene.velocity, 3, d_velocity); + upload_tensor(scene.affine, d_affine); + upload_tensor(scene.defgrad, d_defgrad); + CHECK(hipMemcpy(d_opacity, scene.opacity.data(), sizeof(float) * n, + hipMemcpyHostToDevice)); + CHECK(hipMemcpy(d_sh, scene.sh.data(), sizeof(float) * SH_COEFFS * 3 * (size_t)n, + hipMemcpyHostToDevice)); + + float2 *d_mean2d; + float4 *d_conic, *d_color, *d_image; + int *d_radii, *d_tile_offsets, *d_tile_list; + CHECK(hipMalloc((void**)&d_mean2d, sizeof(float2) * n)); + CHECK(hipMalloc((void**)&d_conic, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_color, sizeof(float4) * n)); + CHECK(hipMalloc((void**)&d_radii, sizeof(int) * n)); + CHECK(hipMalloc((void**)&d_image, sizeof(float4) * (size_t)width * height)); + CHECK(hipMalloc((void**)&d_tile_offsets, sizeof(int) * (num_tiles + 1))); + + const int particle_blocks = (n + P2G_BLOCK - 1) / P2G_BLOCK; + const int cell_blocks = (MPM_CELLS + GRID_BLOCK - 1) / GRID_BLOCK; + const int preprocess_blocks = (n + PREPROCESS_BLOCK - 1) / PREPROCESS_BLOCK; + + auto mpm_step = [&]() { + CHECK(hipMemset(d_grid, 0, sizeof(float4) * MPM_CELLS)); + mpm_p2g_kernel<<>>( + n, mpm, d_chunk_start, d_chunk_block, d_mean, d_velocity, d_affine, + d_defgrad, (float*)d_grid); + mpm_grid_kernel<<>>(mpm, d_grid); + mpm_g2p_kernel<<>>( + n, mpm, d_mean, d_velocity, d_affine, d_defgrad, d_grid); + }; + + // --- stage 1: one MPM step, verified against the reference --------------- + mpm_step(); + CHECK(hipGetLastError()); + + { + std::vector ref_grid(4 * (size_t)MPM_CELLS); + reference_p2g(ref_scene, mpm, ref_grid.data()); + reference_grid_update(mpm, ref_grid.data()); + reference_g2p(ref_scene, mpm, ref_grid.data()); + + std::vector got(4 * (size_t)n); + CHECK(hipMemcpy(got.data(), d_mean, sizeof(float4) * n, hipMemcpyDeviceToHost)); + int mpm_errors = 0; + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(got[4 * (size_t)i + k], ref_scene.mean[4 * (size_t)i + k], 1e-4f)) + mpm_errors++; + CHECK(hipMemcpy(got.data(), d_velocity, sizeof(float4) * n, hipMemcpyDeviceToHost)); + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(got[4 * (size_t)i + k], ref_scene.velocity[3 * (size_t)i + k], 1e-3f)) + mpm_errors++; + printf("MPM step: %s\n", mpm_errors == 0 ? "PASS" : "FAIL"); + } + + CHECK(hipDeviceSynchronize()); + auto start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) mpm_step(); + CHECK(hipDeviceSynchronize()); + auto end = std::chrono::steady_clock::now(); + auto time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the MPM step (p2g, grid, g2p): %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 2: the 4D preprocess, from the undeformed scene --------------- + upload_vec4(scene.mean, 4, d_mean); + + preprocess_kernel<<>>( + n, time, cam, d_mean, d_scale, d_quat_l, d_quat_r, d_opacity, d_sh, + d_mean2d, d_conic, d_color, d_radii); + CHECK(hipGetLastError()); + + std::vector ref_mean2d(2 * (size_t)n), ref_conic(4 * (size_t)n), + ref_color(4 * (size_t)n); + { + std::vector ref_radii(n); + reference_preprocess(scene, cam, time, ref_mean2d.data(), ref_conic.data(), + ref_color.data(), ref_radii.data()); + + CHECK(hipMemcpy(h_mean2d.data(), d_mean2d, sizeof(float2) * n, hipMemcpyDeviceToHost)); + CHECK(hipMemcpy(h_conic.data(), d_conic, sizeof(float4) * n, hipMemcpyDeviceToHost)); + CHECK(hipMemcpy(h_color.data(), d_color, sizeof(float4) * n, hipMemcpyDeviceToHost)); + CHECK(hipMemcpy(h_radii.data(), d_radii, sizeof(int) * n, hipMemcpyDeviceToHost)); + + int pre_errors = 0; + int visible = 0; + for (int i = 0; i < n && pre_errors == 0; i++) { + if (ref_radii[i] > 0) visible++; + // the extent is the ceiling of a float, so it may land either side of + // an integer when the host and the device contract differently + if (abs(h_radii[i] - ref_radii[i]) > 1) pre_errors++; + if (ref_radii[i] == 0 || h_radii[i] == 0) continue; + for (int k = 0; k < 2; k++) + if (!close_enough(h_mean2d[2 * (size_t)i + k], ref_mean2d[2 * (size_t)i + k], 1e-3f)) + pre_errors++; + for (int k = 0; k < 4; k++) { + if (!close_enough(h_conic[4 * (size_t)i + k], ref_conic[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + if (!close_enough(h_color[4 * (size_t)i + k], ref_color[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + } + } + printf("4D preprocess (%d of %d gaussians visible): %s\n", visible, n, + pre_errors == 0 ? "PASS" : "FAIL"); + + // The tile lists are host side setup, shared by the reference and the + // device rasterizer. Stage 3 is fed the reference splats, so that it is + // verified on its own rather than against the rounding of stage 2. + build_tile_lists(cam, n, ref_mean2d.data(), ref_color.data(), + ref_radii.data(), tile_offsets, tile_list); + reference_render(cam, ref_mean2d.data(), ref_conic.data(), ref_color.data(), + tile_offsets.data(), tile_list.data(), h_ref_image.data()); + } + + CHECK(hipDeviceSynchronize()); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) { + preprocess_kernel<<>>( + n, time, cam, d_mean, d_scale, d_quat_l, d_quat_r, d_opacity, d_sh, + d_mean2d, d_conic, d_color, d_radii); + } + CHECK(hipDeviceSynchronize()); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the 4D preprocess kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 3: the tile rasterizer --------------------------------------- + CHECK(hipMemcpy(d_mean2d, ref_mean2d.data(), sizeof(float2) * n, + hipMemcpyHostToDevice)); + CHECK(hipMemcpy(d_conic, ref_conic.data(), sizeof(float4) * n, + hipMemcpyHostToDevice)); + CHECK(hipMemcpy(d_color, ref_color.data(), sizeof(float4) * n, + hipMemcpyHostToDevice)); + + const size_t list_size = tile_list.empty() ? 1 : tile_list.size(); + CHECK(hipMalloc((void**)&d_tile_list, sizeof(int) * list_size)); + CHECK(hipMemcpy(d_tile_offsets, tile_offsets.data(), sizeof(int) * (num_tiles + 1), + hipMemcpyHostToDevice)); + if (!tile_list.empty()) + CHECK(hipMemcpy(d_tile_list, tile_list.data(), sizeof(int) * tile_list.size(), + hipMemcpyHostToDevice)); + + printf("Gaussian instances after tiling: %zu (%.1f per tile)\n", + tile_list.size(), (double)tile_list.size() / num_tiles); + + const dim3 render_block(BLOCK_X, BLOCK_Y); + const dim3 render_grid(cam.tiles_x, cam.tiles_y); + + CHECK(hipMemset(d_image, 0, sizeof(float4) * (size_t)width * height)); + render_kernel<<>>( + cam, d_mean2d, d_conic, d_color, d_tile_offsets, d_tile_list, d_image); + CHECK(hipGetLastError()); + CHECK(hipMemcpy(h_image.data(), d_image, sizeof(float4) * (size_t)width * height, + hipMemcpyDeviceToHost)); + + int render_errors = 0; + for (size_t k = 0; k < h_image.size() && render_errors == 0; k++) + if (!close_enough(h_image[k], h_ref_image[k], 1e-3f)) render_errors++; + printf("Rasterizer: %s\n", render_errors == 0 ? "PASS" : "FAIL"); + + CHECK(hipDeviceSynchronize()); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) { + render_kernel<<>>( + cam, d_mean2d, d_conic, d_color, d_tile_offsets, d_tile_list, d_image); + } + CHECK(hipDeviceSynchronize()); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the rasterizer kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + CHECK(hipFree(d_mean)); CHECK(hipFree(d_scale)); + CHECK(hipFree(d_quat_l)); CHECK(hipFree(d_quat_r)); + CHECK(hipFree(d_velocity)); CHECK(hipFree(d_opacity)); + CHECK(hipFree(d_sh)); CHECK(hipFree(d_affine)); CHECK(hipFree(d_defgrad)); + CHECK(hipFree(d_grid)); CHECK(hipFree(d_mean2d)); CHECK(hipFree(d_conic)); + CHECK(hipFree(d_color)); CHECK(hipFree(d_radii)); CHECK(hipFree(d_image)); + CHECK(hipFree(d_tile_offsets)); CHECK(hipFree(d_tile_list)); + + return 0; +} diff --git a/src/gsplat4d-omp/CMakeLists.txt b/src/gsplat4d-omp/CMakeLists.txt new file mode 100644 index 000000000..61998fdee --- /dev/null +++ b/src/gsplat4d-omp/CMakeLists.txt @@ -0,0 +1,11 @@ +# gsplat4d-omp/CMakeLists.txt + +add_hecbench_benchmark( + NAME gsplat4d + MODEL omp + SOURCES main.cpp + CATEGORIES simulation + TEST_ARGS 200000 1280 720 100 + TEST_REGEX "(?s)MPM step: PASS.*4D preprocess \([0-9]+ of [0-9]+ gaussians visible\): PASS.*Rasterizer: PASS" + TEST_TIMEOUT 300 +) diff --git a/src/gsplat4d-omp/Makefile b/src/gsplat4d-omp/Makefile new file mode 100644 index 000000000..a5de43672 --- /dev/null +++ b/src/gsplat4d-omp/Makefile @@ -0,0 +1,62 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = icpx +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../gsplat4d-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS +=-fiopenmp -fopenmp-targets=spir64 -D__STRICT_ANSI__ +else + CFLAGS +=-qopenmp +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../gsplat4d-cuda/reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-omp/Makefile.aomp b/src/gsplat4d-omp/Makefile.aomp new file mode 100644 index 000000000..1c9ea3ef1 --- /dev/null +++ b/src/gsplat4d-omp/Makefile.aomp @@ -0,0 +1,66 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = clang++ +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +ARCH = gfx906 +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../gsplat4d-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS += -target x86_64-pc-linux-gnu \ + -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa \ + -Xopenmp-target=amdgcn-amd-amdhsa \ + -march=$(ARCH) +else + CFLAGS +=-fopenmp +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile.aomp + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../gsplat4d-cuda/reference.h Makefile.aomp + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-omp/Makefile.nvc b/src/gsplat4d-omp/Makefile.nvc new file mode 100644 index 000000000..833091ae6 --- /dev/null +++ b/src/gsplat4d-omp/Makefile.nvc @@ -0,0 +1,63 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = nvc++ +OPTIMIZE = yes +DEBUG = no +DEVICE = gpu +SM ?= cc70 +LAUNCHER ?= + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../gsplat4d-cuda + +# Linker Flags +LDFLAGS = + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(DEVICE),gpu) + CFLAGS +=-Minfo -mp=gpu -gpu=$(SM)#,fastmath +else + CFLAGS += +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile.nvc + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../gsplat4d-cuda/reference.h Makefile.nvc + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-omp/main.cpp b/src/gsplat4d-omp/main.cpp new file mode 100644 index 000000000..bf6f5d7ff --- /dev/null +++ b/src/gsplat4d-omp/main.cpp @@ -0,0 +1,854 @@ +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#define GRID_BLOCK 256 +#define P2G_BLOCK 128 +#define PREPROCESS_BLOCK 128 + +#pragma omp declare target +static inline void quad_weights(float fx, float& w0, float& w1, float& w2) +{ + const float a = 1.5f - fx; + const float b = fx - 1.0f; + const float c = fx - 0.5f; + w0 = 0.5f * a * a; + w1 = 0.75f - b * b; + w2 = 0.5f * c * c; +} +#pragma omp end declare target + +// --------------------------------------------------------------------------- +// stage 1: MLS-MPM +// +// One team scatters one chunk of particles, all of which belong to the same +// MPM_BLOCK^3 cell block, so the whole stencil footprint fits the team local +// tile below and the scatter only reaches the global grid at the flush. A +// particle that has drifted out of its block since the binning still lands +// correctly through the global fallback. +// --------------------------------------------------------------------------- + +static void mpm_p2g(int n, int num_chunks, + const int* chunk_start, const int* chunk_block, + const float* mean, const float* velocity, + const float* affine, const float* defgrad, float* grid, + float dt, float dx, float inv_dx, float p_vol, float p_mass, + float mu_lame, float lambda_lame) +{ + #pragma omp target teams num_teams(num_chunks) thread_limit(P2G_BLOCK) + { + float tile[MPM_TILE_CELLS * 4]; + #pragma omp parallel num_threads(P2G_BLOCK) + { + const int team = omp_get_team_num(); + const int nteams = omp_get_num_teams(); + const int lid = omp_get_thread_num(); + const int nthreads = omp_get_num_threads(); + + // The runtime may create fewer than num_chunks teams, so each team walks + // a strided range of chunks rather than assuming a one-to-one mapping. + // Every thread of a team runs the same iteration count, so the team + // barriers below stay aligned. + for (int chunk = team; chunk < num_chunks; chunk += nteams) { + for (int k = lid; k < MPM_TILE_CELLS * 4; k += nthreads) tile[k] = 0.0f; + + const int begin = chunk_start[chunk]; + const int end = chunk_start[chunk + 1]; + const int gb = chunk_block[chunk]; + const int obz = (gb % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int oby = ((gb / MPM_BLOCKS_PER_DIM) % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int obx = (gb / (MPM_BLOCKS_PER_DIM * MPM_BLOCKS_PER_DIM)) * MPM_BLOCK - 1; + + #pragma omp barrier + + for (int i = begin + lid; i < end; i += nthreads) { + float C[9], F0[9], F[9]; + for (int k = 0; k < 9; k++) C[k] = affine[(size_t)k * n + i]; + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + for (int k = 0; k < 3; k++) acc += dt * C[a * 3 + k] * F0[k * 3 + b]; + F[a * 3 + b] = acc; + } + + const float det = + F[0] * (F[4] * F[8] - F[5] * F[7]) - + F[1] * (F[3] * F[8] - F[5] * F[6]) + + F[2] * (F[3] * F[7] - F[4] * F[6]); + const float safe_J = + (fabsf(det) < 1e-6f) ? ((det < 0.0f) ? -1e-6f : 1e-6f) : det; + const float inv_J = 1.0f / safe_J; + + const float cof[9] = { + (F[4] * F[8] - F[5] * F[7]), -(F[3] * F[8] - F[5] * F[6]), (F[3] * F[7] - F[4] * F[6]), + -(F[1] * F[8] - F[2] * F[7]), (F[0] * F[8] - F[2] * F[6]), -(F[0] * F[7] - F[1] * F[6]), + (F[1] * F[5] - F[2] * F[4]), -(F[0] * F[5] - F[2] * F[3]), (F[0] * F[4] - F[1] * F[3]) }; + + const float coeff = (lambda_lame * logf(fabsf(safe_J)) - mu_lame) * inv_J; + float P[9]; + for (int k = 0; k < 9; k++) P[k] = mu_lame * F[k] + coeff * cof[k]; + + const float s = -dt * p_vol * 4.0f * inv_dx * inv_dx; + float aff[9]; + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) { + float acc = 0.0f; + for (int k = 0; k < 3; k++) acc += P[a * 3 + k] * F[b * 3 + k]; + aff[a * 3 + b] = s * acc + p_mass * C[a * 3 + b]; + } + + // F is only used for the stress above; it is deliberately not + // persisted here. The single per-step deformation-gradient update is + // applied once, in mpm_g2p, from the original F0 (matching the host + // reference; writing F back here would advance the gradient twice). + + const float gx = mean[4 * (size_t)i + 0] * inv_dx; + const float gy = mean[4 * (size_t)i + 1] * inv_dx; + const float gz = mean[4 * (size_t)i + 2] * inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + const float mv[3] = { p_mass * velocity[4 * (size_t)i + 0], + p_mass * velocity[4 * (size_t)i + 1], + p_mass * velocity[4 * (size_t)i + 2] }; + + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * dx; + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * dx; + const float wxy = wx[a] * wy[b]; + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float dpz = ((float)c - fz) * dx; + const float w = wxy * wz[c]; + + float val[4]; + for (int k = 0; k < 3; k++) + val[k] = w * (mv[k] + aff[k * 3 + 0] * dpx + + aff[k * 3 + 1] * dpy + aff[k * 3 + 2] * dpz); + val[3] = w * p_mass; + + const int lx = ix - obx, ly = iy - oby, lz = iz - obz; + if ((unsigned)lx < MPM_TILE && (unsigned)ly < MPM_TILE && + (unsigned)lz < MPM_TILE) { + const int o = 4 * ((lx * MPM_TILE + ly) * MPM_TILE + lz); + for (int k = 0; k < 4; k++) { + #pragma omp atomic update + tile[o + k] += val[k]; + } + } else { + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + for (int k = 0; k < 4; k++) { + #pragma omp atomic update + grid[o + k] += val[k]; + } + } + } + } + } + } + + #pragma omp barrier + + for (int c = lid; c < MPM_TILE_CELLS; c += nthreads) { + const int lz = c % MPM_TILE; + const int ly = (c / MPM_TILE) % MPM_TILE; + const int lx = c / (MPM_TILE * MPM_TILE); + const int ix = obx + lx, iy = oby + ly, iz = obz + lz; + if ((unsigned)ix >= MPM_GRID || (unsigned)iy >= MPM_GRID || + (unsigned)iz >= MPM_GRID) + continue; + if (tile[4 * c + 0] == 0.0f && tile[4 * c + 1] == 0.0f && + tile[4 * c + 2] == 0.0f && tile[4 * c + 3] == 0.0f) + continue; + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + for (int k = 0; k < 4; k++) { + #pragma omp atomic update + grid[o + k] += tile[4 * c + k]; + } + } + + // finish the flush before the next chunk re-zeros the shared tile + #pragma omp barrier + } + } + } +} + +static void mpm_grid_update(float* grid, float dt, float gravity, int boundary) +{ + #pragma omp target teams distribute parallel for num_threads(GRID_BLOCK) + for (int cell = 0; cell < MPM_CELLS; cell++) { + const size_t o = 4 * (size_t)cell; + const float mass = grid[o + 3]; + if (mass <= 0.0f) { + grid[o + 0] = 0.0f; + grid[o + 1] = 0.0f; + grid[o + 2] = 0.0f; + continue; + } + const float inv_mass = 1.0f / mass; + float vx = grid[o + 0] * inv_mass; + float vy = grid[o + 1] * inv_mass + dt * gravity; + float vz = grid[o + 2] * inv_mass; + + const int iz = cell % MPM_GRID; + const int iy = (cell / MPM_GRID) % MPM_GRID; + const int ix = cell / (MPM_GRID * MPM_GRID); + + if (ix < boundary && vx < 0.0f) vx = 0.0f; + if (ix >= MPM_GRID - boundary && vx > 0.0f) vx = 0.0f; + if (iy < boundary && vy < 0.0f) vy = 0.0f; + if (iy >= MPM_GRID - boundary && vy > 0.0f) vy = 0.0f; + if (iz < boundary && vz < 0.0f) vz = 0.0f; + if (iz >= MPM_GRID - boundary && vz > 0.0f) vz = 0.0f; + + grid[o + 0] = vx; + grid[o + 1] = vy; + grid[o + 2] = vz; + } +} + +static void mpm_g2p(int n, float* mean, float* velocity, float* affine, + float* defgrad, const float* grid, + float dt, float dx, float inv_dx) +{ + #pragma omp target teams distribute parallel for thread_limit(P2G_BLOCK) + for (int i = 0; i < n; i++) { + const float gx = mean[4 * (size_t)i + 0] * inv_dx; + const float gy = mean[4 * (size_t)i + 1] * inv_dx; + const float gz = mean[4 * (size_t)i + 2] * inv_dx; + const int bx = (int)floorf(gx - 0.5f); + const int by = (int)floorf(gy - 0.5f); + const int bz = (int)floorf(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + float nv[3] = { 0.0f, 0.0f, 0.0f }; + float nC[9] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * dx; + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * dx; + const float wxy = wx[a] * wy[b]; + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float w = wxy * wz[c]; + const size_t o = 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + const float gv[3] = { grid[o + 0], grid[o + 1], grid[o + 2] }; + const float dpos[3] = { dpx, dpy, ((float)c - fz) * dx }; + for (int k = 0; k < 3; k++) { + nv[k] += w * gv[k]; + const float wg = 4.0f * inv_dx * inv_dx * w * gv[k]; + for (int l = 0; l < 3; l++) nC[k * 3 + l] += wg * dpos[l]; + } + } + } + } + + float F0[9]; + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + for (int a = 0; a < 3; a++) + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + for (int k = 0; k < 3; k++) acc += dt * nC[a * 3 + k] * F0[k * 3 + b]; + defgrad[(size_t)(a * 3 + b) * n + i] = acc; + } + for (int k = 0; k < 9; k++) affine[(size_t)k * n + i] = nC[k]; + for (int k = 0; k < 3; k++) { + velocity[4 * (size_t)i + k] = nv[k]; + mean[4 * (size_t)i + k] += dt * nv[k]; + } + } +} + +static void mpm_step(int n, int num_chunks, const int* chunk_start, + const int* chunk_block, float* mean, float* velocity, + float* affine, float* defgrad, float* grid, + const MpmParams p) +{ + #pragma omp target teams distribute parallel for num_threads(GRID_BLOCK) + for (size_t k = 0; k < 4 * (size_t)MPM_CELLS; k++) grid[k] = 0.0f; + + mpm_p2g(n, num_chunks, chunk_start, chunk_block, mean, velocity, affine, + defgrad, grid, p.dt, p.dx, p.inv_dx, p.particle_volume, + p.particle_mass, p.mu, p.lambda); + mpm_grid_update(grid, p.dt, p.gravity, p.boundary); + mpm_g2p(n, mean, velocity, affine, defgrad, grid, p.dt, p.dx, p.inv_dx); +} + +// --------------------------------------------------------------------------- +// stage 2: condition the 4D Gaussian on t, project, shade +// --------------------------------------------------------------------------- + +static void preprocess(int n, float time, const float* mean, const float* scale, + const float* quat_l, const float* quat_r, + const float* opacity, const float* sh, + const float* view, const float* cam_pos, + float focal_x, float focal_y, float center_x, + float center_y, float near_plane, int width, int height, + float* mean2d, float* conic, float* color, int* radii) +{ + #pragma omp target teams distribute parallel for thread_limit(PREPROCESS_BLOCK) + for (int i = 0; i < n; i++) { + radii[i] = 0; + + const float lw = quat_l[4 * (size_t)i + 0], lx = quat_l[4 * (size_t)i + 1]; + const float ly = quat_l[4 * (size_t)i + 2], lz = quat_l[4 * (size_t)i + 3]; + const float rw = quat_r[4 * (size_t)i + 0], rx = quat_r[4 * (size_t)i + 1]; + const float ry = quat_r[4 * (size_t)i + 2], rz = quat_r[4 * (size_t)i + 3]; + + const float L[16] = { + lw, -lx, -ly, -lz, + lx, lw, -lz, ly, + ly, lz, lw, -lx, + lz, -ly, lx, lw }; + const float R[16] = { + rw, -rx, -ry, -rz, + rx, rw, rz, -ry, + ry, -rz, rw, rx, + rz, ry, -rx, rw }; + + float M[16]; + for (int a = 0; a < 4; a++) + for (int b = 0; b < 4; b++) { + float acc = 0.0f; + for (int k = 0; k < 4; k++) acc += L[a * 4 + k] * R[k * 4 + b]; + M[a * 4 + b] = acc; + } + + // the quaternion basis is ordered (t, x, y, z) + const float sx = scale[4 * (size_t)i + 0], sy = scale[4 * (size_t)i + 1]; + const float sz = scale[4 * (size_t)i + 2], stime = scale[4 * (size_t)i + 3]; + const float s2[4] = { stime * stime, sx * sx, sy * sy, sz * sz }; + + float st[3], sxyz[6]; + for (int a = 0; a < 3; a++) { + float acc = 0.0f; + for (int k = 0; k < 4; k++) acc += M[(a + 1) * 4 + k] * s2[k] * M[k]; + st[a] = acc; + } + float sigma_tt = 0.0f; + for (int k = 0; k < 4; k++) sigma_tt += M[k] * s2[k] * M[k]; + + int idx = 0; + for (int a = 0; a < 3; a++) + for (int b = a; b < 3; b++) { + float acc = 0.0f; + for (int k = 0; k < 4; k++) + acc += M[(a + 1) * 4 + k] * s2[k] * M[(b + 1) * 4 + k]; + sxyz[idx++] = acc; + } + + const float inv_tt = 1.0f / sigma_tt; + const float tdiff = time - mean[4 * (size_t)i + 3]; + + const float alpha = opacity[i] * expf(-0.5f * tdiff * tdiff * inv_tt); + if (alpha < MIN_OPACITY) continue; + + const float mean3[3] = { mean[4 * (size_t)i + 0] + tdiff * inv_tt * st[0], + mean[4 * (size_t)i + 1] + tdiff * inv_tt * st[1], + mean[4 * (size_t)i + 2] + tdiff * inv_tt * st[2] }; + const float cov3[6] = { sxyz[0] - st[0] * st[0] * inv_tt, + sxyz[1] - st[0] * st[1] * inv_tt, + sxyz[2] - st[0] * st[2] * inv_tt, + sxyz[3] - st[1] * st[1] * inv_tt, + sxyz[4] - st[1] * st[2] * inv_tt, + sxyz[5] - st[2] * st[2] * inv_tt }; + + const float px = mean3[0] - cam_pos[0]; + const float py = mean3[1] - cam_pos[1]; + const float pz = mean3[2] - cam_pos[2]; + float t[3]; + for (int a = 0; a < 3; a++) + t[a] = view[a * 3 + 0] * px + view[a * 3 + 1] * py + view[a * 3 + 2] * pz; + + if (t[2] < near_plane) continue; + + const float inv_z = 1.0f / t[2]; + const float inv_z2 = inv_z * inv_z; + + const float lim_x = 1.3f * (0.5f * width) / focal_x; + const float lim_y = 1.3f * (0.5f * height) / focal_y; + const float tx = t[2] * fminf(lim_x, fmaxf(-lim_x, t[0] * inv_z)); + const float ty = t[2] * fminf(lim_y, fmaxf(-lim_y, t[1] * inv_z)); + + const float J[6] = { focal_x * inv_z, 0.0f, -focal_x * tx * inv_z2, + 0.0f, focal_y * inv_z, -focal_y * ty * inv_z2 }; + + float T[6]; + for (int a = 0; a < 2; a++) + for (int b = 0; b < 3; b++) + T[a * 3 + b] = J[a * 3 + 0] * view[0 * 3 + b] + + J[a * 3 + 1] * view[1 * 3 + b] + + J[a * 3 + 2] * view[2 * 3 + b]; + + const float S[9] = { cov3[0], cov3[1], cov3[2], + cov3[1], cov3[3], cov3[4], + cov3[2], cov3[4], cov3[5] }; + float TS[6]; + for (int a = 0; a < 2; a++) + for (int b = 0; b < 3; b++) + TS[a * 3 + b] = T[a * 3 + 0] * S[0 * 3 + b] + T[a * 3 + 1] * S[1 * 3 + b] + + T[a * 3 + 2] * S[2 * 3 + b]; + + float ca = 0.3f, cb = 0.0f, cc = 0.3f; + for (int k = 0; k < 3; k++) { + ca += TS[0 * 3 + k] * T[0 * 3 + k]; + cb += TS[0 * 3 + k] * T[1 * 3 + k]; + cc += TS[1 * 3 + k] * T[1 * 3 + k]; + } + + const float det = ca * cc - cb * cb; + if (det <= 0.0f) continue; + const float inv_det = 1.0f / det; + + const float mid = 0.5f * (ca + cc); + const float disc = sqrtf(fmaxf(0.1f, mid * mid - det)); + const float radius = ceilf(3.0f * sqrtf(fmaxf(mid + disc, mid - disc))); + + const float mx = focal_x * t[0] * inv_z + center_x; + const float my = focal_y * t[1] * inv_z + center_y; + + if (mx + radius < 0.0f || mx - radius > (float)width || + my + radius < 0.0f || my - radius > (float)height) + continue; + + const float inv_len = 1.0f / sqrtf(px * px + py * py + pz * pz); + const float ddx = px * inv_len, ddy = py * inv_len, ddz = pz * inv_len; + + float basis[SH_COEFFS]; + basis[0] = SH_C0; + basis[1] = -SH_C1 * ddy; + basis[2] = SH_C1 * ddz; + basis[3] = -SH_C1 * ddx; + + const float xx = ddx * ddx, yy = ddy * ddy, zz = ddz * ddz; + const float xy = ddx * ddy, yz = ddy * ddz, xz = ddx * ddz; + basis[4] = SH_C2_0 * xy; + basis[5] = SH_C2_1 * yz; + basis[6] = SH_C2_2 * (2.0f * zz - xx - yy); + basis[7] = SH_C2_3 * xz; + basis[8] = SH_C2_4 * (xx - yy); + basis[9] = SH_C3_0 * ddy * (3.0f * xx - yy); + basis[10] = SH_C3_1 * xy * ddz; + basis[11] = SH_C3_2 * ddy * (4.0f * zz - xx - yy); + basis[12] = SH_C3_3 * ddz * (2.0f * zz - 3.0f * xx - 3.0f * yy); + basis[13] = SH_C3_4 * ddx * (4.0f * zz - xx - yy); + basis[14] = SH_C3_5 * ddz * (xx - yy); + basis[15] = SH_C3_6 * ddx * (xx - 3.0f * yy); + + float rgb[3] = { 0.0f, 0.0f, 0.0f }; + for (int c = 0; c < SH_COEFFS; c++) + for (int ch = 0; ch < 3; ch++) + rgb[ch] += basis[c] * sh[((size_t)c * 3 + ch) * n + i]; + + radii[i] = (int)radius; + mean2d[2 * (size_t)i + 0] = mx; + mean2d[2 * (size_t)i + 1] = my; + conic[4 * (size_t)i + 0] = cc * inv_det; + conic[4 * (size_t)i + 1] = -cb * inv_det; + conic[4 * (size_t)i + 2] = ca * inv_det; + conic[4 * (size_t)i + 3] = alpha; + color[4 * (size_t)i + 0] = fmaxf(rgb[0] + 0.5f, 0.0f); + color[4 * (size_t)i + 1] = fmaxf(rgb[1] + 0.5f, 0.0f); + color[4 * (size_t)i + 2] = fmaxf(rgb[2] + 0.5f, 0.0f); + color[4 * (size_t)i + 3] = t[2]; + } +} + +// --------------------------------------------------------------------------- +// stage 3: tile rasterizer +// +// One team owns one 16x16 tile and stages a batch of Gaussians in the team +// local arrays, so each Gaussian is read once per tile rather than once per +// pixel, and the team leaves as soon as all of its pixels are saturated. +// --------------------------------------------------------------------------- + +static void render(int num_tiles, int tiles_x, int width, int height, + const float* mean2d, const float* conic, const float* color, + const int* offsets, const int* list, float* image) +{ + #pragma omp target teams num_teams(num_tiles) thread_limit(BLOCK_SIZE) + { + float s_xy[2 * BLOCK_SIZE]; + float s_co[4 * BLOCK_SIZE]; + float s_color[4 * BLOCK_SIZE]; + // per-pixel accumulation state, kept in team-shared memory so that all + // team barriers live outside the per-pixel loop + float p_trans[BLOCK_SIZE]; + float p_r[BLOCK_SIZE]; + float p_g[BLOCK_SIZE]; + float p_b[BLOCK_SIZE]; + #pragma omp parallel num_threads(BLOCK_SIZE) + { + const int team = omp_get_team_num(); + const int nteams = omp_get_num_teams(); + const int lane = omp_get_thread_num(); + const int nthreads = omp_get_num_threads(); + + // The runtime may create fewer than num_tiles teams and fewer than + // BLOCK_SIZE threads, so each team strides over tiles and the round count + // is derived from the thread count. Every thread of a team executes the + // same number of rounds, so the barriers below stay aligned regardless of + // how BLOCK_SIZE divides the thread count. + for (int tile = team; tile < num_tiles; tile += nteams) { + const int tx = tile % tiles_x, ty = tile / tiles_x; + const int begin = offsets[tile]; + const int end = offsets[tile + 1]; + const int total = end - begin; + const int rounds = (total + nthreads - 1) / nthreads; + + for (int pixel = lane; pixel < BLOCK_SIZE; pixel += nthreads) { + p_trans[pixel] = 1.0f; + p_r[pixel] = 0.0f; + p_g[pixel] = 0.0f; + p_b[pixel] = 0.0f; + } + #pragma omp barrier + + for (int round = 0; round < rounds; round++) { + const int fetch = begin + round * nthreads + lane; + if (fetch < end) { + const int gid = list[fetch]; + s_xy[2 * lane + 0] = mean2d[2 * (size_t)gid + 0]; + s_xy[2 * lane + 1] = mean2d[2 * (size_t)gid + 1]; + for (int k = 0; k < 4; k++) { + s_co[4 * lane + k] = conic[4 * (size_t)gid + k]; + s_color[4 * lane + k] = color[4 * (size_t)gid + k]; + } + } + #pragma omp barrier + + const int remaining = total - round * nthreads; + const int count = (nthreads < remaining) ? nthreads : remaining; + + for (int pixel = lane; pixel < BLOCK_SIZE; pixel += nthreads) { + const int x = tx * BLOCK_X + pixel % BLOCK_X; + const int y = ty * BLOCK_Y + pixel / BLOCK_X; + if (x >= width || y >= height) continue; + + float transmittance = p_trans[pixel]; + if (transmittance < MIN_TRANSMITTANCE) continue; + + const float pixf_x = (float)x + 0.5f; + const float pixf_y = (float)y + 0.5f; + float r = p_r[pixel], g = p_g[pixel], b = p_b[pixel]; + + for (int j = 0; j < count; j++) { + const float ddx = s_xy[2 * j + 0] - pixf_x; + const float ddy = s_xy[2 * j + 1] - pixf_y; + const float power = -0.5f * (s_co[4 * j + 0] * ddx * ddx + + s_co[4 * j + 2] * ddy * ddy) - + s_co[4 * j + 1] * ddx * ddy; + if (power > 0.0f) continue; + + const float a = s_co[4 * j + 3] * expf(power); + const float alpha = (a < 0.99f) ? a : 0.99f; + if (alpha < MIN_OPACITY) continue; + + const float weight = alpha * transmittance; + r += s_color[4 * j + 0] * weight; + g += s_color[4 * j + 1] * weight; + b += s_color[4 * j + 2] * weight; + + transmittance *= 1.0f - alpha; + if (transmittance < MIN_TRANSMITTANCE) break; + } + + p_trans[pixel] = transmittance; + p_r[pixel] = r; + p_g[pixel] = g; + p_b[pixel] = b; + } + // finish shading against this batch before it is overwritten + #pragma omp barrier + } + + for (int pixel = lane; pixel < BLOCK_SIZE; pixel += nthreads) { + const int x = tx * BLOCK_X + pixel % BLOCK_X; + const int y = ty * BLOCK_Y + pixel / BLOCK_X; + if (x >= width || y >= height) continue; + const size_t o = 4 * ((size_t)y * width + x); + image[o + 0] = p_r[pixel]; + image[o + 1] = p_g[pixel]; + image[o + 2] = p_b[pixel]; + image[o + 3] = 1.0f - p_trans[pixel]; + } + // finish writes/reads of the shared state before the next tile reuses it + #pragma omp barrier + } + } + } +} + +// --------------------------------------------------------------------------- + +int main(int argc, char* argv[]) +{ + if (argc != 5) { + printf("Usage: %s \n", + argv[0]); + return 1; + } + + const int n = atoi(argv[1]); + const int width = atoi(argv[2]); + const int height = atoi(argv[3]); + const int repeat = atoi(argv[4]); + + if (n <= 0 || width <= 0 || height <= 0 || repeat <= 0) { + printf("Error: number of gaussians, image width, image height, and repeat " + "must all be positive integers (got n=%d, width=%d, height=%d, " + "repeat=%d)\n", n, width, height, repeat); + return 1; + } + + Camera cam; + setup_camera(width, height, cam); + MpmParams mpm; + setup_mpm(mpm); + const float render_time = 0.5f; + + Scene scene; + std::vector h_image, h_ref_image; + std::vector h_mean, h_velocity, h_affine, h_defgrad, h_grid; + std::vector h_mean2d, h_conic, h_color; + std::vector h_radii; + std::vector tile_offsets, tile_list; + try { + generate_scene(n, scene); + h_mean.resize((size_t)4 * n); + h_velocity.resize((size_t)4 * n); + h_affine.resize((size_t)9 * n); + h_defgrad.resize((size_t)9 * n); + h_grid.resize(4 * (size_t)MPM_CELLS); + h_mean2d.resize((size_t)2 * n); + h_conic.resize((size_t)4 * n); + h_color.resize((size_t)4 * n); + h_radii.resize(n); + h_image.resize(4 * (size_t)width * height); + h_ref_image.resize(4 * (size_t)width * height); + } catch (const std::bad_alloc&) { + printf("Failed to allocate the host buffers for %d gaussians and a %d x %d " + "image\n", n, width, height); + return 1; + } + + const int num_tiles = cam.tiles_x * cam.tiles_y; + printf("Gaussians: %d, image: %d x %d (%d tiles), MPM grid: %d^3\n", + n, width, height, num_tiles, MPM_GRID); + + Scene ref_scene = scene; + + // the device sees the same layouts as the other backends: xyzw vectors for + // the poses, component major for the 3x3 tensors + for (int i = 0; i < n; i++) { + for (int k = 0; k < 4; k++) h_mean[4 * (size_t)i + k] = scene.mean[4 * (size_t)i + k]; + for (int k = 0; k < 3; k++) h_velocity[4 * (size_t)i + k] = scene.velocity[3 * (size_t)i + k]; + h_velocity[4 * (size_t)i + 3] = 0.0f; + for (int k = 0; k < 9; k++) { + h_affine[(size_t)k * n + i] = scene.affine[9 * (size_t)i + k]; + h_defgrad[(size_t)k * n + i] = scene.defgrad[9 * (size_t)i + k]; + } + } + + float* mean = h_mean.data(); + float* velocity = h_velocity.data(); + float* affine = h_affine.data(); + float* defgrad = h_defgrad.data(); + float* grid = h_grid.data(); + const float* scale = scene.scale.data(); + const float* quat_l = scene.quat_l.data(); + const float* quat_r = scene.quat_r.data(); + const float* opacity = scene.opacity.data(); + const float* sh = scene.sh.data(); + float* mean2d = h_mean2d.data(); + float* conic = h_conic.data(); + float* color = h_color.data(); + int* radii = h_radii.data(); + float* image = h_image.data(); + + const int num_chunks = (int)scene.chunk_block.size(); + const int* chunk_start = scene.chunk_start.data(); + const int* chunk_block = scene.chunk_block.data(); + + const size_t sh_size = (size_t)SH_COEFFS * 3 * n; + const size_t image_size = 4 * (size_t)width * height; + + float view[9], cam_pos[3]; + for (int k = 0; k < 9; k++) view[k] = cam.view[k]; + for (int k = 0; k < 3; k++) cam_pos[k] = cam.cam_pos[k]; + + #pragma omp target data \ + map(tofrom: mean[0:4*(size_t)n], velocity[0:4*(size_t)n], \ + affine[0:9*(size_t)n], defgrad[0:9*(size_t)n]) \ + map(alloc: grid[0:4*(size_t)MPM_CELLS]) \ + map(to: scale[0:4*(size_t)n], quat_l[0:4*(size_t)n], quat_r[0:4*(size_t)n], \ + opacity[0:n], sh[0:sh_size], view[0:9], cam_pos[0:3], \ + chunk_start[0:num_chunks+1], chunk_block[0:num_chunks]) \ + map(tofrom: mean2d[0:2*(size_t)n], conic[0:4*(size_t)n], \ + color[0:4*(size_t)n], radii[0:n]) + { + + // --- stage 1: one MPM step, verified against the reference --------------- + mpm_step(n, num_chunks, chunk_start, chunk_block, mean, velocity, affine, + defgrad, grid, mpm); + + { + std::vector ref_grid(4 * (size_t)MPM_CELLS); + reference_p2g(ref_scene, mpm, ref_grid.data()); + reference_grid_update(mpm, ref_grid.data()); + reference_g2p(ref_scene, mpm, ref_grid.data()); + + #pragma omp target update from(mean[0:4*(size_t)n], velocity[0:4*(size_t)n]) + + int mpm_errors = 0; + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(mean[4 * (size_t)i + k], ref_scene.mean[4 * (size_t)i + k], 1e-4f)) + mpm_errors++; + for (int i = 0; i < n && mpm_errors == 0; i++) + for (int k = 0; k < 3; k++) + if (!close_enough(velocity[4 * (size_t)i + k], ref_scene.velocity[3 * (size_t)i + k], 1e-3f)) + mpm_errors++; + printf("MPM step: %s\n", mpm_errors == 0 ? "PASS" : "FAIL"); + } + + auto start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) + mpm_step(n, num_chunks, chunk_start, chunk_block, mean, velocity, affine, + defgrad, grid, mpm); + auto end = std::chrono::steady_clock::now(); + auto time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the MPM step (p2g, grid, g2p): %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 2: the 4D preprocess, from the undeformed scene --------------- + for (int i = 0; i < n; i++) + for (int k = 0; k < 4; k++) + mean[4 * (size_t)i + k] = scene.mean[4 * (size_t)i + k]; + #pragma omp target update to(mean[0:4*(size_t)n]) + + preprocess(n, render_time, mean, scale, quat_l, quat_r, opacity, sh, view, + cam_pos, cam.focal_x, cam.focal_y, cam.center_x, cam.center_y, + cam.near_plane, width, height, mean2d, conic, color, radii); + + std::vector ref_mean2d(2 * (size_t)n), ref_conic(4 * (size_t)n), + ref_color(4 * (size_t)n); + { + std::vector ref_radii(n); + reference_preprocess(scene, cam, render_time, ref_mean2d.data(), + ref_conic.data(), ref_color.data(), ref_radii.data()); + + #pragma omp target update from(mean2d[0:2*(size_t)n], conic[0:4*(size_t)n], \ + color[0:4*(size_t)n], radii[0:n]) + + int pre_errors = 0; + int visible = 0; + for (int i = 0; i < n && pre_errors == 0; i++) { + if (ref_radii[i] > 0) visible++; + if (abs(radii[i] - ref_radii[i]) > 1) pre_errors++; + if (ref_radii[i] == 0 || radii[i] == 0) continue; + for (int k = 0; k < 2; k++) + if (!close_enough(mean2d[2 * (size_t)i + k], ref_mean2d[2 * (size_t)i + k], 1e-3f)) + pre_errors++; + for (int k = 0; k < 4; k++) { + if (!close_enough(conic[4 * (size_t)i + k], ref_conic[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + if (!close_enough(color[4 * (size_t)i + k], ref_color[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + } + } + printf("4D preprocess (%d of %d gaussians visible): %s\n", visible, n, + pre_errors == 0 ? "PASS" : "FAIL"); + + build_tile_lists(cam, n, ref_mean2d.data(), ref_color.data(), + ref_radii.data(), tile_offsets, tile_list); + reference_render(cam, ref_mean2d.data(), ref_conic.data(), ref_color.data(), + tile_offsets.data(), tile_list.data(), h_ref_image.data()); + } + + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) + preprocess(n, render_time, mean, scale, quat_l, quat_r, opacity, sh, view, + cam_pos, cam.focal_x, cam.focal_y, cam.center_x, cam.center_y, + cam.near_plane, width, height, mean2d, conic, color, radii); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the 4D preprocess kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 3: the tile rasterizer --------------------------------------- + for (size_t k = 0; k < 2 * (size_t)n; k++) mean2d[k] = ref_mean2d[k]; + for (size_t k = 0; k < 4 * (size_t)n; k++) conic[k] = ref_conic[k]; + for (size_t k = 0; k < 4 * (size_t)n; k++) color[k] = ref_color[k]; + #pragma omp target update to(mean2d[0:2*(size_t)n], conic[0:4*(size_t)n], \ + color[0:4*(size_t)n]) + + if (tile_list.empty()) tile_list.push_back(0); + const int list_size = (int)tile_list.size(); + const int* offsets = tile_offsets.data(); + const int* list = tile_list.data(); + + printf("Gaussian instances after tiling: %d (%.1f per tile)\n", + list_size, (double)list_size / num_tiles); + + #pragma omp target data map(to: offsets[0:num_tiles+1], list[0:list_size]) \ + map(from: image[0:image_size]) + { + render(num_tiles, cam.tiles_x, width, height, mean2d, conic, color, offsets, + list, image); + + #pragma omp target update from(image[0:image_size]) + + { + int render_errors = 0; + for (size_t k = 0; k < image_size && render_errors == 0; k++) + if (!close_enough(image[k], h_ref_image[k], 1e-3f)) render_errors++; + printf("Rasterizer: %s\n", render_errors == 0 ? "PASS" : "FAIL"); + } + + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) + render(num_tiles, cam.tiles_x, width, height, mean2d, conic, color, offsets, + list, image); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the rasterizer kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + } + + } + + return 0; +} diff --git a/src/gsplat4d-sycl/CMakeLists.txt b/src/gsplat4d-sycl/CMakeLists.txt new file mode 100644 index 000000000..b1669a2fc --- /dev/null +++ b/src/gsplat4d-sycl/CMakeLists.txt @@ -0,0 +1,11 @@ +# gsplat4d-sycl/CMakeLists.txt + +add_hecbench_benchmark( + NAME gsplat4d + MODEL sycl + SOURCES main.cpp + CATEGORIES simulation + TEST_ARGS 200000 1280 720 100 + TEST_REGEX "(?s)MPM step: PASS.*4D preprocess \([0-9]+ of [0-9]+ gaussians visible\): PASS.*Rasterizer: PASS" + TEST_TIMEOUT 300 +) diff --git a/src/gsplat4d-sycl/Makefile b/src/gsplat4d-sycl/Makefile new file mode 100644 index 000000000..ab623c2cc --- /dev/null +++ b/src/gsplat4d-sycl/Makefile @@ -0,0 +1,82 @@ +#=============================================================================== +# User Options +#=============================================================================== + +# Compiler can be set below, or via environment variable +CC = icpx +OPTIMIZE = yes +DEBUG = no +LAUNCHER ?= + +GPU = yes +CUDA = no +CUDA_ARCH ?= sm_70 +HIP = no +HIP_ARCH ?= gfx908 +#GCC_TOOLCHAIN = "/auto/software/gcc/x86_64/gcc-9.1.0/" + +#=============================================================================== +# Program name & source code list +#=============================================================================== + +program = main + +source = main.cpp + +obj = $(source:.cpp=.o) + +#=============================================================================== +# Sets Flags +#=============================================================================== + +# Standard Flags +CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -I../gsplat4d-cuda \ + -fsycl --gcc-toolchain=$(GCC_TOOLCHAIN) + +ifeq ($(VENDOR), AdaptiveCpp) + CFLAGS_TMP := $(CFLAGS) + CFLAGS = $(filter-out -fsycl, $(CFLAGS_TMP)) +endif + +# Linker Flags +LDFLAGS = + +ifeq ($(CUDA), yes) + CFLAGS += -fsycl-targets=nvptx64-nvidia-cuda \ + -Xsycl-target-backend --cuda-gpu-arch=$(CUDA_ARCH) +endif + +ifeq ($(HIP), yes) + CFLAGS += -fsycl-targets=amdgcn-amd-amdhsa \ + -Xsycl-target-backend --offload-arch=$(HIP_ARCH) +endif + +# Debug Flags +ifeq ($(DEBUG),yes) + CFLAGS += -g -DDEBUG + LDFLAGS += -g +endif + +# Optimization Flags +ifeq ($(OPTIMIZE),yes) + CFLAGS += -O3 +endif + +ifeq ($(GPU),yes) + CFLAGS +=-DUSE_GPU +endif +#=============================================================================== +# Targets to Build +#=============================================================================== + +$(program): $(obj) Makefile + $(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS) + +%.o: %.cpp ../gsplat4d-cuda/reference.h Makefile + $(CC) $(CFLAGS) -c $< -o $@ + +clean: + rm -rf $(program) $(obj) + +run: $(program) + $(LAUNCHER) ./$(program) 200000 1280 720 100 diff --git a/src/gsplat4d-sycl/main.cpp b/src/gsplat4d-sycl/main.cpp new file mode 100644 index 000000000..9465fa4be --- /dev/null +++ b/src/gsplat4d-sycl/main.cpp @@ -0,0 +1,951 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include "reference.h" + +#define P2G_BLOCK 128 +#define GRID_BLOCK 256 +#define PREPROCESS_BLOCK 128 + +// 16 byte vectors, so that a work item loads a pose or a splat with one +// 128-bit access +struct alignas(16) Float4 { float x, y, z, w; }; +struct alignas(8) Float2 { float x, y; }; + +using AtomicGlobal = + sycl::atomic_ref; +using AtomicLocal = + sycl::atomic_ref; + +static inline void quad_weights(float fx, float& w0, float& w1, float& w2) +{ + const float a = 1.5f - fx; + const float b = fx - 1.0f; + const float c = fx - 0.5f; + w0 = 0.5f * a * a; + w1 = 0.75f - b * b; + w2 = 0.5f * c * c; +} + +// --------------------------------------------------------------------------- +// stage 1: MLS-MPM +// --------------------------------------------------------------------------- + +// One work group scatters one chunk of particles, all of which belong to the +// same MPM_BLOCK^3 cell block, so their stencil footprint fits a MPM_TILE^3 +// tile in local memory. A particle that has drifted out of its block since +// the binning still lands correctly through the global fallback. +static void mpm_p2g(sycl::nd_item<1>& item, float* tile, int n, const MpmParams p, + const int* __restrict chunk_start, + const int* __restrict chunk_block, + const Float4* __restrict mean, + const Float4* __restrict velocity, + const float* __restrict affine_in, + const float* __restrict defgrad, + float* __restrict grid) +{ + auto group = item.get_group(); + const int lid = item.get_local_id(0); + + for (int k = lid; k < MPM_TILE_CELLS * 4; k += P2G_BLOCK) tile[k] = 0.0f; + + const int chunk = item.get_group(0); + const int begin = chunk_start[chunk]; + const int end = chunk_start[chunk + 1]; + + const int gb = chunk_block[chunk]; + const int obz = (gb % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int oby = ((gb / MPM_BLOCKS_PER_DIM) % MPM_BLOCKS_PER_DIM) * MPM_BLOCK - 1; + const int obx = (gb / (MPM_BLOCKS_PER_DIM * MPM_BLOCKS_PER_DIM)) * MPM_BLOCK - 1; + + sycl::group_barrier(group); + + for (int i = begin + lid; i < end; i += P2G_BLOCK) { + const Float4 x = mean[i]; + const Float4 v = velocity[i]; + + float C[9], F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) C[k] = affine_in[(size_t)k * n + i]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + float F[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = sycl::fma(p.dt * C[a * 3 + k], F0[k * 3 + b], acc); + F[a * 3 + b] = acc; + } + + const float det = + F[0] * (F[4] * F[8] - F[5] * F[7]) - + F[1] * (F[3] * F[8] - F[5] * F[6]) + + F[2] * (F[3] * F[7] - F[4] * F[6]); + const float safe_J = + (sycl::fabs(det) < 1e-6f) ? ((det < 0.0f) ? -1e-6f : 1e-6f) : det; + const float inv_J = 1.0f / safe_J; + + const float cof[9] = { + (F[4] * F[8] - F[5] * F[7]), -(F[3] * F[8] - F[5] * F[6]), (F[3] * F[7] - F[4] * F[6]), + -(F[1] * F[8] - F[2] * F[7]), (F[0] * F[8] - F[2] * F[6]), -(F[0] * F[7] - F[1] * F[6]), + (F[1] * F[5] - F[2] * F[4]), -(F[0] * F[5] - F[2] * F[3]), (F[0] * F[4] - F[1] * F[3]) }; + + const float coeff = + (p.lambda * sycl::native::log(sycl::fabs(safe_J)) - p.mu) * inv_J; + float P[9]; + #pragma unroll + for (int k = 0; k < 9; k++) P[k] = sycl::fma(coeff, cof[k], p.mu * F[k]); + + const float s = -p.dt * p.particle_volume * 4.0f * p.inv_dx * p.inv_dx; + float affine[9]; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = sycl::fma(P[a * 3 + k], F[b * 3 + k], acc); + affine[a * 3 + b] = sycl::fma(s, acc, p.particle_mass * C[a * 3 + b]); + } + + // F is only used for the stress above; it is deliberately not persisted + // here. The single per-step deformation-gradient update is applied once, in + // mpm_g2p, from the original F0 (this matches the host reference; writing F + // back here would advance the gradient twice per step). + + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)sycl::floor(gx - 0.5f); + const int by = (int)sycl::floor(gy - 0.5f); + const int bz = (int)sycl::floor(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + const float mv[3] = { p.particle_mass * v.x, p.particle_mass * v.y, + p.particle_mass * v.z }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float dpz = ((float)c - fz) * p.dx; + const float w = wxy * wz[c]; + + float val[4]; + #pragma unroll + for (int k = 0; k < 3; k++) { + float impulse = affine[k * 3 + 0] * dpx; + impulse = sycl::fma(affine[k * 3 + 1], dpy, impulse); + impulse = sycl::fma(affine[k * 3 + 2], dpz, impulse); + val[k] = w * (mv[k] + impulse); + } + val[3] = w * p.particle_mass; + + const int lx = ix - obx, ly = iy - oby, lz = iz - obz; + if ((unsigned)lx < MPM_TILE && (unsigned)ly < MPM_TILE && + (unsigned)lz < MPM_TILE) { + float* cell = tile + 4 * ((lx * MPM_TILE + ly) * MPM_TILE + lz); + #pragma unroll + for (int k = 0; k < 4; k++) AtomicLocal(cell[k]).fetch_add(val[k]); + } else { + float* cell = + grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + #pragma unroll + for (int k = 0; k < 4; k++) AtomicGlobal(cell[k]).fetch_add(val[k]); + } + } + } + } + } + + sycl::group_barrier(group); + + for (int c = lid; c < MPM_TILE_CELLS; c += P2G_BLOCK) { + const int lz = c % MPM_TILE; + const int ly = (c / MPM_TILE) % MPM_TILE; + const int lx = c / (MPM_TILE * MPM_TILE); + const int ix = obx + lx, iy = oby + ly, iz = obz + lz; + if ((unsigned)ix >= MPM_GRID || (unsigned)iy >= MPM_GRID || + (unsigned)iz >= MPM_GRID) + continue; + + const float* acc = tile + 4 * c; + if (acc[0] == 0.0f && acc[1] == 0.0f && acc[2] == 0.0f && acc[3] == 0.0f) + continue; + + float* dst = grid + 4 * ((size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz); + #pragma unroll + for (int k = 0; k < 4; k++) AtomicGlobal(dst[k]).fetch_add(acc[k]); + } +} + +static void mpm_grid_update(sycl::nd_item<1>& item, const MpmParams p, + Float4* __restrict grid) +{ + const int cell = item.get_global_id(0); + if (cell >= MPM_CELLS) return; + + Float4 g = grid[cell]; + if (g.w <= 0.0f) { + grid[cell] = Float4{ 0.0f, 0.0f, 0.0f, g.w }; + return; + } + + const float inv_mass = 1.0f / g.w; + float vx = g.x * inv_mass; + float vy = sycl::fma(p.dt, p.gravity, g.y * inv_mass); + float vz = g.z * inv_mass; + + const int iz = cell % MPM_GRID; + const int iy = (cell / MPM_GRID) % MPM_GRID; + const int ix = cell / (MPM_GRID * MPM_GRID); + + if (ix < p.boundary && vx < 0.0f) vx = 0.0f; + if (ix >= MPM_GRID - p.boundary && vx > 0.0f) vx = 0.0f; + if (iy < p.boundary && vy < 0.0f) vy = 0.0f; + if (iy >= MPM_GRID - p.boundary && vy > 0.0f) vy = 0.0f; + if (iz < p.boundary && vz < 0.0f) vz = 0.0f; + if (iz >= MPM_GRID - p.boundary && vz > 0.0f) vz = 0.0f; + + grid[cell] = Float4{ vx, vy, vz, g.w }; +} + +static void mpm_g2p(sycl::nd_item<1>& item, int n, const MpmParams p, + Float4* __restrict mean, + Float4* __restrict velocity, + float* __restrict affine_out, + float* __restrict defgrad, + const Float4* __restrict grid) +{ + const int i = item.get_global_id(0); + if (i >= n) return; + + const Float4 x = mean[i]; + const float gx = x.x * p.inv_dx, gy = x.y * p.inv_dx, gz = x.z * p.inv_dx; + const int bx = (int)sycl::floor(gx - 0.5f); + const int by = (int)sycl::floor(gy - 0.5f); + const int bz = (int)sycl::floor(gz - 0.5f); + const float fx = gx - bx, fy = gy - by, fz = gz - bz; + + float wx[3], wy[3], wz[3]; + quad_weights(fx, wx[0], wx[1], wx[2]); + quad_weights(fy, wy[0], wy[1], wy[2]); + quad_weights(fz, wz[0], wz[1], wz[2]); + + float nv[3] = { 0.0f, 0.0f, 0.0f }; + float nC[9] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f }; + + #pragma unroll + for (int a = 0; a < 3; a++) { + const int ix = bx + a; + if (ix < 0 || ix >= MPM_GRID) continue; + const float dpx = ((float)a - fx) * p.dx; + #pragma unroll + for (int b = 0; b < 3; b++) { + const int iy = by + b; + if (iy < 0 || iy >= MPM_GRID) continue; + const float dpy = ((float)b - fy) * p.dx; + const float wxy = wx[a] * wy[b]; + #pragma unroll + for (int c = 0; c < 3; c++) { + const int iz = bz + c; + if (iz < 0 || iz >= MPM_GRID) continue; + const float w = wxy * wz[c]; + const Float4 g = grid[(size_t)(ix * MPM_GRID + iy) * MPM_GRID + iz]; + const float gv[3] = { g.x, g.y, g.z }; + const float dpos[3] = { dpx, dpy, ((float)c - fz) * p.dx }; + + #pragma unroll + for (int k = 0; k < 3; k++) { + nv[k] = sycl::fma(w, gv[k], nv[k]); + const float wg = 4.0f * p.inv_dx * p.inv_dx * w * gv[k]; + #pragma unroll + for (int l = 0; l < 3; l++) + nC[k * 3 + l] = sycl::fma(wg, dpos[l], nC[k * 3 + l]); + } + } + } + } + + float F0[9]; + #pragma unroll + for (int k = 0; k < 9; k++) F0[k] = defgrad[(size_t)k * n + i]; + + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = 0; b < 3; b++) { + float acc = F0[a * 3 + b]; + #pragma unroll + for (int k = 0; k < 3; k++) + acc = sycl::fma(p.dt * nC[a * 3 + k], F0[k * 3 + b], acc); + defgrad[(size_t)(a * 3 + b) * n + i] = acc; + } + + #pragma unroll + for (int k = 0; k < 9; k++) affine_out[(size_t)k * n + i] = nC[k]; + + velocity[i] = Float4{ nv[0], nv[1], nv[2], 0.0f }; + mean[i] = Float4{ sycl::fma(p.dt, nv[0], x.x), sycl::fma(p.dt, nv[1], x.y), + sycl::fma(p.dt, nv[2], x.z), x.w }; +} + +// --------------------------------------------------------------------------- +// stage 2: condition the 4D Gaussian on t, project, shade +// --------------------------------------------------------------------------- + +static void preprocess(sycl::nd_item<1>& item, int n, float time, const Camera cam, + const Float4* __restrict mean4, + const Float4* __restrict scale4, + const Float4* __restrict quat_l, + const Float4* __restrict quat_r, + const float* __restrict opacity_in, + const float* __restrict sh, + Float2* __restrict mean2d, + Float4* __restrict conic_opacity, + Float4* __restrict color_depth, + int* __restrict radii) +{ + const int i = item.get_global_id(0); + if (i >= n) return; + + radii[i] = 0; + + const Float4 mu = mean4[i]; + const Float4 sc = scale4[i]; + const Float4 ql = quat_l[i]; + const Float4 qr = quat_r[i]; + + const float lw = ql.x, lx = ql.y, ly = ql.z, lz = ql.w; + const float rw = qr.x, rx = qr.y, ry = qr.z, rz = qr.w; + + const float L[16] = { + lw, -lx, -ly, -lz, + lx, lw, -lz, ly, + ly, lz, lw, -lx, + lz, -ly, lx, lw }; + const float R[16] = { + rw, -rx, -ry, -rz, + rx, rw, rz, -ry, + ry, -rz, rw, rx, + rz, ry, -rx, rw }; + + float M[16]; + #pragma unroll + for (int a = 0; a < 4; a++) + #pragma unroll + for (int b = 0; b < 4; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) acc = sycl::fma(L[a * 4 + k], R[k * 4 + b], acc); + M[a * 4 + b] = acc; + } + + // the quaternion basis is ordered (t, x, y, z) + const float s2[4] = { sc.w * sc.w, sc.x * sc.x, sc.y * sc.y, sc.z * sc.z }; + + float st[3], sxyz[6]; + #pragma unroll + for (int a = 0; a < 3; a++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) + acc = sycl::fma(M[(a + 1) * 4 + k] * s2[k], M[k], acc); + st[a] = acc; + } + float sigma_tt = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) sigma_tt = sycl::fma(M[k] * s2[k], M[k], sigma_tt); + + int idx = 0; + #pragma unroll + for (int a = 0; a < 3; a++) + #pragma unroll + for (int b = a; b < 3; b++) { + float acc = 0.0f; + #pragma unroll + for (int k = 0; k < 4; k++) + acc = sycl::fma(M[(a + 1) * 4 + k] * s2[k], M[(b + 1) * 4 + k], acc); + sxyz[idx++] = acc; + } + + const float inv_tt = 1.0f / sigma_tt; + const float dt = time - mu.w; + + const float alpha = opacity_in[i] * sycl::native::exp(-0.5f * dt * dt * inv_tt); + if (alpha < MIN_OPACITY) return; + + const float mean3[3] = { sycl::fma(dt * inv_tt, st[0], mu.x), + sycl::fma(dt * inv_tt, st[1], mu.y), + sycl::fma(dt * inv_tt, st[2], mu.z) }; + const float cov3[6] = { sxyz[0] - st[0] * st[0] * inv_tt, + sxyz[1] - st[0] * st[1] * inv_tt, + sxyz[2] - st[0] * st[2] * inv_tt, + sxyz[3] - st[1] * st[1] * inv_tt, + sxyz[4] - st[1] * st[2] * inv_tt, + sxyz[5] - st[2] * st[2] * inv_tt }; + + const float px = mean3[0] - cam.cam_pos[0]; + const float py = mean3[1] - cam.cam_pos[1]; + const float pz = mean3[2] - cam.cam_pos[2]; + float t[3]; + #pragma unroll + for (int a = 0; a < 3; a++) + t[a] = sycl::fma(cam.view[a * 3 + 0], px, + sycl::fma(cam.view[a * 3 + 1], py, cam.view[a * 3 + 2] * pz)); + + if (t[2] < cam.near_plane) return; + + const float inv_z = 1.0f / t[2]; + const float inv_z2 = inv_z * inv_z; + + const float lim_x = 1.3f * (0.5f * cam.width) / cam.focal_x; + const float lim_y = 1.3f * (0.5f * cam.height) / cam.focal_y; + const float tx = t[2] * sycl::fmin(lim_x, sycl::fmax(-lim_x, t[0] * inv_z)); + const float ty = t[2] * sycl::fmin(lim_y, sycl::fmax(-lim_y, t[1] * inv_z)); + + const float J[6] = { cam.focal_x * inv_z, 0.0f, -cam.focal_x * tx * inv_z2, + 0.0f, cam.focal_y * inv_z, -cam.focal_y * ty * inv_z2 }; + + float T[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + T[a * 3 + b] = sycl::fma(J[a * 3 + 0], cam.view[0 * 3 + b], + sycl::fma(J[a * 3 + 1], cam.view[1 * 3 + b], + J[a * 3 + 2] * cam.view[2 * 3 + b])); + + const float S[9] = { cov3[0], cov3[1], cov3[2], + cov3[1], cov3[3], cov3[4], + cov3[2], cov3[4], cov3[5] }; + float TS[6]; + #pragma unroll + for (int a = 0; a < 2; a++) + #pragma unroll + for (int b = 0; b < 3; b++) + TS[a * 3 + b] = sycl::fma(T[a * 3 + 0], S[0 * 3 + b], + sycl::fma(T[a * 3 + 1], S[1 * 3 + b], + T[a * 3 + 2] * S[2 * 3 + b])); + + float ca = 0.3f, cb = 0.0f, cc = 0.3f; + #pragma unroll + for (int k = 0; k < 3; k++) { + ca = sycl::fma(TS[0 * 3 + k], T[0 * 3 + k], ca); + cb = sycl::fma(TS[0 * 3 + k], T[1 * 3 + k], cb); + cc = sycl::fma(TS[1 * 3 + k], T[1 * 3 + k], cc); + } + + const float det = sycl::fma(ca, cc, -cb * cb); + if (det <= 0.0f) return; + const float inv_det = 1.0f / det; + + const float mid = 0.5f * (ca + cc); + const float disc = sycl::sqrt(sycl::fmax(0.1f, sycl::fma(mid, mid, -det))); + const float radius = sycl::ceil(3.0f * sycl::sqrt(sycl::fmax(mid + disc, mid - disc))); + + const float mx = sycl::fma(cam.focal_x * inv_z, t[0], cam.center_x); + const float my = sycl::fma(cam.focal_y * inv_z, t[1], cam.center_y); + + if (mx + radius < 0.0f || mx - radius > (float)cam.width || + my + radius < 0.0f || my - radius > (float)cam.height) + return; + + const float inv_len = sycl::rsqrt(sycl::fma(px, px, sycl::fma(py, py, pz * pz))); + const float dx = px * inv_len, dy = py * inv_len, dz = pz * inv_len; + + float basis[SH_COEFFS]; + basis[0] = SH_C0; + basis[1] = -SH_C1 * dy; + basis[2] = SH_C1 * dz; + basis[3] = -SH_C1 * dx; + + const float xx = dx * dx, yy = dy * dy, zz = dz * dz; + const float xy = dx * dy, yz = dy * dz, xz = dx * dz; + basis[4] = SH_C2_0 * xy; + basis[5] = SH_C2_1 * yz; + basis[6] = SH_C2_2 * (2.0f * zz - xx - yy); + basis[7] = SH_C2_3 * xz; + basis[8] = SH_C2_4 * (xx - yy); + basis[9] = SH_C3_0 * dy * (3.0f * xx - yy); + basis[10] = SH_C3_1 * xy * dz; + basis[11] = SH_C3_2 * dy * (4.0f * zz - xx - yy); + basis[12] = SH_C3_3 * dz * (2.0f * zz - 3.0f * xx - 3.0f * yy); + basis[13] = SH_C3_4 * dx * (4.0f * zz - xx - yy); + basis[14] = SH_C3_5 * dz * (xx - yy); + basis[15] = SH_C3_6 * dx * (xx - 3.0f * yy); + + float rgb[3] = { 0.0f, 0.0f, 0.0f }; + #pragma unroll + for (int c = 0; c < SH_COEFFS; c++) { + const float w = basis[c]; + #pragma unroll + for (int ch = 0; ch < 3; ch++) + rgb[ch] = sycl::fma(w, sh[((size_t)c * 3 + ch) * n + i], rgb[ch]); + } + + radii[i] = (int)radius; + mean2d[i] = Float2{ mx, my }; + conic_opacity[i] = Float4{ cc * inv_det, -cb * inv_det, ca * inv_det, alpha }; + color_depth[i] = Float4{ sycl::fmax(rgb[0] + 0.5f, 0.0f), + sycl::fmax(rgb[1] + 0.5f, 0.0f), + sycl::fmax(rgb[2] + 0.5f, 0.0f), t[2] }; +} + +// --------------------------------------------------------------------------- +// stage 3: tile rasterizer +// --------------------------------------------------------------------------- + +static void render(sycl::nd_item<2>& item, Float2* s_xy, Float4* s_co, + Float4* s_color, const Camera cam, + const Float2* __restrict mean2d, + const Float4* __restrict conic_opacity, + const Float4* __restrict color_depth, + const int* __restrict tile_offsets, + const int* __restrict tile_list, + Float4* __restrict image) +{ + const int tile = item.get_group(0) * cam.tiles_x + item.get_group(1); + const int x = item.get_global_id(1); + const int y = item.get_global_id(0); + const int lane = item.get_local_id(0) * BLOCK_X + item.get_local_id(1); + + const bool valid = (x < cam.width) && (y < cam.height); + bool done = !valid; + + const float pixf_x = (float)x + 0.5f; + const float pixf_y = (float)y + 0.5f; + + const int begin = tile_offsets[tile]; + const int end = tile_offsets[tile + 1]; + const int rounds = (end - begin + BLOCK_SIZE - 1) / BLOCK_SIZE; + + float transmittance = 1.0f; + float r = 0.0f, g = 0.0f, b = 0.0f; + + int todo = end - begin; + auto group = item.get_group(); + + for (int round = 0; round < rounds; round++, todo -= BLOCK_SIZE) { + // every pixel of the tile is saturated, so no further Gaussian can + // contribute to this work group + if (sycl::all_of_group(group, done)) break; + + const int fetch = begin + round * BLOCK_SIZE + lane; + if (fetch < end) { + const int gid = tile_list[fetch]; + s_xy[lane] = mean2d[gid]; + s_co[lane] = conic_opacity[gid]; + s_color[lane] = color_depth[gid]; + } + sycl::group_barrier(group); + + const int count = sycl::min(BLOCK_SIZE, todo); + for (int j = 0; j < count && !done; j++) { + const Float2 xy = s_xy[j]; + const Float4 co = s_co[j]; + const float dx = xy.x - pixf_x; + const float dy = xy.y - pixf_y; + + const float power = + -0.5f * sycl::fma(co.x, dx * dx, co.z * dy * dy) - co.y * dx * dy; + if (power > 0.0f) continue; + + const float alpha = sycl::fmin(0.99f, co.w * sycl::native::exp(power)); + if (alpha < MIN_OPACITY) continue; + + const float weight = alpha * transmittance; + const Float4 c = s_color[j]; + r = sycl::fma(c.x, weight, r); + g = sycl::fma(c.y, weight, g); + b = sycl::fma(c.z, weight, b); + + transmittance *= 1.0f - alpha; + if (transmittance < MIN_TRANSMITTANCE) done = true; + } + sycl::group_barrier(group); + } + + if (valid) + image[(size_t)y * cam.width + x] = Float4{ r, g, b, 1.0f - transmittance }; +} + +// --------------------------------------------------------------------------- + +static int run(int argc, char* argv[]) +{ + if (argc != 5) { + printf("Usage: %s \n", + argv[0]); + return 1; + } + + const int n = atoi(argv[1]); + const int width = atoi(argv[2]); + const int height = atoi(argv[3]); + const int repeat = atoi(argv[4]); + + if (n <= 0 || width <= 0 || height <= 0 || repeat <= 0) { + printf("Error: number of gaussians, image width, image height, and repeat " + "must all be positive integers (got n=%d, width=%d, height=%d, " + "repeat=%d)\n", n, width, height, repeat); + return 1; + } + + Camera cam; + setup_camera(width, height, cam); + MpmParams mpm; + setup_mpm(mpm); + const float time = 0.5f; + + Scene scene; + std::vector h_image, h_ref_image; + std::vector h_mean2d, h_conic, h_color; + std::vector h_radii; + std::vector tile_offsets, tile_list; + try { + generate_scene(n, scene); + h_mean2d.resize((size_t)2 * n); + h_conic.resize((size_t)4 * n); + h_color.resize((size_t)4 * n); + h_radii.resize(n); + h_image.resize(4 * (size_t)width * height); + h_ref_image.resize(4 * (size_t)width * height); + } catch (const std::bad_alloc&) { + printf("Failed to allocate the host buffers for %d gaussians and a %d x %d " + "image\n", n, width, height); + return 1; + } + + const int num_tiles = cam.tiles_x * cam.tiles_y; + printf("Gaussians: %d, image: %d x %d (%d tiles), MPM grid: %d^3\n", + n, width, height, num_tiles, MPM_GRID); + + Scene ref_scene = scene; + + // held by value in an optional: a default constructed queue would run the + // default selector, which throws outside the handler when no device exists + std::optional queue; + try { +#ifdef USE_GPU + queue.emplace(sycl::gpu_selector_v, sycl::property::queue::in_order()); +#else + queue.emplace(sycl::cpu_selector_v, sycl::property::queue::in_order()); +#endif + } catch (const sycl::exception& e) { + printf("Failed to select a SYCL device: %s\n", e.what()); + return 1; + } + sycl::queue& q = *queue; + + const int num_chunks = (int)scene.chunk_block.size(); + + Float4* d_mean = sycl::malloc_device(n, q); + Float4* d_scale = sycl::malloc_device(n, q); + Float4* d_quat_l = sycl::malloc_device(n, q); + Float4* d_quat_r = sycl::malloc_device(n, q); + Float4* d_velocity = sycl::malloc_device(n, q); + float* d_opacity = sycl::malloc_device(n, q); + float* d_sh = sycl::malloc_device((size_t)SH_COEFFS * 3 * n, q); + float* d_affine = sycl::malloc_device((size_t)9 * n, q); + float* d_defgrad = sycl::malloc_device((size_t)9 * n, q); + Float4* d_grid = sycl::malloc_device(MPM_CELLS, q); + int* d_chunk_start = sycl::malloc_device(num_chunks + 1, q); + int* d_chunk_block = sycl::malloc_device(num_chunks, q); + Float2* d_mean2d = sycl::malloc_device(n, q); + Float4* d_conic = sycl::malloc_device(n, q); + Float4* d_color = sycl::malloc_device(n, q); + int* d_radii = sycl::malloc_device(n, q); + Float4* d_image = sycl::malloc_device((size_t)width * height, q); + int* d_tile_offsets = sycl::malloc_device(num_tiles + 1, q); + + if (d_mean == nullptr || d_scale == nullptr || d_quat_l == nullptr || + d_quat_r == nullptr || d_velocity == nullptr || d_opacity == nullptr || + d_sh == nullptr || d_affine == nullptr || d_defgrad == nullptr || + d_grid == nullptr || d_chunk_start == nullptr || d_chunk_block == nullptr || + d_mean2d == nullptr || d_conic == nullptr || d_color == nullptr || + d_radii == nullptr || d_image == nullptr || d_tile_offsets == nullptr) { + printf("Failed to allocate the device buffers for %d gaussians and a " + "%d x %d image\n", n, width, height); + return 1; + } + + std::vector pack(n); + auto upload_vec4 = [&](const std::vector& src, int comps, Float4* dst) { + for (int i = 0; i < n; i++) + pack[i] = Float4{ src[(size_t)comps * i + 0], src[(size_t)comps * i + 1], + src[(size_t)comps * i + 2], + (comps == 4) ? src[(size_t)comps * i + 3] : 0.0f }; + q.memcpy(dst, pack.data(), sizeof(Float4) * n).wait(); + }; + + std::vector pack9((size_t)9 * n); + auto upload_tensor = [&](const std::vector& src, float* dst) { + for (int k = 0; k < 9; k++) + for (int i = 0; i < n; i++) pack9[(size_t)k * n + i] = src[9 * (size_t)i + k]; + q.memcpy(dst, pack9.data(), sizeof(float) * 9 * (size_t)n).wait(); + }; + + upload_vec4(scene.mean, 4, d_mean); + upload_vec4(scene.scale, 4, d_scale); + upload_vec4(scene.quat_l, 4, d_quat_l); + upload_vec4(scene.quat_r, 4, d_quat_r); + upload_vec4(scene.velocity, 3, d_velocity); + upload_tensor(scene.affine, d_affine); + upload_tensor(scene.defgrad, d_defgrad); + q.memcpy(d_opacity, scene.opacity.data(), sizeof(float) * n); + q.memcpy(d_sh, scene.sh.data(), sizeof(float) * SH_COEFFS * 3 * (size_t)n); + q.memcpy(d_chunk_start, scene.chunk_start.data(), sizeof(int) * (num_chunks + 1)); + q.memcpy(d_chunk_block, scene.chunk_block.data(), sizeof(int) * num_chunks); + //q.wait_and_throw(); + + const int cell_groups = (MPM_CELLS + GRID_BLOCK - 1) / GRID_BLOCK; + const int particle_groups = (n + P2G_BLOCK - 1) / P2G_BLOCK; + + auto mpm_step = [&]() { + q.memset(d_grid, 0, sizeof(Float4) * MPM_CELLS); + q.submit([&](sycl::handler& cgh) { + sycl::local_accessor tile(sycl::range<1>(MPM_TILE_CELLS * 4), cgh); + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t)num_chunks * P2G_BLOCK), + sycl::range<1>(P2G_BLOCK)), + [=](sycl::nd_item<1> item) { + mpm_p2g(item, tile.get_multi_ptr().get(), + n, mpm, d_chunk_start, d_chunk_block, d_mean, d_velocity, + d_affine, d_defgrad, (float*)d_grid); + }); + }); + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t)cell_groups * GRID_BLOCK), + sycl::range<1>(GRID_BLOCK)), + [=](sycl::nd_item<1> item) { mpm_grid_update(item, mpm, d_grid); }); + }); + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t)particle_groups * P2G_BLOCK), + sycl::range<1>(P2G_BLOCK)), + [=](sycl::nd_item<1> item) { + mpm_g2p(item, n, mpm, d_mean, d_velocity, d_affine, d_defgrad, d_grid); + }); + }); + }; + + // --- stage 1 ------------------------------------------------------------ + mpm_step(); + //q.wait_and_throw(); + + { + std::vector ref_grid(4 * (size_t)MPM_CELLS); + reference_p2g(ref_scene, mpm, ref_grid.data()); + reference_grid_update(mpm, ref_grid.data()); + reference_g2p(ref_scene, mpm, ref_grid.data()); + + std::vector got(n); + q.memcpy(got.data(), d_mean, sizeof(Float4) * n).wait(); + int mpm_errors = 0; + for (int i = 0; i < n && mpm_errors == 0; i++) { + const float* r = &ref_scene.mean[4 * (size_t)i]; + if (!close_enough(got[i].x, r[0], 1e-4f) || + !close_enough(got[i].y, r[1], 1e-4f) || + !close_enough(got[i].z, r[2], 1e-4f)) + mpm_errors++; + } + q.memcpy(got.data(), d_velocity, sizeof(Float4) * n).wait(); + for (int i = 0; i < n && mpm_errors == 0; i++) { + const float* r = &ref_scene.velocity[3 * (size_t)i]; + if (!close_enough(got[i].x, r[0], 1e-3f) || + !close_enough(got[i].y, r[1], 1e-3f) || + !close_enough(got[i].z, r[2], 1e-3f)) + mpm_errors++; + } + printf("MPM step: %s\n", mpm_errors == 0 ? "PASS" : "FAIL"); + } + + q.wait_and_throw(); + auto start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) mpm_step(); + q.wait_and_throw(); + auto end = std::chrono::steady_clock::now(); + auto time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the MPM step (p2g, grid, g2p): %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 2 ------------------------------------------------------------ + upload_vec4(scene.mean, 4, d_mean); + + auto preprocess_launch = [&]() { + q.submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<1>( + sycl::range<1>((size_t)((n + PREPROCESS_BLOCK - 1) / PREPROCESS_BLOCK) * + PREPROCESS_BLOCK), + sycl::range<1>(PREPROCESS_BLOCK)), + [=](sycl::nd_item<1> item) { + preprocess(item, n, time, cam, d_mean, d_scale, d_quat_l, d_quat_r, + d_opacity, d_sh, d_mean2d, d_conic, d_color, d_radii); + }); + }); + }; + + preprocess_launch(); + //q.wait_and_throw(); + + std::vector ref_mean2d(2 * (size_t)n), ref_conic(4 * (size_t)n), + ref_color(4 * (size_t)n); + { + std::vector ref_radii(n); + reference_preprocess(scene, cam, time, ref_mean2d.data(), ref_conic.data(), + ref_color.data(), ref_radii.data()); + + q.memcpy(h_mean2d.data(), d_mean2d, sizeof(Float2) * n); + q.memcpy(h_conic.data(), d_conic, sizeof(Float4) * n); + q.memcpy(h_color.data(), d_color, sizeof(Float4) * n); + q.memcpy(h_radii.data(), d_radii, sizeof(int) * n); + //q.wait_and_throw(); + + int pre_errors = 0; + int visible = 0; + for (int i = 0; i < n && pre_errors == 0; i++) { + if (ref_radii[i] > 0) visible++; + if (abs(h_radii[i] - ref_radii[i]) > 1) pre_errors++; + if (ref_radii[i] == 0 || h_radii[i] == 0) continue; + for (int k = 0; k < 2; k++) + if (!close_enough(h_mean2d[2 * (size_t)i + k], ref_mean2d[2 * (size_t)i + k], 1e-3f)) + pre_errors++; + for (int k = 0; k < 4; k++) { + if (!close_enough(h_conic[4 * (size_t)i + k], ref_conic[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + if (!close_enough(h_color[4 * (size_t)i + k], ref_color[4 * (size_t)i + k], 1e-3f)) + pre_errors++; + } + } + printf("4D preprocess (%d of %d gaussians visible): %s\n", visible, n, + pre_errors == 0 ? "PASS" : "FAIL"); + + build_tile_lists(cam, n, ref_mean2d.data(), ref_color.data(), + ref_radii.data(), tile_offsets, tile_list); + reference_render(cam, ref_mean2d.data(), ref_conic.data(), ref_color.data(), + tile_offsets.data(), tile_list.data(), h_ref_image.data()); + } + + q.wait_and_throw(); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) preprocess_launch(); + q.wait_and_throw(); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the 4D preprocess kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + // --- stage 3 ------------------------------------------------------------ + q.memcpy(d_mean2d, ref_mean2d.data(), sizeof(Float2) * n); + q.memcpy(d_conic, ref_conic.data(), sizeof(Float4) * n); + q.memcpy(d_color, ref_color.data(), sizeof(Float4) * n); + + const size_t list_size = tile_list.empty() ? 1 : tile_list.size(); + int* d_tile_list = sycl::malloc_device(list_size, q); + if (d_tile_list == nullptr) { + printf("Failed to allocate the device buffer for %zu gaussian instances\n", + list_size); + return 1; + } + q.memcpy(d_tile_offsets, tile_offsets.data(), sizeof(int) * (num_tiles + 1)); + if (!tile_list.empty()) + q.memcpy(d_tile_list, tile_list.data(), sizeof(int) * tile_list.size()); + //q.wait_and_throw(); + + printf("Gaussian instances after tiling: %zu (%.1f per tile)\n", + tile_list.size(), (double)tile_list.size() / num_tiles); + + const sycl::range<2> global_range((size_t)cam.tiles_y * BLOCK_Y, + (size_t)cam.tiles_x * BLOCK_X); + const sycl::range<2> local_range(BLOCK_Y, BLOCK_X); + + auto render_launch = [&]() { + q.submit([&](sycl::handler& cgh) { + sycl::local_accessor s_xy(sycl::range<1>(BLOCK_SIZE), cgh); + sycl::local_accessor s_co(sycl::range<1>(BLOCK_SIZE), cgh); + sycl::local_accessor s_color(sycl::range<1>(BLOCK_SIZE), cgh); + cgh.parallel_for( + sycl::nd_range<2>(global_range, local_range), + [=](sycl::nd_item<2> item) { + render(item, s_xy.get_multi_ptr().get(), + s_co.get_multi_ptr().get(), + s_color.get_multi_ptr().get(), + cam, d_mean2d, d_conic, d_color, d_tile_offsets, d_tile_list, + d_image); + }); + }); + }; + + q.memset(d_image, 0, sizeof(Float4) * (size_t)width * height).wait(); + render_launch(); + q.memcpy(h_image.data(), d_image, sizeof(Float4) * (size_t)width * height).wait(); + + { + int render_errors = 0; + for (size_t k = 0; k < h_image.size() && render_errors == 0; k++) + if (!close_enough(h_image[k], h_ref_image[k], 1e-3f)) render_errors++; + printf("Rasterizer: %s\n", render_errors == 0 ? "PASS" : "FAIL"); + } + + q.wait_and_throw(); + start = std::chrono::steady_clock::now(); + for (int r = 0; r < repeat; r++) render_launch(); + q.wait_and_throw(); + end = std::chrono::steady_clock::now(); + time_ns = std::chrono::duration_cast(end - start).count(); + printf("Average execution time of the rasterizer kernel: %f (us)\n", + time_ns * 1e-3 / repeat); + + sycl::free(d_mean, q); sycl::free(d_scale, q); sycl::free(d_quat_l, q); + sycl::free(d_quat_r, q); sycl::free(d_velocity, q); sycl::free(d_opacity, q); + sycl::free(d_sh, q); sycl::free(d_affine, q); sycl::free(d_defgrad, q); + sycl::free(d_grid, q); sycl::free(d_chunk_start, q); sycl::free(d_chunk_block, q); + sycl::free(d_mean2d, q); sycl::free(d_conic, q); sycl::free(d_color, q); + sycl::free(d_radii, q); sycl::free(d_image, q); sycl::free(d_tile_offsets, q); + sycl::free(d_tile_list, q); + + return 0; +} + +int main(int argc, char* argv[]) +{ + // allocation and submission can throw too, so the whole run is guarded + try { + return run(argc, argv); + } catch (const sycl::exception& e) { + printf("SYCL error: %s\n", e.what()); + return 1; + } +}