← Back to Step 1 in blog post
inria → s00s00 → s01s01 → s02s02 → s03s03 → s04s04 → s05s05 → s06

s00_baseline → s01_literature

Step 1: Consolidating What the Literature Already Knows  ·  +1420 / −1548 lines across 7 changed files

Core files: csrc/auxiliary.h +737/−124csrc/forward.h +18/−27csrc/forward.cu.cc +188/−380csrc/backward.h +25/−26csrc/backward.cu.cc +362/−624csrc/rasterizer_impl.h +43/−51csrc/rasterizer_impl.cu.cc +47/−316
csrc/auxiliary.h +737 −124
s00_baseline/csrc/auxiliary.h
s01_literature/csrc/auxiliary.h
↕ Expand 9 unchanged lines (lines 1–9)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#ifndef CUDA_RASTERIZER_AUXILIARY_H_INCLUDED12#ifndef CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
13#define CUDA_RASTERIZER_AUXILIARY_H_INCLUDED13#define CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
14
15#include <cuda_runtime.h>
16
17#include <cmath>
18#include <cstdint>
1419
15#include "config.h"20#include "config.h"
16#include "stdio.h"21#include "stdio.h"
1722
↕ Expand 10 unchanged lines (lines 23–32)
18#define BLOCK_SIZE (BLOCK_X * BLOCK_Y)23#define BLOCK_SIZE (BLOCK_X * BLOCK_Y)
19#define NUM_WARPS (BLOCK_SIZE / 32)24#define NUM_WARPS (BLOCK_SIZE / 32)
2025
21// Spherical harmonics coefficients26// Spherical harmonics coefficients
22__device__ const float SH_C0 = 0.28209479177387814f;27__device__ const float SH_C0 = 0.28209479177387814f;
23__device__ const float SH_C1 = 0.4886025119029199f;28__device__ const float SH_C1 = 0.4886025119029199f;
24__device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,29__device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,
25 0.31539156525252005f, -1.0925484305920792f,30 0.31539156525252005f, -1.0925484305920792f,
26 0.5462742152960396f};31 0.5462742152960396f};
27__device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,32__device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,
28 -0.4570457994644658f, 0.3731763325901154f,33 -0.4570457994644658f, 0.3731763325901154f,
29 -0.4570457994644658f, 1.445305721320277f,34 -0.4570457994644658f, 1.445305721320277f,
30 -0.5900435899266435f};35 -0.5900435899266435f};
3136
37namespace gs {
38
39#define GS_HD __host__ __device__ __forceinline__
40#define GS_D __device__ __forceinline__
41
42GS_HD float3 operator+(float3 a, float3 b) {
43 return {a.x + b.x, a.y + b.y, a.z + b.z};
44}
45GS_HD float3 operator-(float3 a, float3 b) {
46 return {a.x - b.x, a.y - b.y, a.z - b.z};
47}
48GS_HD float3 operator*(float s, float3 a) {
49 return {s * a.x, s * a.y, s * a.z};
50}
51GS_HD float3& operator+=(float3& a, float3 b) {
52 a.x += b.x;
53 a.y += b.y;
54 a.z += b.z;
55 return a;
56}
32__forceinline__ __device__ float ndc2Pix(float v, int S) {57GS_HD float dot(float3 a, float3 b) {
33 return ((v + 1.0) * S - 1.0) * 0.5;58 return a.x * b.x + a.y * b.y + a.z * b.z;
34}59}
3560
36__forceinline__ __device__ void getRect(const float2 p, int max_radius,
37 uint2& rect_min, uint2& rect_max,
38 dim3 grid) {
39 rect_min = {min(grid.x, max((int)0, (int)((p.x - max_radius) / BLOCK_X))),
40 min(grid.y, max((int)0, (int)((p.y - max_radius) / BLOCK_Y)))};
61// Row-major 3x3 matrix.
41 rect_max = {62struct Mat3 {
42 min(grid.x,
43 max((int)0, (int)((p.x + max_radius + BLOCK_X - 1) / BLOCK_X))),
44 min(grid.y,
45 max((int)0, (int)((p.y + max_radius + BLOCK_Y - 1) / BLOCK_Y)))};
63 float m[3][3];
64 GS_HD float3 row(int i) const { return {m[i][0], m[i][1], m[i][2]}; }
65 GS_HD float3 col(int j) const { return {m[0][j], m[1][j], m[2][j]}; }
66};
67
68GS_HD float3 Mul(const Mat3& a, float3 v) {
69 return {dot(a.row(0), v), dot(a.row(1), v), dot(a.row(2), v)};
70}
71// a^T v.
72GS_HD float3 MulT(const Mat3& a, float3 v) {
73 return {dot(a.col(0), v), dot(a.col(1), v), dot(a.col(2), v)};
46}74}
4775
48__forceinline__ __device__ void getRect(const float2 p, int2 ext_rect,
49 uint2& rect_min, uint2& rect_max,
50 dim3 grid) {
51 rect_min = {min(grid.x, max((int)0, (int)((p.x - ext_rect.x) / BLOCK_X))),
52 min(grid.y, max((int)0, (int)((p.y - ext_rect.y) / BLOCK_Y)))};
53 rect_max = {
54 min(grid.x,
55 max((int)0, (int)((p.x + ext_rect.x + BLOCK_X - 1) / BLOCK_X))),
56 min(grid.y,
57 max((int)0, (int)((p.y + ext_rect.y + BLOCK_Y - 1) / BLOCK_Y)))};
76// Reinterprets bits; used to pack small integers into float records.
77GS_HD float AsFloat(uint32_t u) {
78 union {
79 uint32_t u;
80 float f;
81 } c{u};
82 return c.f;
83}
84GS_HD uint32_t AsUint(float f) {
85 union {
86 float f;
87 uint32_t u;
88 } c{f};
89 return c.u;
58}90}
5991
60__forceinline__ __device__ float3 transformPoint4x3(const float3& p,
92// Number of floats in the packed camera buffer built by `camera.py`:
93// [R (9, row-major world-to-camera), t (3), fx, fy, cx, cy, position (3), 0].
94inline constexpr int kCameraFloats = 20;
95
96// Pinhole camera with OpenCV axes (+x right, +y down, +z forward). Pixel (i, j)
97// is sampled at coordinate (j, i), as in the Inria rasterizer.
61 const float* matrix) {98struct Camera {
99 Mat3 rotation; // World-to-camera rotation.
62 float3 transformed = {100 float3 translation;
63 matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12],
64 matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13],
65 matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14],
66 };
101 float fx, fy, cx, cy;
102 float3 position; // Camera center in world space.
103 int width, height;
104
105 GS_HD static Camera Load(const float* p, int width, int height) {
106 Camera c;
107 for (int i = 0; i < 3; ++i)
108 for (int j = 0; j < 3; ++j) c.rotation.m[i][j] = p[3 * i + j];
109 c.translation = {p[9], p[10], p[11]};
110 c.fx = p[12];
111 c.fy = p[13];
112 c.cx = p[14];
113 c.cy = p[15];
114 c.position = {p[16], p[17], p[18]};
115 c.width = width;
116 c.height = height;
67 return transformed;117 return c;
68}118 }
69119
70__forceinline__ __device__ float4 transformPoint4x4(const float3& p,
120 GS_HD float3 ToView(float3 p) const { return Mul(rotation, p) + translation; }
121 // Tangent of the half field of view, as the Inria rasterizer defines it.
122 GS_HD float TanHalfFovX() const { return width / (2.0f * fx); }
123 GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }
124};
125
126// Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block
127// rasterizes one tile.
128inline constexpr int kTileSize = 16;
129inline constexpr int kTilePixels = kTileSize * kTileSize;
130
131struct TileGrid {
132 int tiles_x, tiles_y;
133 GS_HD static TileGrid For(int width, int height) {
134 return {(width + kTileSize - 1) / kTileSize,
135 (height + kTileSize - 1) / kTileSize};
136 }
137 GS_HD int num_tiles() const { return tiles_x * tiles_y; }
138};
139
140// Smallest alpha that contributes to a pixel; smaller ones are skipped.
141inline constexpr float kMinAlpha = 1.0f / 255.0f;
71 const float* matrix) {142inline constexpr float kMaxAlpha = 0.99f;
143// Transmittance below which a pixel stops accumulating.
144inline constexpr float kMinTransmittance = 1e-4f;
145
146// One splat as stored in memory: three float4 so the rasterizer issues three
147// 16-byte loads. Layout:
148// p0 = (x, y, conic.a, conic.b)
149// p1 = (conic.c, opacity, limit, radius) radius 0 => culled
150// p2 = (r, g, b, min_power)
151inline constexpr int kSplatFloats = 12;
152struct alignas(16) Splat {
72 float4 transformed = {153 float4 p0, p1, p2;
73 matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12],
74 matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13],
75 matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14],
76 matrix[3] * p.x + matrix[7] * p.y + matrix[11] * p.z + matrix[15]};
154 GS_HD float2 xy() const { return {p0.x, p0.y}; }
155 GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }
156 GS_HD float opacity() const { return p1.y; }
157 GS_HD float limit() const { return p1.z; }
158 GS_HD int radius() const { return static_cast<int>(p1.w); }
159 GS_HD float3 rgb() const { return {p2.x, p2.y, p2.z}; }
160};
161
162// Gradient w.r.t. a splat's rasterized quantities, accumulated by the raster
163// backward pass: (dx, dy, da, db, dc, dopacity, dr, dg, db).
164inline constexpr int kSplatGradFloats = 9;
165
166// Half-open tile rectangle [x0, x1) x [y0, y1).
167struct TileRect {
168 int x0, y0, x1, y1;
169 GS_HD int area() const { return (x1 - x0) * (y1 - y0); }
170};
171
172// Tile rectangle of a splat of integer bounding `radius`, identical to the
77 return transformed;173// reference rasterizer.
174GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {
175 const float r = static_cast<float>(radius);
176 auto lo = [](float v, int n) {
177 return min(n, max(0, static_cast<int>(v / kTileSize)));
178 };
179 auto hi = [](float v, int n) {
180 return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));
181 };
182 return {lo(xy.x - r, grid.tiles_x), lo(xy.y - r, grid.tiles_y),
183 hi(xy.x + r, grid.tiles_x), hi(xy.y + r, grid.tiles_y)};
78}184}
79185
80__forceinline__ __device__ float3 transformVec4x3(const float3& p,
81 const float* matrix) {
82 float3 transformed = {
83 matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z,
84 matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z,
85 matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z,
186// How splats are assigned to tiles.
187enum class Culling : int {
188 // Every tile of the bounding rectangle (reference behaviour).
189 kRect = 0,
190 // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives
191 // the same image: skipped (splat, tile) pairs never contribute.
192 kEllipse = 1,
86 };193};
87 return transformed;
88}
89194
90__forceinline__ __device__ float3
91transformVec4x3Transpose(const float3& p, const float* matrix) {
92 float3 transformed = {
93 matrix[0] * p.x + matrix[1] * p.y + matrix[2] * p.z,
94 matrix[4] * p.x + matrix[5] * p.y + matrix[6] * p.z,
95 matrix[8] * p.x + matrix[9] * p.y + matrix[10] * p.z,
96 };
97 return transformed;
98}
99
100__forceinline__ __device__ float dnormvdz(float3 v, float3 dv) {
101 float sum2 = v.x * v.x + v.y * v.y + v.z * v.z;
102 float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
103 float dnormvdz =
104 (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) *
105 invsum32;
106 return dnormvdz;
107}
108
109__forceinline__ __device__ float3 dnormvdv(float3 v, float3 dv) {
110 float sum2 = v.x * v.x + v.y * v.y + v.z * v.z;
111 float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
112
113 float3 dnormvdv;
114 dnormvdv.x =
115 ((+sum2 - v.x * v.x) * dv.x - v.y * v.x * dv.y - v.z * v.x * dv.z) *
116 invsum32;
117 dnormvdv.y =
118 (-v.x * v.y * dv.x + (sum2 - v.y * v.y) * dv.y - v.z * v.y * dv.z) *
119 invsum32;
120 dnormvdv.z =
121 (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) *
122 invsum32;
123 return dnormvdv;
124}
125
126__forceinline__ __device__ float4 dnormvdv(float4 v, float4 dv) {
127 float sum2 = v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
128 float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
129
130 float4 vdv = {v.x * dv.x, v.y * dv.y, v.z * dv.z, v.w * dv.w};
131 float vdv_sum = vdv.x + vdv.y + vdv.z + vdv.w;
132 float4 dnormvdv;
133 dnormvdv.x = ((sum2 - v.x * v.x) * dv.x - v.x * (vdv_sum - vdv.x)) * invsum32;
134 dnormvdv.y = ((sum2 - v.y * v.y) * dv.y - v.y * (vdv_sum - vdv.y)) * invsum32;
135 dnormvdv.z = ((sum2 - v.z * v.z) * dv.z - v.z * (vdv_sum - vdv.z)) * invsum32;
136 dnormvdv.w = ((sum2 - v.w * v.w) * dv.w - v.w * (vdv_sum - vdv.w)) * invsum32;
137 return dnormvdv;
138}
139
140__forceinline__ __device__ float sigmoid(float x) {
141 return 1.0f / (1.0f + expf(-x));
142}
143
144__forceinline__ __device__ bool in_frustum(int idx, const float* orig_points,
145 const float* viewmatrix,
146 const float* projmatrix,
195// alpha = o exp(-q/2) >= 1/255 <=> q <= 2 ln(255 o). Writes the bound on q
196// with margins that keep tests conservative under float rounding of the
197// rasterizer's power. Returns false if no pixel can reach 1/255.
147 bool prefiltered, float3& p_view) {198GS_HD bool AlphaLimit(float opacity, float& limit) {
148 float3 p_orig = {orig_points[3 * idx], orig_points[3 * idx + 1],
149 orig_points[3 * idx + 2]};
150
151 // Bring points to screen space
152 float4 p_hom = transformPoint4x4(p_orig, projmatrix);
153 float p_w = 1.0f / (p_hom.w + 0.0000001f);
154 float3 p_proj = {p_hom.x * p_w, p_hom.y * p_w, p_hom.z * p_w};
155 p_view = transformPoint4x3(p_orig, viewmatrix);
156
157 if (p_view.z <= 0.2f) // || ((p_proj.x < -1.3 || p_proj.x > 1.3 || p_proj.y <
158 // -1.3 || p_proj.y > 1.3)))
159 {
160 if (prefiltered) {
161 printf(
162 "Point is filtered although prefiltered is set. This shouldn't "
163 "happen!");
164 __trap();
165 }
199 const float o255 = opacity * 255.0f;
166 return false;200 if (!(o255 >= 0.999f)) return false;
167 }
201 limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;
168 return true;202 return true;
169}203}
170204
205// Pad (pixels) on ellipse extents, against rounding.
206inline constexpr float kSpanPadPx = 0.01f;
207
208// Precomputed splat-level ellipse constants for fast band intersection tests
209// of {q(d) <= limit}, where q(d) = a dx^2 + 2 b dx dy + c dy^2.
210struct EllipseBounds {
211 float a, b, c;
212 float det;
213 float inv_a;
214 float a_limit;
215 float y_ext;
216 float dy_peak;
217
218 GS_D static bool Init(float3 conic, float opacity, float cached_limit,
219 EllipseBounds& out) {
220 float limit = cached_limit;
221 if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;
222 out.a = conic.x;
223 out.b = conic.y;
224 out.c = conic.z;
225 out.det = out.a * out.c - out.b * out.b;
226 if (!(out.det > 0.0f && out.a > 0.0f && out.c > 0.0f)) return false;
227 out.inv_a = 1.0f / out.a;
228 out.a_limit = out.a * limit;
229 const float limit_over_det = limit / out.det;
230 out.y_ext = sqrtf(out.a * limit_over_det);
231 out.dy_peak = out.b * sqrtf(limit_over_det / out.c);
232 return true;
233 }
234
235 // Span [x0_px, x1_px] of integer pixel columns hit across integer pixel rows
236 // Y in [y0_px, y1_px] (inclusive). Returns false if no integer pixel hits.
237 GS_D bool SpanForPixelRows(float2 xy, int y0_px, int y1_px, int& x0_px,
238 int& x1_px) const {
239 const int ylo_i =
240 max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));
241 const int yhi_i =
242 min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));
243 if (ylo_i > yhi_i) return false;
244 const float ylo = static_cast<float>(ylo_i) - xy.y;
245 const float yhi = static_cast<float>(yhi_i) - xy.y;
246 // Boundary x(dy) = (-b dy +- sqrt(a limit - det dy^2)) / a; the right
247 // (left) branch is concave (convex) and peaks at dy = -+dy_peak.
248 const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));
249 const float dyl = fminf(yhi, fmaxf(ylo, dy_peak));
250 const float dyr2 = dyr * dyr, dyl2 = dyl * dyl;
251 const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));
252 const float rl =
253 (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));
254 const float hi = (-b * dyr + rr) * inv_a;
255 const float lo = (-b * dyl - rl) * inv_a;
256 x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));
257 x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));
258 return x0_px <= x1_px;
259 }
260};
261
262// Visits every tile row `ty` with the half-open span [tx0, tx1) of tiles that
263// splat `s` (not culled) is binned to. Rows may be empty (tx0 >= tx1).
264template <Culling kCulling, typename Fn>
265GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,
266 int height, Fn&& fn) {
267 const float2 xy = s.xy();
268 const TileRect rect = RectFor(xy, s.radius(), grid);
269 if constexpr (kCulling == Culling::kRect) {
270 for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);
271 } else {
272 EllipseBounds eb;
273 if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;
274 const int min_y =
275 max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));
276 const int max_y =
277 min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));
278 if (min_y > max_y) return;
279 const int ty0 = max(rect.y0, min_y / kTileSize);
280 const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);
281 for (int ty = ty0; ty < ty1; ++ty) {
282 int x0_px, x1_px;
283 if (!eb.SpanForPixelRows(xy, ty * kTileSize,
284 min((ty + 1) * kTileSize, height) - 1, x0_px,
285 x1_px)) {
286 continue;
287 }
288 x0_px = max(0, x0_px);
289 x1_px = min(width - 1, x1_px);
290 if (x0_px > x1_px) continue;
291 const int tx0 = max(rect.x0, x0_px / kTileSize);
292 const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);
293 if (tx0 < tx1) fn(ty, tx0, tx1);
294 }
295 }
296}
297
298namespace warp {
299
300inline constexpr unsigned kFullMask = 0xffffffffu;
301
302// Standard warp butterfly reduction of an array of `N` per-lane values:
303// 5 steps of `__shfl_xor_sync` for each element (`5 * N` shuffles total),
304// after which lane 0 (and every other lane) holds the warp sum of `v[0..N)`.
305template <int N>
306GS_D void ButterflyReduce(float (&v)[N]) {
307#pragma unroll
308 for (int offset = 16; offset >= 1; offset >>= 1) {
309#pragma unroll
310 for (int i = 0; i < N; ++i) {
311 v[i] += __shfl_xor_sync(kFullMask, v[i], offset);
312 }
313 }
314}
315
316} // namespace warp
317
318inline constexpr int kShCoeffs = 16; // Degree 3.
319inline constexpr float kNearPlane = 0.2f;
320inline constexpr float kLowPass = 0.3f;
321inline constexpr float kFovClamp = 1.3f;
322
323// Pointers to the trainable parameters, all activated, row-major.
324struct GaussianParams {
325 const float* means; // (N, 3)
326 const float* scales; // (N, 3)
327 const float* quats; // (N, 4) as (w, x, y, z), not normalized.
328 const float* opacities; // (N,)
329 const void* sh; // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).
330 int n;
331};
332
333// Gradients w.r.t. `GaussianParams`, same layout and SH storage type.
334struct GaussianGrads {
335 float* means;
336 float* scales;
337 float* quats;
338 float* opacities;
339 void* sh;
340};
341
342// Screen-space quantities needed by rasterization. `conic` is the inverse 2D
343// covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.
344struct Projection {
345 float2 xy;
346 float depth;
347 float3 conic;
348 float opacity;
349 float limit;
350 float3 rgb;
351 int radius; // 0 means culled.
352};
353
354// ----------------------------------------------------------------------------
355// SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian
356// and may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.
357
358using ShBf16 = uint16_t; // bfloat16 bits.
359
360template <typename T>
361struct ShIo;
362
363template <>
364struct ShIo<float> {
365 GS_D static void LoadRow(const float* row, float3 coeffs[kShCoeffs]) {
366 const float4* p = reinterpret_cast<const float4*>(row);
367 float raw[3 * kShCoeffs];
368#pragma unroll
369 for (int k = 0; k < 12; ++k) {
370 const float4 q = __ldg(p + k);
371 raw[4 * k] = q.x;
372 raw[4 * k + 1] = q.y;
373 raw[4 * k + 2] = q.z;
374 raw[4 * k + 3] = q.w;
375 }
376#pragma unroll
377 for (int c = 0; c < kShCoeffs; ++c) {
378 coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
379 }
380 }
381 GS_D static void StoreRow(float* row, const float3 dsh[kShCoeffs]) {
382 float raw[3 * kShCoeffs];
383#pragma unroll
384 for (int c = 0; c < kShCoeffs; ++c) {
385 raw[3 * c] = dsh[c].x;
386 raw[3 * c + 1] = dsh[c].y;
387 raw[3 * c + 2] = dsh[c].z;
388 }
389 float4* p = reinterpret_cast<float4*>(row);
390#pragma unroll
391 for (int k = 0; k < 12; ++k) {
392 p[k] = {raw[4 * k], raw[4 * k + 1], raw[4 * k + 2], raw[4 * k + 3]};
393 }
394 }
395};
396
397template <>
398struct ShIo<ShBf16> {
399 GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
400 GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
401 // Round to nearest even (inputs are finite).
402 GS_D static uint32_t Round(float f) {
403 const uint32_t u = __float_as_uint(f);
404 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
405 }
406 GS_D static uint32_t Pack(float lo, float hi) {
407 return Round(lo) | (Round(hi) << 16);
408 }
409 GS_D static void LoadRow(const ShBf16* row, float3 coeffs[kShCoeffs]) {
410 const uint2* p = reinterpret_cast<const uint2*>(row);
411 float raw[3 * kShCoeffs];
412#pragma unroll
413 for (int k = 0; k < 12; ++k) {
414 const uint2 w = __ldg(p + k);
415 raw[4 * k] = Lo(w.x);
416 raw[4 * k + 1] = Hi(w.x);
417 raw[4 * k + 2] = Lo(w.y);
418 raw[4 * k + 3] = Hi(w.y);
419 }
420#pragma unroll
421 for (int c = 0; c < kShCoeffs; ++c) {
422 coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
423 }
424 }
425 GS_D static void StoreRow(ShBf16* row, const float3 dsh[kShCoeffs]) {
426 float raw[3 * kShCoeffs];
427#pragma unroll
428 for (int c = 0; c < kShCoeffs; ++c) {
429 raw[3 * c] = dsh[c].x;
430 raw[3 * c + 1] = dsh[c].y;
431 raw[3 * c + 2] = dsh[c].z;
432 }
433 uint2* p = reinterpret_cast<uint2*>(row);
434#pragma unroll
435 for (int k = 0; k < 12; ++k) {
436 p[k] = {Pack(raw[4 * k], raw[4 * k + 1]),
437 Pack(raw[4 * k + 2], raw[4 * k + 3])};
438 }
439 }
440};
441
442template <typename ShT>
443GS_D const ShT* ShRow(const void* sh, int i) {
444 return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
445}
446template <typename ShT>
447GS_D ShT* ShRow(void* sh, int i) {
448 return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
449}
450
451// ----------------------------------------------------------------------------
452// Spherical harmonics (degree 3, Inria sign convention).
453
454namespace sh {
455inline constexpr float kC0 = 0.28209479177387814f;
456inline constexpr float kC1 = 0.4886025119029199f;
457inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
458 0.31539156525252005f, -1.0925484305920792f,
459 0.5462742152960396f};
460inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
461 -0.4570457994644658f, 0.3731763325901154f,
462 -0.4570457994644658f, 1.445305721320277f,
463 -0.5900435899266435f};
464
465// Basis values for a unit direction.
466GS_D void Basis(float3 d, float b[kShCoeffs]) {
467 const float x = d.x, y = d.y, z = d.z;
468 const float xx = x * x, yy = y * y, zz = z * z;
469 b[0] = kC0;
470 b[1] = -kC1 * y;
471 b[2] = kC1 * z;
472 b[3] = -kC1 * x;
473 b[4] = kC2[0] * x * y;
474 b[5] = kC2[1] * y * z;
475 b[6] = kC2[2] * (2.0f * zz - xx - yy);
476 b[7] = kC2[3] * x * z;
477 b[8] = kC2[4] * (xx - yy);
478 b[9] = kC3[0] * y * (3.0f * xx - yy);
479 b[10] = kC3[1] * x * y * z;
480 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
481 b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
482 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
483 b[14] = kC3[5] * z * (xx - yy);
484 b[15] = kC3[6] * x * (xx - 3.0f * yy);
485}
486
487// Partial derivatives of every basis function w.r.t. (x, y, z) of the unit
488// direction (treated as free coordinates).
489GS_D void BasisGrad(float3 d, float3 g[kShCoeffs]) {
490 const float x = d.x, y = d.y, z = d.z;
491 const float xx = x * x, yy = y * y, zz = z * z;
492 g[0] = {0.0f, 0.0f, 0.0f};
493 g[1] = {0.0f, -kC1, 0.0f};
494 g[2] = {0.0f, 0.0f, kC1};
495 g[3] = {-kC1, 0.0f, 0.0f};
496 g[4] = kC2[0] * float3{y, x, 0.0f};
497 g[5] = kC2[1] * float3{0.0f, z, y};
498 g[6] = kC2[2] * float3{-2.0f * x, -2.0f * y, 4.0f * z};
499 g[7] = kC2[3] * float3{z, 0.0f, x};
500 g[8] = kC2[4] * float3{2.0f * x, -2.0f * y, 0.0f};
501 g[9] = kC3[0] * float3{6.0f * x * y, 3.0f * (xx - yy), 0.0f};
502 g[10] = kC3[1] * float3{y * z, x * z, x * y};
503 g[11] =
504 kC3[2] * float3{-2.0f * x * y, 4.0f * zz - xx - 3.0f * yy, 8.0f * y * z};
505 g[12] = kC3[3] * float3{-6.0f * x * z, -6.0f * y * z,
506 6.0f * zz - 3.0f * xx - 3.0f * yy};
507 g[13] =
508 kC3[4] * float3{4.0f * zz - 3.0f * xx - yy, -2.0f * x * y, 8.0f * x * z};
509 g[14] = kC3[5] * float3{2.0f * x * z, -2.0f * y * z, xx - yy};
510 g[15] = kC3[6] * float3{3.0f * (xx - yy), -6.0f * x * y, 0.0f};
511}
512
513// Pre-clamp color 0.5 + sum_k b_k sh_k, loading the full 16-basis SH row.
514template <typename ShT>
515GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
516 float3 coeffs[kShCoeffs];
517 ShIo<ShT>::LoadRow(sh_row, coeffs);
518 float3 c = {0.5f, 0.5f, 0.5f};
519#pragma unroll
520 for (int k = 0; k < kShCoeffs; ++k) {
521 c += b[k] * coeffs[k];
522 }
523 return c;
524}
525} // namespace sh
526
527// ----------------------------------------------------------------------------
528// Covariance helpers.
529
530// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
531GS_D Mat3 QuatToRotation(float4 q) {
532 const float w = q.x, x = q.y, y = q.z, z = q.w;
533 Mat3 r;
534 r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
535 r.m[0][1] = 2.0f * (x * y - w * z);
536 r.m[0][2] = 2.0f * (x * z + w * y);
537 r.m[1][0] = 2.0f * (x * y + w * z);
538 r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
539 r.m[1][2] = 2.0f * (y * z - w * x);
540 r.m[2][0] = 2.0f * (x * z - w * y);
541 r.m[2][1] = 2.0f * (y * z + w * x);
542 r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
543 return r;
544}
545
546// Symmetric 3x3 matrix.
547struct Sym3 {
548 float xx, xy, xz, yy, yz, zz;
549 GS_D float3 Mul(float3 v) const {
550 return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
551 xz * v.x + yz * v.y + zz * v.z};
552 }
553};
554
555// Sigma = A A^T with A = R diag(s).
556GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
557 Mat3 a;
558 const float sv[3] = {s.x, s.y, s.z};
559#pragma unroll
560 for (int i = 0; i < 3; ++i)
561#pragma unroll
562 for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
563 return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
564 dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
565 dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
566}
567
568// Intermediate values of the projection shared by forward and VJP.
569struct ProjectionState {
570 float3 t; // View-space mean.
571 float tx_c, ty_c; // Clamped x, y used by the Jacobian.
572 bool x_clamped, y_clamped;
573 float3 m0, m1; // Rows of M = J W (2x3).
574 Sym3 sigma; // World-space 3D covariance.
575 float a, b, c; // Dilated 2D covariance.
576};
577
578// Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
579GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
580 float4 quat, ProjectionState& st) {
581 st.t = t;
582 const float tz = t.z;
583 const float limx = kFovClamp * cam.TanHalfFovX();
584 const float limy = kFovClamp * cam.TanHalfFovY();
585 const float txtz = t.x / tz, tytz = t.y / tz;
586 st.x_clamped = txtz < -limx || txtz > limx;
587 st.y_clamped = tytz < -limy || tytz > limy;
588 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
589 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
590 const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
591 const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
592 const Mat3& w = cam.rotation;
593 st.m0 = j00 * w.row(0) + j02 * w.row(2);
594 st.m1 = j11 * w.row(1) + j12 * w.row(2);
595 st.sigma = Covariance3D(QuatToRotation(quat), scale);
596 const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
597 st.a = dot(st.m0, s0) + kLowPass;
598 st.b = dot(st.m0, s1);
599 st.c = dot(st.m1, s1) + kLowPass;
600}
601
602// Computes everything up to the 2D covariance. Returns false if culled by the
603// near plane.
604GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
605 float4 quat, ProjectionState& st) {
606 const float3 t = cam.ToView(mean);
607 if (t.z <= kNearPlane) return false;
608 ProjectCovarianceFromView(cam, t, scale, quat, st);
609 return true;
610}
611
612// Integer bounding radius (pixels) of the 2D covariance, as in the reference.
613GS_D int BoundingRadius(float a, float b, float c) {
614 const float det = a * c - b * b;
615 const float mid = 0.5f * (a + c);
616 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
617 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
618}
619
620GS_D float2 ProjectMean(const Camera& cam, float3 t) {
621 return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
622}
623
624GS_D float3 LoadMean(const GaussianParams& p, int i) {
625 return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
626}
627
628// Unit direction from the camera center to `mean`, and 1 / distance.
629GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
630 const float3 v = mean - cam.position;
631 inv_len = rsqrtf(dot(v, v));
632 return inv_len * v;
633}
634
635// Gradient of the loss w.r.t. one splat's screen-space quantities.
636struct ProjectionGrad {
637 float2 xy;
638 float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.
639 float opacity;
640 float3 rgb;
641};
642
643// Image geometry passed by value to every kernel.
644struct Frame {
645 const float* camera; // Device pointer to the packed camera.
646 int width, height;
647 TileGrid grid;
648 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
649};
650
651// ---------------------------------------------------------------------------
652// Per-Gaussian kernels run one thread per Gaussian.
653
654inline constexpr int kGaussianThreads = 128;
655
656inline int GaussianBlocks(int64_t n) {
657 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
658}
659
660// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
661inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
662inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
663
664GS_HD float3 Power2Conic(float a, float b, float c) {
665 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
666}
667
668GS_D float Exp2Approx(float x) {
669 float y;
670 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
671 return y;
672}
673
674// ---------------------------------------------------------------------------
675// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into
676// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
677// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
678
679template <int kRows>
680struct RasterLayout {
681 static_assert(kRows == 1 || kRows == 2 || kRows == 4);
682 static constexpr int kThreads = kTilePixels / kRows;
683 static constexpr int kWarps = kThreads / 32;
684 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
685 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
686 static constexpr int kWarpsX = kTileSize / kWarpWidth;
687 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
688
689 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
690 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
691};
692
693struct RasterPixels {
694 int x, y0;
695 template <int kRows>
696 GS_D static RasterPixels For(int tid) {
697 using L = RasterLayout<kRows>;
698 const int warp = tid / 32, lane = tid % 32;
699 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
700 lane % L::kWarpWidth,
701 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
702 (lane / L::kWarpWidth) * kRows};
703 }
704};
705
706// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
707// pixel region of warp w of this block's tile.
708template <int kRows>
709GS_D uint32_t WarpMask(float4 p0, float4 p1) {
710 using L = RasterLayout<kRows>;
711 constexpr int kWarpsY = L::kWarps / L::kWarpsX;
712 float limit = p1.z;
713 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
714 const float a = p0.z, b = p0.w, c = p1.x;
715 const float det = a * c - b * b;
716 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
717 const float l_over_det = limit / det;
718 const float rx = sqrtf(c * l_over_det);
719 const float ry = sqrtf(a * l_over_det);
720 const float2 xy = {p0.x, p0.y};
721 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
722 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
723 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
724 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
725 if (x0 > x1 || y0 > y1) return 0;
726 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
727 uint32_t mx = 0u;
728#pragma unroll
729 for (int wx = 0; wx < L::kWarpsX; ++wx) {
730 const int wx0 = tx + wx * L::kWarpWidth;
731 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
732 }
733 if (mx == 0u) return 0;
734 uint32_t my = 0u;
735#pragma unroll
736 for (int wy = 0; wy < kWarpsY; ++wy) {
737 const int wy0 = ty + wy * L::kWarpRows;
738 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
739 }
740 if (my == 0u) return 0;
741 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
742 uint32_t mask = 0u;
743#pragma unroll
744 for (int wy = 0; wy < kWarpsY; ++wy) {
745 if (!((my >> wy) & 1u)) continue;
746#pragma unroll
747 for (int wx = 0; wx < L::kWarpsX; ++wx) {
748 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
749 }
750 }
751 return mask;
752 }
753 EllipseBounds eb;
754 eb.a = a;
755 eb.b = b;
756 eb.c = c;
757 eb.det = det;
758 eb.inv_a = 1.0f / a;
759 eb.a_limit = a * limit;
760 eb.y_ext = ry;
761 eb.dy_peak = b * sqrtf(l_over_det / c);
762 uint32_t mask = 0u;
763#pragma unroll
764 for (int wy = 0; wy < kWarpsY; ++wy) {
765 if (!((my >> wy) & 1u)) continue;
766 const int wy0 = ty + wy * L::kWarpRows;
767 int x0_px, x1_px;
768 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
769 continue;
770 }
771#pragma unroll
772 for (int wx = 0; wx < L::kWarpsX; ++wx) {
773 const int wx0 = tx + wx * L::kWarpWidth;
774 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
775 mask |= 1u << (wy * L::kWarpsX + wx);
776 }
777 }
778 }
779 return mask;
780}
781
782} // namespace gs
783
171#endif784#endif
↕ Expand 9 unchanged lines (lines 1–9)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #ifndef CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
1313 #define CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
14+
15+#include <cuda_runtime.h>
16+
17+#include <cmath>
18+#include <cstdint>
1419
1520 #include "config.h"
1621 #include "stdio.h"
1722
↕ Expand 10 unchanged lines (lines 23–32)
1823 #define BLOCK_SIZE (BLOCK_X * BLOCK_Y)
1924 #define NUM_WARPS (BLOCK_SIZE / 32)
2025
2126 // Spherical harmonics coefficients
2227 __device__ const float SH_C0 = 0.28209479177387814f;
2328 __device__ const float SH_C1 = 0.4886025119029199f;
2429 __device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,
2530 0.31539156525252005f, -1.0925484305920792f,
2631 0.5462742152960396f};
2732 __device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,
2833 -0.4570457994644658f, 0.3731763325901154f,
2934 -0.4570457994644658f, 1.445305721320277f,
3035 -0.5900435899266435f};
3136
37+namespace gs {
38+
39+#define GS_HD __host__ __device__ __forceinline__
40+#define GS_D __device__ __forceinline__
41+
42+GS_HD float3 operator+(float3 a, float3 b) {
43+ return {a.x + b.x, a.y + b.y, a.z + b.z};
44+}
45+GS_HD float3 operator-(float3 a, float3 b) {
46+ return {a.x - b.x, a.y - b.y, a.z - b.z};
47+}
48+GS_HD float3 operator*(float s, float3 a) {
49+ return {s * a.x, s * a.y, s * a.z};
50+}
51+GS_HD float3& operator+=(float3& a, float3 b) {
52+ a.x += b.x;
53+ a.y += b.y;
54+ a.z += b.z;
55+ return a;
56+}
32−__forceinline__ __device__ float ndc2Pix(float v, int S) {
57+GS_HD float dot(float3 a, float3 b) {
33− return ((v + 1.0) * S - 1.0) * 0.5;
58+ return a.x * b.x + a.y * b.y + a.z * b.z;
3459 }
3560
36−__forceinline__ __device__ void getRect(const float2 p, int max_radius,
37− uint2& rect_min, uint2& rect_max,
38− dim3 grid) {
39− rect_min = {min(grid.x, max((int)0, (int)((p.x - max_radius) / BLOCK_X))),
40− min(grid.y, max((int)0, (int)((p.y - max_radius) / BLOCK_Y)))};
61+// Row-major 3x3 matrix.
41− rect_max = {
62+struct Mat3 {
42− min(grid.x,
43− max((int)0, (int)((p.x + max_radius + BLOCK_X - 1) / BLOCK_X))),
44− min(grid.y,
45− max((int)0, (int)((p.y + max_radius + BLOCK_Y - 1) / BLOCK_Y)))};
63+ float m[3][3];
64+ GS_HD float3 row(int i) const { return {m[i][0], m[i][1], m[i][2]}; }
65+ GS_HD float3 col(int j) const { return {m[0][j], m[1][j], m[2][j]}; }
66+};
67+
68+GS_HD float3 Mul(const Mat3& a, float3 v) {
69+ return {dot(a.row(0), v), dot(a.row(1), v), dot(a.row(2), v)};
70+}
71+// a^T v.
72+GS_HD float3 MulT(const Mat3& a, float3 v) {
73+ return {dot(a.col(0), v), dot(a.col(1), v), dot(a.col(2), v)};
4674 }
4775
48−__forceinline__ __device__ void getRect(const float2 p, int2 ext_rect,
49− uint2& rect_min, uint2& rect_max,
50− dim3 grid) {
51− rect_min = {min(grid.x, max((int)0, (int)((p.x - ext_rect.x) / BLOCK_X))),
52− min(grid.y, max((int)0, (int)((p.y - ext_rect.y) / BLOCK_Y)))};
53− rect_max = {
54− min(grid.x,
55− max((int)0, (int)((p.x + ext_rect.x + BLOCK_X - 1) / BLOCK_X))),
56− min(grid.y,
57− max((int)0, (int)((p.y + ext_rect.y + BLOCK_Y - 1) / BLOCK_Y)))};
76+// Reinterprets bits; used to pack small integers into float records.
77+GS_HD float AsFloat(uint32_t u) {
78+ union {
79+ uint32_t u;
80+ float f;
81+ } c{u};
82+ return c.f;
83+}
84+GS_HD uint32_t AsUint(float f) {
85+ union {
86+ float f;
87+ uint32_t u;
88+ } c{f};
89+ return c.u;
5890 }
5991
60−__forceinline__ __device__ float3 transformPoint4x3(const float3& p,
92+// Number of floats in the packed camera buffer built by `camera.py`:
93+// [R (9, row-major world-to-camera), t (3), fx, fy, cx, cy, position (3), 0].
94+inline constexpr int kCameraFloats = 20;
95+
96+// Pinhole camera with OpenCV axes (+x right, +y down, +z forward). Pixel (i, j)
97+// is sampled at coordinate (j, i), as in the Inria rasterizer.
61− const float* matrix) {
98+struct Camera {
99+ Mat3 rotation; // World-to-camera rotation.
62− float3 transformed = {
100+ float3 translation;
63− matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12],
64− matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13],
65− matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14],
66− };
101+ float fx, fy, cx, cy;
102+ float3 position; // Camera center in world space.
103+ int width, height;
104+
105+ GS_HD static Camera Load(const float* p, int width, int height) {
106+ Camera c;
107+ for (int i = 0; i < 3; ++i)
108+ for (int j = 0; j < 3; ++j) c.rotation.m[i][j] = p[3 * i + j];
109+ c.translation = {p[9], p[10], p[11]};
110+ c.fx = p[12];
111+ c.fy = p[13];
112+ c.cx = p[14];
113+ c.cy = p[15];
114+ c.position = {p[16], p[17], p[18]};
115+ c.width = width;
116+ c.height = height;
67− return transformed;
117+ return c;
68118⇥ }
69119
70−__forceinline__ __device__ float4 transformPoint4x4(const float3& p,
120+ GS_HD float3 ToView(float3 p) const { return Mul(rotation, p) + translation; }
121+ // Tangent of the half field of view, as the Inria rasterizer defines it.
122+ GS_HD float TanHalfFovX() const { return width / (2.0f * fx); }
123+ GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }
124+};
125+
126+// Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block
127+// rasterizes one tile.
128+inline constexpr int kTileSize = 16;
129+inline constexpr int kTilePixels = kTileSize * kTileSize;
130+
131+struct TileGrid {
132+ int tiles_x, tiles_y;
133+ GS_HD static TileGrid For(int width, int height) {
134+ return {(width + kTileSize - 1) / kTileSize,
135+ (height + kTileSize - 1) / kTileSize};
136+ }
137+ GS_HD int num_tiles() const { return tiles_x * tiles_y; }
138+};
139+
140+// Smallest alpha that contributes to a pixel; smaller ones are skipped.
141+inline constexpr float kMinAlpha = 1.0f / 255.0f;
71− const float* matrix) {
142+inline constexpr float kMaxAlpha = 0.99f;
143+// Transmittance below which a pixel stops accumulating.
144+inline constexpr float kMinTransmittance = 1e-4f;
145+
146+// One splat as stored in memory: three float4 so the rasterizer issues three
147+// 16-byte loads. Layout:
148+// p0 = (x, y, conic.a, conic.b)
149+// p1 = (conic.c, opacity, limit, radius) radius 0 => culled
150+// p2 = (r, g, b, min_power)
151+inline constexpr int kSplatFloats = 12;
152+struct alignas(16) Splat {
72− float4 transformed = {
153+ float4 p0, p1, p2;
73− matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z + matrix[12],
74− matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z + matrix[13],
75− matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z + matrix[14],
76− matrix[3] * p.x + matrix[7] * p.y + matrix[11] * p.z + matrix[15]};
154+ GS_HD float2 xy() const { return {p0.x, p0.y}; }
155+ GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }
156+ GS_HD float opacity() const { return p1.y; }
157+ GS_HD float limit() const { return p1.z; }
158+ GS_HD int radius() const { return static_cast<int>(p1.w); }
159+ GS_HD float3 rgb() const { return {p2.x, p2.y, p2.z}; }
160+};
161+
162+// Gradient w.r.t. a splat's rasterized quantities, accumulated by the raster
163+// backward pass: (dx, dy, da, db, dc, dopacity, dr, dg, db).
164+inline constexpr int kSplatGradFloats = 9;
165+
166+// Half-open tile rectangle [x0, x1) x [y0, y1).
167+struct TileRect {
168+ int x0, y0, x1, y1;
169+ GS_HD int area() const { return (x1 - x0) * (y1 - y0); }
170+};
171+
172+// Tile rectangle of a splat of integer bounding `radius`, identical to the
77− return transformed;
173+// reference rasterizer.
174+GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {
175+ const float r = static_cast<float>(radius);
176+ auto lo = [](float v, int n) {
177+ return min(n, max(0, static_cast<int>(v / kTileSize)));
178+ };
179+ auto hi = [](float v, int n) {
180+ return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));
181+ };
182+ return {lo(xy.x - r, grid.tiles_x), lo(xy.y - r, grid.tiles_y),
183+ hi(xy.x + r, grid.tiles_x), hi(xy.y + r, grid.tiles_y)};
78184 }
79185
80−__forceinline__ __device__ float3 transformVec4x3(const float3& p,
81− const float* matrix) {
82− float3 transformed = {
83− matrix[0] * p.x + matrix[4] * p.y + matrix[8] * p.z,
84− matrix[1] * p.x + matrix[5] * p.y + matrix[9] * p.z,
85− matrix[2] * p.x + matrix[6] * p.y + matrix[10] * p.z,
186+// How splats are assigned to tiles.
187+enum class Culling : int {
188+ // Every tile of the bounding rectangle (reference behaviour).
189+ kRect = 0,
190+ // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives
191+ // the same image: skipped (splat, tile) pairs never contribute.
192+ kEllipse = 1,
86193⇥};
87− return transformed;
88−}
89194
90−__forceinline__ __device__ float3
91−transformVec4x3Transpose(const float3& p, const float* matrix) {
92− float3 transformed = {
93− matrix[0] * p.x + matrix[1] * p.y + matrix[2] * p.z,
94− matrix[4] * p.x + matrix[5] * p.y + matrix[6] * p.z,
95− matrix[8] * p.x + matrix[9] * p.y + matrix[10] * p.z,
96− };
97− return transformed;
98−}
99−
100−__forceinline__ __device__ float dnormvdz(float3 v, float3 dv) {
101− float sum2 = v.x * v.x + v.y * v.y + v.z * v.z;
102− float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
103− float dnormvdz =
104− (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) *
105− invsum32;
106− return dnormvdz;
107−}
108−
109−__forceinline__ __device__ float3 dnormvdv(float3 v, float3 dv) {
110− float sum2 = v.x * v.x + v.y * v.y + v.z * v.z;
111− float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
112−
113− float3 dnormvdv;
114− dnormvdv.x =
115− ((+sum2 - v.x * v.x) * dv.x - v.y * v.x * dv.y - v.z * v.x * dv.z) *
116− invsum32;
117− dnormvdv.y =
118− (-v.x * v.y * dv.x + (sum2 - v.y * v.y) * dv.y - v.z * v.y * dv.z) *
119− invsum32;
120− dnormvdv.z =
121− (-v.x * v.z * dv.x - v.y * v.z * dv.y + (sum2 - v.z * v.z) * dv.z) *
122− invsum32;
123− return dnormvdv;
124−}
125−
126−__forceinline__ __device__ float4 dnormvdv(float4 v, float4 dv) {
127− float sum2 = v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
128− float invsum32 = 1.0f / sqrt(sum2 * sum2 * sum2);
129−
130− float4 vdv = {v.x * dv.x, v.y * dv.y, v.z * dv.z, v.w * dv.w};
131− float vdv_sum = vdv.x + vdv.y + vdv.z + vdv.w;
132− float4 dnormvdv;
133− dnormvdv.x = ((sum2 - v.x * v.x) * dv.x - v.x * (vdv_sum - vdv.x)) * invsum32;
134− dnormvdv.y = ((sum2 - v.y * v.y) * dv.y - v.y * (vdv_sum - vdv.y)) * invsum32;
135− dnormvdv.z = ((sum2 - v.z * v.z) * dv.z - v.z * (vdv_sum - vdv.z)) * invsum32;
136− dnormvdv.w = ((sum2 - v.w * v.w) * dv.w - v.w * (vdv_sum - vdv.w)) * invsum32;
137− return dnormvdv;
138−}
139−
140−__forceinline__ __device__ float sigmoid(float x) {
141− return 1.0f / (1.0f + expf(-x));
142−}
143−
144−__forceinline__ __device__ bool in_frustum(int idx, const float* orig_points,
145− const float* viewmatrix,
146− const float* projmatrix,
195+// alpha = o exp(-q/2) >= 1/255 <=> q <= 2 ln(255 o). Writes the bound on q
196+// with margins that keep tests conservative under float rounding of the
197+// rasterizer's power. Returns false if no pixel can reach 1/255.
147− bool prefiltered, float3& p_view) {
198+GS_HD bool AlphaLimit(float opacity, float& limit) {
148− float3 p_orig = {orig_points[3 * idx], orig_points[3 * idx + 1],
149− orig_points[3 * idx + 2]};
150−
151− // Bring points to screen space
152− float4 p_hom = transformPoint4x4(p_orig, projmatrix);
153− float p_w = 1.0f / (p_hom.w + 0.0000001f);
154− float3 p_proj = {p_hom.x * p_w, p_hom.y * p_w, p_hom.z * p_w};
155− p_view = transformPoint4x3(p_orig, viewmatrix);
156−
157− if (p_view.z <= 0.2f) // || ((p_proj.x < -1.3 || p_proj.x > 1.3 || p_proj.y <
158− // -1.3 || p_proj.y > 1.3)))
159− {
160− if (prefiltered) {
161− printf(
162− "Point is filtered although prefiltered is set. This shouldn't "
163− "happen!");
164− __trap();
165− }
199+ const float o255 = opacity * 255.0f;
166− return false;
200+ if (!(o255 >= 0.999f)) return false;
167− }
201+ limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;
168202 return true;
169203 }
170204
205+// Pad (pixels) on ellipse extents, against rounding.
206+inline constexpr float kSpanPadPx = 0.01f;
207+
208+// Precomputed splat-level ellipse constants for fast band intersection tests
209+// of {q(d) <= limit}, where q(d) = a dx^2 + 2 b dx dy + c dy^2.
210+struct EllipseBounds {
211+ float a, b, c;
212+ float det;
213+ float inv_a;
214+ float a_limit;
215+ float y_ext;
216+ float dy_peak;
217+
218+ GS_D static bool Init(float3 conic, float opacity, float cached_limit,
219+ EllipseBounds& out) {
220+ float limit = cached_limit;
221+ if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;
222+ out.a = conic.x;
223+ out.b = conic.y;
224+ out.c = conic.z;
225+ out.det = out.a * out.c - out.b * out.b;
226+ if (!(out.det > 0.0f && out.a > 0.0f && out.c > 0.0f)) return false;
227+ out.inv_a = 1.0f / out.a;
228+ out.a_limit = out.a * limit;
229+ const float limit_over_det = limit / out.det;
230+ out.y_ext = sqrtf(out.a * limit_over_det);
231+ out.dy_peak = out.b * sqrtf(limit_over_det / out.c);
232+ return true;
233+ }
234+
235+ // Span [x0_px, x1_px] of integer pixel columns hit across integer pixel rows
236+ // Y in [y0_px, y1_px] (inclusive). Returns false if no integer pixel hits.
237+ GS_D bool SpanForPixelRows(float2 xy, int y0_px, int y1_px, int& x0_px,
238+ int& x1_px) const {
239+ const int ylo_i =
240+ max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));
241+ const int yhi_i =
242+ min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));
243+ if (ylo_i > yhi_i) return false;
244+ const float ylo = static_cast<float>(ylo_i) - xy.y;
245+ const float yhi = static_cast<float>(yhi_i) - xy.y;
246+ // Boundary x(dy) = (-b dy +- sqrt(a limit - det dy^2)) / a; the right
247+ // (left) branch is concave (convex) and peaks at dy = -+dy_peak.
248+ const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));
249+ const float dyl = fminf(yhi, fmaxf(ylo, dy_peak));
250+ const float dyr2 = dyr * dyr, dyl2 = dyl * dyl;
251+ const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));
252+ const float rl =
253+ (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));
254+ const float hi = (-b * dyr + rr) * inv_a;
255+ const float lo = (-b * dyl - rl) * inv_a;
256+ x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));
257+ x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));
258+ return x0_px <= x1_px;
259+ }
260+};
261+
262+// Visits every tile row `ty` with the half-open span [tx0, tx1) of tiles that
263+// splat `s` (not culled) is binned to. Rows may be empty (tx0 >= tx1).
264+template <Culling kCulling, typename Fn>
265+GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,
266+ int height, Fn&& fn) {
267+ const float2 xy = s.xy();
268+ const TileRect rect = RectFor(xy, s.radius(), grid);
269+ if constexpr (kCulling == Culling::kRect) {
270+ for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);
271+ } else {
272+ EllipseBounds eb;
273+ if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;
274+ const int min_y =
275+ max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));
276+ const int max_y =
277+ min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));
278+ if (min_y > max_y) return;
279+ const int ty0 = max(rect.y0, min_y / kTileSize);
280+ const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);
281+ for (int ty = ty0; ty < ty1; ++ty) {
282+ int x0_px, x1_px;
283+ if (!eb.SpanForPixelRows(xy, ty * kTileSize,
284+ min((ty + 1) * kTileSize, height) - 1, x0_px,
285+ x1_px)) {
286+ continue;
287+ }
288+ x0_px = max(0, x0_px);
289+ x1_px = min(width - 1, x1_px);
290+ if (x0_px > x1_px) continue;
291+ const int tx0 = max(rect.x0, x0_px / kTileSize);
292+ const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);
293+ if (tx0 < tx1) fn(ty, tx0, tx1);
294+ }
295+ }
296+}
297+
298+namespace warp {
299+
300+inline constexpr unsigned kFullMask = 0xffffffffu;
301+
302+// Standard warp butterfly reduction of an array of `N` per-lane values:
303+// 5 steps of `__shfl_xor_sync` for each element (`5 * N` shuffles total),
304+// after which lane 0 (and every other lane) holds the warp sum of `v[0..N)`.
305+template <int N>
306+GS_D void ButterflyReduce(float (&v)[N]) {
307+#pragma unroll
308+ for (int offset = 16; offset >= 1; offset >>= 1) {
309+#pragma unroll
310+ for (int i = 0; i < N; ++i) {
311+ v[i] += __shfl_xor_sync(kFullMask, v[i], offset);
312+ }
313+ }
314+}
315+
316+} // namespace warp
317+
318+inline constexpr int kShCoeffs = 16; // Degree 3.
319+inline constexpr float kNearPlane = 0.2f;
320+inline constexpr float kLowPass = 0.3f;
321+inline constexpr float kFovClamp = 1.3f;
322+
323+// Pointers to the trainable parameters, all activated, row-major.
324+struct GaussianParams {
325+ const float* means; // (N, 3)
326+ const float* scales; // (N, 3)
327+ const float* quats; // (N, 4) as (w, x, y, z), not normalized.
328+ const float* opacities; // (N,)
329+ const void* sh; // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).
330+ int n;
331+};
332+
333+// Gradients w.r.t. `GaussianParams`, same layout and SH storage type.
334+struct GaussianGrads {
335+ float* means;
336+ float* scales;
337+ float* quats;
338+ float* opacities;
339+ void* sh;
340+};
341+
342+// Screen-space quantities needed by rasterization. `conic` is the inverse 2D
343+// covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.
344+struct Projection {
345+ float2 xy;
346+ float depth;
347+ float3 conic;
348+ float opacity;
349+ float limit;
350+ float3 rgb;
351+ int radius; // 0 means culled.
352+};
353+
354+// ----------------------------------------------------------------------------
355+// SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian
356+// and may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.
357+
358+using ShBf16 = uint16_t; // bfloat16 bits.
359+
360+template <typename T>
361+struct ShIo;
362+
363+template <>
364+struct ShIo<float> {
365+ GS_D static void LoadRow(const float* row, float3 coeffs[kShCoeffs]) {
366+ const float4* p = reinterpret_cast<const float4*>(row);
367+ float raw[3 * kShCoeffs];
368+#pragma unroll
369+ for (int k = 0; k < 12; ++k) {
370+ const float4 q = __ldg(p + k);
371+ raw[4 * k] = q.x;
372+ raw[4 * k + 1] = q.y;
373+ raw[4 * k + 2] = q.z;
374+ raw[4 * k + 3] = q.w;
375+ }
376+#pragma unroll
377+ for (int c = 0; c < kShCoeffs; ++c) {
378+ coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
379+ }
380+ }
381+ GS_D static void StoreRow(float* row, const float3 dsh[kShCoeffs]) {
382+ float raw[3 * kShCoeffs];
383+#pragma unroll
384+ for (int c = 0; c < kShCoeffs; ++c) {
385+ raw[3 * c] = dsh[c].x;
386+ raw[3 * c + 1] = dsh[c].y;
387+ raw[3 * c + 2] = dsh[c].z;
388+ }
389+ float4* p = reinterpret_cast<float4*>(row);
390+#pragma unroll
391+ for (int k = 0; k < 12; ++k) {
392+ p[k] = {raw[4 * k], raw[4 * k + 1], raw[4 * k + 2], raw[4 * k + 3]};
393+ }
394+ }
395+};
396+
397+template <>
398+struct ShIo<ShBf16> {
399+ GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
400+ GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
401+ // Round to nearest even (inputs are finite).
402+ GS_D static uint32_t Round(float f) {
403+ const uint32_t u = __float_as_uint(f);
404+ return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
405+ }
406+ GS_D static uint32_t Pack(float lo, float hi) {
407+ return Round(lo) | (Round(hi) << 16);
408+ }
409+ GS_D static void LoadRow(const ShBf16* row, float3 coeffs[kShCoeffs]) {
410+ const uint2* p = reinterpret_cast<const uint2*>(row);
411+ float raw[3 * kShCoeffs];
412+#pragma unroll
413+ for (int k = 0; k < 12; ++k) {
414+ const uint2 w = __ldg(p + k);
415+ raw[4 * k] = Lo(w.x);
416+ raw[4 * k + 1] = Hi(w.x);
417+ raw[4 * k + 2] = Lo(w.y);
418+ raw[4 * k + 3] = Hi(w.y);
419+ }
420+#pragma unroll
421+ for (int c = 0; c < kShCoeffs; ++c) {
422+ coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
423+ }
424+ }
425+ GS_D static void StoreRow(ShBf16* row, const float3 dsh[kShCoeffs]) {
426+ float raw[3 * kShCoeffs];
427+#pragma unroll
428+ for (int c = 0; c < kShCoeffs; ++c) {
429+ raw[3 * c] = dsh[c].x;
430+ raw[3 * c + 1] = dsh[c].y;
431+ raw[3 * c + 2] = dsh[c].z;
432+ }
433+ uint2* p = reinterpret_cast<uint2*>(row);
434+#pragma unroll
435+ for (int k = 0; k < 12; ++k) {
436+ p[k] = {Pack(raw[4 * k], raw[4 * k + 1]),
437+ Pack(raw[4 * k + 2], raw[4 * k + 3])};
438+ }
439+ }
440+};
441+
442+template <typename ShT>
443+GS_D const ShT* ShRow(const void* sh, int i) {
444+ return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
445+}
446+template <typename ShT>
447+GS_D ShT* ShRow(void* sh, int i) {
448+ return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
449+}
450+
451+// ----------------------------------------------------------------------------
452+// Spherical harmonics (degree 3, Inria sign convention).
453+
454+namespace sh {
455+inline constexpr float kC0 = 0.28209479177387814f;
456+inline constexpr float kC1 = 0.4886025119029199f;
457+inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
458+ 0.31539156525252005f, -1.0925484305920792f,
459+ 0.5462742152960396f};
460+inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
461+ -0.4570457994644658f, 0.3731763325901154f,
462+ -0.4570457994644658f, 1.445305721320277f,
463+ -0.5900435899266435f};
464+
465+// Basis values for a unit direction.
466+GS_D void Basis(float3 d, float b[kShCoeffs]) {
467+ const float x = d.x, y = d.y, z = d.z;
468+ const float xx = x * x, yy = y * y, zz = z * z;
469+ b[0] = kC0;
470+ b[1] = -kC1 * y;
471+ b[2] = kC1 * z;
472+ b[3] = -kC1 * x;
473+ b[4] = kC2[0] * x * y;
474+ b[5] = kC2[1] * y * z;
475+ b[6] = kC2[2] * (2.0f * zz - xx - yy);
476+ b[7] = kC2[3] * x * z;
477+ b[8] = kC2[4] * (xx - yy);
478+ b[9] = kC3[0] * y * (3.0f * xx - yy);
479+ b[10] = kC3[1] * x * y * z;
480+ b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
481+ b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
482+ b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
483+ b[14] = kC3[5] * z * (xx - yy);
484+ b[15] = kC3[6] * x * (xx - 3.0f * yy);
485+}
486+
487+// Partial derivatives of every basis function w.r.t. (x, y, z) of the unit
488+// direction (treated as free coordinates).
489+GS_D void BasisGrad(float3 d, float3 g[kShCoeffs]) {
490+ const float x = d.x, y = d.y, z = d.z;
491+ const float xx = x * x, yy = y * y, zz = z * z;
492+ g[0] = {0.0f, 0.0f, 0.0f};
493+ g[1] = {0.0f, -kC1, 0.0f};
494+ g[2] = {0.0f, 0.0f, kC1};
495+ g[3] = {-kC1, 0.0f, 0.0f};
496+ g[4] = kC2[0] * float3{y, x, 0.0f};
497+ g[5] = kC2[1] * float3{0.0f, z, y};
498+ g[6] = kC2[2] * float3{-2.0f * x, -2.0f * y, 4.0f * z};
499+ g[7] = kC2[3] * float3{z, 0.0f, x};
500+ g[8] = kC2[4] * float3{2.0f * x, -2.0f * y, 0.0f};
501+ g[9] = kC3[0] * float3{6.0f * x * y, 3.0f * (xx - yy), 0.0f};
502+ g[10] = kC3[1] * float3{y * z, x * z, x * y};
503+ g[11] =
504+ kC3[2] * float3{-2.0f * x * y, 4.0f * zz - xx - 3.0f * yy, 8.0f * y * z};
505+ g[12] = kC3[3] * float3{-6.0f * x * z, -6.0f * y * z,
506+ 6.0f * zz - 3.0f * xx - 3.0f * yy};
507+ g[13] =
508+ kC3[4] * float3{4.0f * zz - 3.0f * xx - yy, -2.0f * x * y, 8.0f * x * z};
509+ g[14] = kC3[5] * float3{2.0f * x * z, -2.0f * y * z, xx - yy};
510+ g[15] = kC3[6] * float3{3.0f * (xx - yy), -6.0f * x * y, 0.0f};
511+}
512+
513+// Pre-clamp color 0.5 + sum_k b_k sh_k, loading the full 16-basis SH row.
514+template <typename ShT>
515+GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
516+ float3 coeffs[kShCoeffs];
517+ ShIo<ShT>::LoadRow(sh_row, coeffs);
518+ float3 c = {0.5f, 0.5f, 0.5f};
519+#pragma unroll
520+ for (int k = 0; k < kShCoeffs; ++k) {
521+ c += b[k] * coeffs[k];
522+ }
523+ return c;
524+}
525+} // namespace sh
526+
527+// ----------------------------------------------------------------------------
528+// Covariance helpers.
529+
530+// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
531+GS_D Mat3 QuatToRotation(float4 q) {
532+ const float w = q.x, x = q.y, y = q.z, z = q.w;
533+ Mat3 r;
534+ r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
535+ r.m[0][1] = 2.0f * (x * y - w * z);
536+ r.m[0][2] = 2.0f * (x * z + w * y);
537+ r.m[1][0] = 2.0f * (x * y + w * z);
538+ r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
539+ r.m[1][2] = 2.0f * (y * z - w * x);
540+ r.m[2][0] = 2.0f * (x * z - w * y);
541+ r.m[2][1] = 2.0f * (y * z + w * x);
542+ r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
543+ return r;
544+}
545+
546+// Symmetric 3x3 matrix.
547+struct Sym3 {
548+ float xx, xy, xz, yy, yz, zz;
549+ GS_D float3 Mul(float3 v) const {
550+ return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
551+ xz * v.x + yz * v.y + zz * v.z};
552+ }
553+};
554+
555+// Sigma = A A^T with A = R diag(s).
556+GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
557+ Mat3 a;
558+ const float sv[3] = {s.x, s.y, s.z};
559+#pragma unroll
560+ for (int i = 0; i < 3; ++i)
561+#pragma unroll
562+ for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
563+ return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
564+ dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
565+ dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
566+}
567+
568+// Intermediate values of the projection shared by forward and VJP.
569+struct ProjectionState {
570+ float3 t; // View-space mean.
571+ float tx_c, ty_c; // Clamped x, y used by the Jacobian.
572+ bool x_clamped, y_clamped;
573+ float3 m0, m1; // Rows of M = J W (2x3).
574+ Sym3 sigma; // World-space 3D covariance.
575+ float a, b, c; // Dilated 2D covariance.
576+};
577+
578+// Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
579+GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
580+ float4 quat, ProjectionState& st) {
581+ st.t = t;
582+ const float tz = t.z;
583+ const float limx = kFovClamp * cam.TanHalfFovX();
584+ const float limy = kFovClamp * cam.TanHalfFovY();
585+ const float txtz = t.x / tz, tytz = t.y / tz;
586+ st.x_clamped = txtz < -limx || txtz > limx;
587+ st.y_clamped = tytz < -limy || tytz > limy;
588+ st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
589+ st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
590+ const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
591+ const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
592+ const Mat3& w = cam.rotation;
593+ st.m0 = j00 * w.row(0) + j02 * w.row(2);
594+ st.m1 = j11 * w.row(1) + j12 * w.row(2);
595+ st.sigma = Covariance3D(QuatToRotation(quat), scale);
596+ const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
597+ st.a = dot(st.m0, s0) + kLowPass;
598+ st.b = dot(st.m0, s1);
599+ st.c = dot(st.m1, s1) + kLowPass;
600+}
601+
602+// Computes everything up to the 2D covariance. Returns false if culled by the
603+// near plane.
604+GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
605+ float4 quat, ProjectionState& st) {
606+ const float3 t = cam.ToView(mean);
607+ if (t.z <= kNearPlane) return false;
608+ ProjectCovarianceFromView(cam, t, scale, quat, st);
609+ return true;
610+}
611+
612+// Integer bounding radius (pixels) of the 2D covariance, as in the reference.
613+GS_D int BoundingRadius(float a, float b, float c) {
614+ const float det = a * c - b * b;
615+ const float mid = 0.5f * (a + c);
616+ const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
617+ return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
618+}
619+
620+GS_D float2 ProjectMean(const Camera& cam, float3 t) {
621+ return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
622+}
623+
624+GS_D float3 LoadMean(const GaussianParams& p, int i) {
625+ return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
626+}
627+
628+// Unit direction from the camera center to `mean`, and 1 / distance.
629+GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
630+ const float3 v = mean - cam.position;
631+ inv_len = rsqrtf(dot(v, v));
632+ return inv_len * v;
633+}
634+
635+// Gradient of the loss w.r.t. one splat's screen-space quantities.
636+struct ProjectionGrad {
637+ float2 xy;
638+ float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.
639+ float opacity;
640+ float3 rgb;
641+};
642+
643+// Image geometry passed by value to every kernel.
644+struct Frame {
645+ const float* camera; // Device pointer to the packed camera.
646+ int width, height;
647+ TileGrid grid;
648+ GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
649+};
650+
651+// ---------------------------------------------------------------------------
652+// Per-Gaussian kernels run one thread per Gaussian.
653+
654+inline constexpr int kGaussianThreads = 128;
655+
656+inline int GaussianBlocks(int64_t n) {
657+ return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
658+}
659+
660+// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
661+inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
662+inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
663+
664+GS_HD float3 Power2Conic(float a, float b, float c) {
665+ return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
666+}
667+
668+GS_D float Exp2Approx(float x) {
669+ float y;
670+ asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
671+ return y;
672+}
673+
674+// ---------------------------------------------------------------------------
675+// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into
676+// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
677+// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
678+
679+template <int kRows>
680+struct RasterLayout {
681+ static_assert(kRows == 1 || kRows == 2 || kRows == 4);
682+ static constexpr int kThreads = kTilePixels / kRows;
683+ static constexpr int kWarps = kThreads / 32;
684+ static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
685+ static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
686+ static constexpr int kWarpsX = kTileSize / kWarpWidth;
687+ static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
688+
689+ GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
690+ GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
691+};
692+
693+struct RasterPixels {
694+ int x, y0;
695+ template <int kRows>
696+ GS_D static RasterPixels For(int tid) {
697+ using L = RasterLayout<kRows>;
698+ const int warp = tid / 32, lane = tid % 32;
699+ return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
700+ lane % L::kWarpWidth,
701+ static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
702+ (lane / L::kWarpWidth) * kRows};
703+ }
704+};
705+
706+// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
707+// pixel region of warp w of this block's tile.
708+template <int kRows>
709+GS_D uint32_t WarpMask(float4 p0, float4 p1) {
710+ using L = RasterLayout<kRows>;
711+ constexpr int kWarpsY = L::kWarps / L::kWarpsX;
712+ float limit = p1.z;
713+ if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
714+ const float a = p0.z, b = p0.w, c = p1.x;
715+ const float det = a * c - b * b;
716+ if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
717+ const float l_over_det = limit / det;
718+ const float rx = sqrtf(c * l_over_det);
719+ const float ry = sqrtf(a * l_over_det);
720+ const float2 xy = {p0.x, p0.y};
721+ const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
722+ const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
723+ const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
724+ const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
725+ if (x0 > x1 || y0 > y1) return 0;
726+ const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
727+ uint32_t mx = 0u;
728+#pragma unroll
729+ for (int wx = 0; wx < L::kWarpsX; ++wx) {
730+ const int wx0 = tx + wx * L::kWarpWidth;
731+ if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
732+ }
733+ if (mx == 0u) return 0;
734+ uint32_t my = 0u;
735+#pragma unroll
736+ for (int wy = 0; wy < kWarpsY; ++wy) {
737+ const int wy0 = ty + wy * L::kWarpRows;
738+ if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
739+ }
740+ if (my == 0u) return 0;
741+ if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
742+ uint32_t mask = 0u;
743+#pragma unroll
744+ for (int wy = 0; wy < kWarpsY; ++wy) {
745+ if (!((my >> wy) & 1u)) continue;
746+#pragma unroll
747+ for (int wx = 0; wx < L::kWarpsX; ++wx) {
748+ if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
749+ }
750+ }
751+ return mask;
752+ }
753+ EllipseBounds eb;
754+ eb.a = a;
755+ eb.b = b;
756+ eb.c = c;
757+ eb.det = det;
758+ eb.inv_a = 1.0f / a;
759+ eb.a_limit = a * limit;
760+ eb.y_ext = ry;
761+ eb.dy_peak = b * sqrtf(l_over_det / c);
762+ uint32_t mask = 0u;
763+#pragma unroll
764+ for (int wy = 0; wy < kWarpsY; ++wy) {
765+ if (!((my >> wy) & 1u)) continue;
766+ const int wy0 = ty + wy * L::kWarpRows;
767+ int x0_px, x1_px;
768+ if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
769+ continue;
770+ }
771+#pragma unroll
772+ for (int wx = 0; wx < L::kWarpsX; ++wx) {
773+ const int wx0 = tx + wx * L::kWarpWidth;
774+ if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
775+ mask |= 1u << (wy * L::kWarpsX + wx);
776+ }
777+ }
778+ }
779+ return mask;
780+}
781+
782+} // namespace gs
783+
171784 #endif
csrc/forward.h +18 −27
s00_baseline/csrc/forward.h
s01_literature/csrc/forward.h
↕ Expand 10 unchanged lines (lines 1–10)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#ifndef CUDA_RASTERIZER_FORWARD_H_INCLUDED12#ifndef CUDA_RASTERIZER_FORWARD_H_INCLUDED
13#define CUDA_RASTERIZER_FORWARD_H_INCLUDED13#define CUDA_RASTERIZER_FORWARD_H_INCLUDED
1414
15#include <cuda.h>15#include "auxiliary.h"
1616
17#include "cuda_runtime.h"
18#include "device_launch_parameters.h"
19#define GLM_FORCE_CUDA
20#include <glm/glm.hpp>
17namespace gs {
2118
22namespace FORWARD {
23// Perform initial steps for each Gaussian prior to rasterization.
24void preprocess(int P, int D, int M, const float* orig_points,
25 const glm::vec3* scales, const float scale_modifier,
26 const glm::vec4* rotations, const float* opacities,
27 const float* shs, bool* clamped, const float* cov3D_precomp,
28 const float* colors_precomp, const float* viewmatrix,
29 const float* projmatrix, const glm::vec3* cam_pos, const int W,
30 int H, const float focal_x, float focal_y, const float tan_fovx,
31 float tan_fovy, const float low_pass_2d_kernel_size,
32 const bool mip_splatting_enabled, int* radii,
33 float2* points_xy_image, float* depths, float* cov3Ds,
34 float* colors, float4* conic_opacity, const dim3 grid,
19template <Culling kCulling, typename ShT>
20__global__ void __launch_bounds__(kGaussianThreads, 8)
21 PreprocessKernel(GaussianParams params, Frame frame,
22 Splat* __restrict__ splats,
35 uint32_t* tiles_touched, bool prefiltered, int2* rects,23 int32_t* __restrict__ tile_counts,
36 float3 boxmin, float3 boxmax);
24 uint32_t* __restrict__ depth_keys,
25 int32_t* __restrict__ ids,
26 int32_t* __restrict__ num_visible);
3727
38// Main rasterization method.
39void render(const dim3 grid, dim3 block, const uint2* ranges,
28__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
29 RasterForwardKernel(const int2* __restrict__ ranges,
40 const uint32_t* point_list, int W, int H,30 const int32_t* __restrict__ point_list,
31 const Splat* __restrict__ splats, Frame frame,
41 const float2* points_xy_image, const float* features,32 float* __restrict__ image, float* __restrict__ final_t,
42 const float4* conic_opacity, float* final_T, uint32_t* n_contrib,
43 const float* bg_color, float* out_color);
33 int32_t* __restrict__ n_contrib);
34
44} // namespace FORWARD35} // namespace gs
4536
46#endif37#endif
↕ Expand 10 unchanged lines (lines 1–10)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #ifndef CUDA_RASTERIZER_FORWARD_H_INCLUDED
1313 #define CUDA_RASTERIZER_FORWARD_H_INCLUDED
1414
15−#include <cuda.h>
15+#include "auxiliary.h"
1616
17−#include "cuda_runtime.h"
18−#include "device_launch_parameters.h"
19−#define GLM_FORCE_CUDA
20−#include <glm/glm.hpp>
17+namespace gs {
2118
22−namespace FORWARD {
23−// Perform initial steps for each Gaussian prior to rasterization.
24−void preprocess(int P, int D, int M, const float* orig_points,
25− const glm::vec3* scales, const float scale_modifier,
26− const glm::vec4* rotations, const float* opacities,
27− const float* shs, bool* clamped, const float* cov3D_precomp,
28− const float* colors_precomp, const float* viewmatrix,
29− const float* projmatrix, const glm::vec3* cam_pos, const int W,
30− int H, const float focal_x, float focal_y, const float tan_fovx,
31− float tan_fovy, const float low_pass_2d_kernel_size,
32− const bool mip_splatting_enabled, int* radii,
33− float2* points_xy_image, float* depths, float* cov3Ds,
34− float* colors, float4* conic_opacity, const dim3 grid,
19+template <Culling kCulling, typename ShT>
20+__global__ void __launch_bounds__(kGaussianThreads, 8)
21+ PreprocessKernel(GaussianParams params, Frame frame,
22+ Splat* __restrict__ splats,
35− uint32_t* tiles_touched, bool prefiltered, int2* rects,
23+ int32_t* __restrict__ tile_counts,
36− float3 boxmin, float3 boxmax);
24+ uint32_t* __restrict__ depth_keys,
25+ int32_t* __restrict__ ids,
26+ int32_t* __restrict__ num_visible);
3727
38−// Main rasterization method.
39−void render(const dim3 grid, dim3 block, const uint2* ranges,
28+__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
29+ RasterForwardKernel(const int2* __restrict__ ranges,
40− const uint32_t* point_list, int W, int H,
30+ const int32_t* __restrict__ point_list,
31+ const Splat* __restrict__ splats, Frame frame,
41− const float2* points_xy_image, const float* features,
32+ float* __restrict__ image, float* __restrict__ final_t,
42− const float4* conic_opacity, float* final_T, uint32_t* n_contrib,
43− const float* bg_color, float* out_color);
33+ int32_t* __restrict__ n_contrib);
34+
44−} // namespace FORWARD
35+} // namespace gs
4536
4637 #endif
csrc/forward.cu.cc +188 −380
s00_baseline/csrc/forward.cu.cc
s01_literature/csrc/forward.cu.cc
↕ Expand 9 unchanged lines (lines 1–9)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#include "forward.h"12#include "forward.h"
1313
14#include <cooperative_groups.h>14#include "auxiliary.h"
15#include <cooperative_groups/reduce.h>
1615
17#include "auxiliary.h"
18namespace cg = cooperative_groups;16namespace gs {
1917
20// Forward method for converting the input spherical harmonics
21// coefficients of each Gaussian to a simple RGB color.
22__device__ glm::vec3 computeColorFromSH(int idx, int deg, int max_coeffs,
18// Screen-space geometry of Gaussian `i`: everything in `Projection` except
19// `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
20// tile-rectangle cull).
21template <Culling kCulling>
22GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
23 float3 mean, int i) {
24 Projection out;
25 out.radius = 0;
23 const glm::vec3* means,26 const float3 t = cam.ToView(mean);
24 glm::vec3 campos, const float* shs,
25 bool* clamped) {
26 // The implementation is loosely based on code for
27 // "Differentiable Point-Based Radiance Fields for
28 // Efficient View Synthesis" by Zhang et al. (2022)
29 glm::vec3 pos = means[idx];
30 glm::vec3 dir = pos - campos;
31 dir = dir / glm::length(dir);
27 if (t.z <= kNearPlane) return out;
28 const float opacity = p.opacities[i];
29 float limit = 0.0f;
30 if (!AlphaLimit(opacity, limit)) {
31 if constexpr (kCulling == Culling::kEllipse) return out;
32 }
33 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
34 p.scales[3 * i + 2]};
35 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
36 ProjectionState st;
37 ProjectCovarianceFromView(cam, t, scale, quat, st);
38 const float det = st.a * st.c - st.b * st.b;
39 if (det == 0.0f) return out;
40 const float inv = 1.0f / det;
41 out.conic = {st.c * inv, -st.b * inv, st.a * inv};
42 out.xy = ProjectMean(cam, st.t);
43 out.depth = st.t.z;
44 out.opacity = opacity;
45 out.limit = limit;
46 out.radius = BoundingRadius(st.a, st.b, st.c);
47 return out;
48}
3249
33 glm::vec3* sh = ((glm::vec3*)shs) + idx * max_coeffs;
34 glm::vec3 result = SH_C0 * sh[0];
50// Color seen along `dir`, clamped at zero, reading from an SH row.
51template <typename ShT>
52GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
53 float basis[kShCoeffs];
54 sh::Basis(dir, basis);
55 const float3 raw = sh::EvalRow(sh_row, basis);
56 return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
57}
3558
36 if (deg > 0) {
37 float x = dir.x;
38 float y = dir.y;
39 float z = dir.z;
40 result = result - SH_C1 * y * sh[1] + SH_C1 * z * sh[2] - SH_C1 * x * sh[3];
59// ---------------------------------------------------------------------------
60// Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
61// evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
62// depth sort via per-thread `atomicAdd`.
4163
64template <Culling kCulling, typename ShT>
65__global__ void __launch_bounds__(kGaussianThreads, 8)
66 PreprocessKernel(GaussianParams params, Frame frame,
67 Splat* __restrict__ splats,
68 int32_t* __restrict__ tile_counts,
69 uint32_t* __restrict__ depth_keys,
70 int32_t* __restrict__ ids,
71 int32_t* __restrict__ num_visible) {
72 const int i = blockIdx.x * blockDim.x + threadIdx.x;
73 if (i >= params.n) return;
74 const Camera cam = frame.LoadCamera();
75 // Geometry first: only Gaussians that reach a tile need their SH row.
76 const float3 mean = LoadMean(params, i);
77 const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
78 Splat s;
79 int count = 0;
80 if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
81 const float min_power =
82 pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;
83 s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
84 s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
85 s.p2 = {0.0f, 0.0f, 0.0f, min_power};
86 ForEachTileRow<kCulling>(
87 s, frame.grid, frame.width, frame.height,
88 [&](int, int tx0, int tx1) { count += tx1 - tx0; });
89 }
42 if (deg > 1) {90 if (count == 0) {
43 float xx = x * x, yy = y * y, zz = z * z;
44 float xy = x * y, yz = y * z, xz = x * z;
45 result = result + SH_C2[0] * xy * sh[4] + SH_C2[1] * yz * sh[5] +
46 SH_C2[2] * (2.0f * zz - xx - yy) * sh[6] +
47 SH_C2[3] * xz * sh[7] + SH_C2[4] * (xx - yy) * sh[8];
91 s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
92 splats[i] = s;
93 tile_counts[i] = 0;
94 return;
95 }
96 const int vis_idx = atomicAdd(num_visible, 1);
97 float inv_len;
98 const float3 rgb =
99 ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
100 s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
101 splats[i] = s;
102 tile_counts[i] = count;
103 depth_keys[vis_idx] = AsUint(pr.depth);
104 ids[vis_idx] = i;
105}
48106
107__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
108 RasterForwardKernel(const int2* __restrict__ ranges,
109 const int32_t* __restrict__ point_list,
110 const Splat* __restrict__ splats, Frame frame,
111 float* __restrict__ image, float* __restrict__ final_t,
112 int32_t* __restrict__ n_contrib) {
113 constexpr int kRasterThreads = RasterLayout<1>::kThreads;
114 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
115 const int tid = threadIdx.x;
116 const int warp = tid / 32;
117 const int lane = tid % 32;
118 const RasterPixels px = RasterPixels::For<1>(tid);
119 const float fx = static_cast<float>(px.x);
120 const float fy = static_cast<float>(px.y0);
121 const int2 range = ranges[tile];
122
123 float t = 1.0f, r = 0.0f, g = 0.0f, b = 0.0f;
124 int last = 0;
125 bool done = px.x >= frame.width || px.y0 >= frame.height;
126
127 __shared__ float4 batch_geom0[kRasterThreads];
128 __shared__ float2 batch_geom1[kRasterThreads];
129 __shared__ float4 batch_rgb_o[kRasterThreads];
130 __shared__ uint8_t batch_mask[kRasterThreads];
131 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
132 for (int start = range.x; start < range.y; start += kRasterThreads) {
133 if (__syncthreads_count(done) == kRasterThreads) break;
134 const int idx = start + tid;
135 if (idx < range.y) {
136 const int32_t id = point_list[idx];
137 const float4 p0 = __ldg(&splats[id].p0);
138 const float4 p1 = __ldg(&splats[id].p1);
139 const float4 p2 = __ldg(&splats[id].p2);
140 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
141 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
142 batch_geom1[tid] = {c2.z, p2.w};
143 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
144 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
145 } else {
146 batch_mask[tid] = 0u;
147 }
148 __syncthreads();
149 const int count = min(kRasterThreads, range.y - start);
150 const int base = start - range.x;
151 bool warp_done = __all_sync(warp::kFullMask, done);
152#pragma unroll
153 for (int c = 0; c < kRasterThreads; c += 32) {
154 if (c < count && !warp_done) {
155 uint32_t mask = __ballot_sync(
156 warp::kFullMask,
157 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
158 while (mask != 0u) {
159 const int bit = __ffs(mask) - 1;
160 mask &= mask - 1u;
161 const int j = c + bit;
49 if (deg > 2) {162 if (!done) {
50 result = result + SH_C3[0] * y * (3.0f * xx - yy) * sh[9] +
51 SH_C3[1] * xy * z * sh[10] +
52 SH_C3[2] * y * (4.0f * zz - xx - yy) * sh[11] +
53 SH_C3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy) * sh[12] +
54 SH_C3[4] * x * (4.0f * zz - xx - yy) * sh[13] +
55 SH_C3[5] * z * (xx - yy) * sh[14] +
163 const float4 g0 = batch_geom0[j];
164 const float2 g1 = batch_geom1[j];
165 const float dx = g0.x - fx;
166 const float dy = g0.y - fy;
167 const float power =
56 SH_C3[6] * x * (xx - 3.0f * yy) * sh[15];168 g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
169 if (power <= 0.0f && power >= g1.y) {
170 const float4 ro = batch_rgb_o[j];
171 const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
172 if (alpha >= kMinAlpha) {
173 const float next_t = t * (1.0f - alpha);
174 if (next_t < kMinTransmittance) {
175 done = true;
176 } else {
177 const float w = alpha * t;
178 r += w * ro.x;
179 g += w * ro.y;
180 b += w * ro.z;
181 t = next_t;
182 last = base + j + 1;
57 }183 }
58 }184 }
59 }185 }
60 result += 0.5f;
61
62 // RGB colors are clamped to positive values. If values are
63 // clamped, we need to keep track of this for the backward pass.
64 clamped[3 * idx + 0] = (result.x < 0);
65 clamped[3 * idx + 1] = (result.y < 0);
66 clamped[3 * idx + 2] = (result.z < 0);
67 return glm::max(result, 0.0f);
68}186 }
69
70// Forward version of 2D covariance matrix computation
71__device__ float4 computeCov2D(const float3& mean, float focal_x, float focal_y,
72 float tan_fovx, float tan_fovy,
73 const float* cov3D, const float* viewmatrix,
74 float kernel_size, bool mip_splatting_enabled) {
75 // The following models the steps outlined by equations 29
76 // and 31 in "EWA Splatting" (Zwicker et al., 2002).
77 // Additionally considers aspect / scaling of viewport.
78 // Transposes used to account for row-/column-major conventions.
79 float3 t = transformPoint4x3(mean, viewmatrix);
80
81 const float limx = 1.3f * tan_fovx;
82 const float limy = 1.3f * tan_fovy;
83 const float txtz = t.x / t.z;
84 const float tytz = t.y / t.z;
85 t.x = min(limx, max(-limx, txtz)) * t.z;
86 t.y = min(limy, max(-limy, tytz)) * t.z;
87
88 glm::mat3 J =
89 glm::mat3(focal_x / t.z, 0.0f, -(focal_x * t.x) / (t.z * t.z), 0.0f,
90 focal_y / t.z, -(focal_y * t.y) / (t.z * t.z), 0, 0, 0);
91
92 glm::mat3 W = glm::mat3(viewmatrix[0], viewmatrix[4], viewmatrix[8],
93 viewmatrix[1], viewmatrix[5], viewmatrix[9],
94 viewmatrix[2], viewmatrix[6], viewmatrix[10]);
95
96 glm::mat3 T = W * J;
97
98 glm::mat3 Vrk = glm::mat3(cov3D[0], cov3D[1], cov3D[2], cov3D[1], cov3D[3],
99 cov3D[4], cov3D[2], cov3D[4], cov3D[5]);
100
101 glm::mat3 cov = glm::transpose(T) * glm::transpose(Vrk) * T;
102
103 float coef;
104 if (mip_splatting_enabled) {
105 // compute the coef of alpha based on the detemintant
106 const float det_0 =
107 max(1e-6, cov[0][0] * cov[1][1] - cov[0][1] * cov[0][1]);
108 const float det_1 =
109 max(1e-6, (cov[0][0] + kernel_size) * (cov[1][1] + kernel_size) -
110 cov[0][1] * cov[0][1]);
111 coef = sqrt(det_0 / (det_1 + 1e-6) + 1e-6);
112
113 if (det_0 <= 1e-6 || det_1 <= 1e-6) {
114 coef = 0.0f;
187 if (__all_sync(warp::kFullMask, done)) {
188 warp_done = true;
189 break;
115 }190 }
116
117 cov[0][0] += kernel_size;
118 cov[1][1] += kernel_size;
119 } else {
120 // Apply low-pass filter: every Gaussian should be at least
121 // one pixel wide/high. Discard 3rd row and column.
122 cov[0][0] += 0.3f;
123 cov[1][1] += 0.3f;
124 coef = 1.0f;
125 }191 }
126 return {float(cov[0][0]), float(cov[0][1]), float(cov[1][1]), coef};
192 warp_done = __all_sync(warp::kFullMask, done);
127}193 }
128
129// Forward method for converting scale and rotation properties of each
130// Gaussian to a 3D covariance matrix in world space. Also takes care
131// of quaternion normalization.
132__device__ void computeCov3D(const glm::vec3 scale, float mod,
133 const glm::vec4 rot, float* cov3D) {
134 // Create scaling matrix
135 glm::mat3 S = glm::mat3(1.0f);
136 S[0][0] = mod * scale.x;
137 S[1][1] = mod * scale.y;
138 S[2][2] = mod * scale.z;
139
140 // Normalize quaternion to get valid rotation
141 glm::vec4 q = rot; // / glm::length(rot);
142 float r = q.x;
143 float x = q.y;
144 float y = q.z;
145 float z = q.w;
146
147 // Compute rotation matrix from quaternion
148 glm::mat3 R = glm::mat3(1.f - 2.f * (y * y + z * z), 2.f * (x * y - r * z),
149 2.f * (x * z + r * y), 2.f * (x * y + r * z),
150 1.f - 2.f * (x * x + z * z), 2.f * (y * z - r * x),
151 2.f * (x * z - r * y), 2.f * (y * z + r * x),
152 1.f - 2.f * (x * x + y * y));
153
154 glm::mat3 M = S * R;
155
156 // Compute 3D world covariance matrix Sigma
157 glm::mat3 Sigma = glm::transpose(M) * M;
158
159 // Covariance is symmetric, only store upper right
160 cov3D[0] = Sigma[0][0];
161 cov3D[1] = Sigma[0][1];
162 cov3D[2] = Sigma[0][2];
163 cov3D[3] = Sigma[1][1];
164 cov3D[4] = Sigma[1][2];
165 cov3D[5] = Sigma[2][2];
166}194 }
167
168// Perform initial steps for each Gaussian prior to rasterization.
169template <int C>
170__global__ void preprocessCUDA(
171 int P, int D, int M, const float* orig_points, const glm::vec3* scales,
172 const float scale_modifier, const glm::vec4* rotations,
173 const float* opacities, const float* shs, bool* clamped,
174 const float* cov3D_precomp, const float* colors_precomp,
175 const float* viewmatrix, const float* projmatrix, const glm::vec3* cam_pos,
176 const int W, int H, const float tan_fovx, float tan_fovy,
177 const float focal_x, float focal_y, const float low_pass_2d_kernel_size,
178 const bool mip_splatting_enabled, int* radii, float2* points_xy_image,
179 float* depths, float* cov3Ds, float* rgb, float4* conic_opacity,
180 const dim3 grid, uint32_t* tiles_touched, bool prefiltered, int2* rects,
181 float3 boxmin, float3 boxmax) {
182 auto idx = cg::this_grid().thread_rank();
183 if (idx >= P) return;
184
185 // Initialize radius and touched tiles to 0. If this isn't changed,
186 // this Gaussian will not be processed further.
187 radii[idx] = 0;
188 tiles_touched[idx] = 0;
189
190 // Perform near culling, quit if outside.
191 float3 p_view;
192 if (!in_frustum(idx, orig_points, viewmatrix, projmatrix, prefiltered,
193 p_view))
194 return;
195
196 // Transform point by projecting
197 float3 p_orig = {orig_points[3 * idx], orig_points[3 * idx + 1],
198 orig_points[3 * idx + 2]};
199
200 if (p_orig.x < boxmin.x || p_orig.y < boxmin.y || p_orig.z < boxmin.z ||
201 p_orig.x > boxmax.x || p_orig.y > boxmax.y || p_orig.z > boxmax.z)
202 return;
203
204 float4 p_hom = transformPoint4x4(p_orig, projmatrix);
205 float p_w = 1.0f / (p_hom.w + 0.0000001f);
206 float3 p_proj = {p_hom.x * p_w, p_hom.y * p_w, p_hom.z * p_w};
207
208 // If 3D covariance matrix is precomputed, use it, otherwise compute
209 // from scaling and rotation parameters.
210 const float* cov3D;
211 if (cov3D_precomp != nullptr) {
212 cov3D = cov3D_precomp + idx * 6;
213 } else {
214 computeCov3D(scales[idx], scale_modifier, rotations[idx], cov3Ds + idx * 6);
215 cov3D = cov3Ds + idx * 6;
216 }195 }
217
218 // Compute 2D screen-space covariance matrix
219 float4 cov =
220 computeCov2D(p_orig, focal_x, focal_y, tan_fovx, tan_fovy, cov3D,
221 viewmatrix, low_pass_2d_kernel_size, mip_splatting_enabled);
222
223 // Invert covariance (EWA algorithm)
224 float det = (cov.x * cov.z - cov.y * cov.y);
225 if (det == 0.0f) return;
226 float det_inv = 1.f / det;
227 float3 conic = {cov.z * det_inv, -cov.y * det_inv, cov.x * det_inv};
228
229 // Compute extent in screen space (by finding eigenvalues of
230 // 2D covariance matrix). Use extent to compute a bounding rectangle
231 // of screen-space tiles that this Gaussian overlaps with. Quit if
232 // rectangle covers 0 tiles.
233
234 float mid = 0.5f * (cov.x + cov.z);
235 float lambda1 = mid + sqrt(max(0.1f, mid * mid - det));
236 float lambda2 = mid - sqrt(max(0.1f, mid * mid - det));
237 float my_radius = ceil(3.f * sqrt(max(lambda1, lambda2)));
238 float2 point_image = {ndc2Pix(p_proj.x, W), ndc2Pix(p_proj.y, H)};
239 uint2 rect_min, rect_max;
240
241 if (rects == nullptr) // More conservative
242 {
243 getRect(point_image, my_radius, rect_min, rect_max, grid);
244 } else // Slightly more aggressive, might need a math cleanup
245 {
196 if (px.x < frame.width && px.y0 < frame.height) {
197 const int pixels = frame.width * frame.height;
246 const int2 my_rect = {(int)ceil(3.f * sqrt(cov.x)),198 const int p = px.y0 * frame.width + px.x;
247 (int)ceil(3.f * sqrt(cov.z))};
248 rects[idx] = my_rect;
249 getRect(point_image, my_rect, rect_min, rect_max, grid);
250 }
251
252 if ((rect_max.x - rect_min.x) * (rect_max.y - rect_min.y) == 0) return;
253
254 // If colors have been precomputed, use them, otherwise convert
255 // spherical harmonics coefficients to RGB color.
256 if (colors_precomp == nullptr) {
257 glm::vec3 result = computeColorFromSH(idx, D, M, (glm::vec3*)orig_points,
258 *cam_pos, shs, clamped);
259 rgb[idx * C + 0] = result.x;199 image[0 * pixels + p] = r;
260 rgb[idx * C + 1] = result.y;200 image[1 * pixels + p] = g;
261 rgb[idx * C + 2] = result.z;201 image[2 * pixels + p] = b;
262 }
263
264 // Store some useful helper data for the next steps.
265 depths[idx] = p_view.z;
266 radii[idx] = my_radius;
267 points_xy_image[idx] = point_image;
268 // Inverse 2D covariance and opacity neatly pack into one float4
269 conic_opacity[idx] = {conic.x, conic.y, conic.z, opacities[idx] * cov.w};
270 tiles_touched[idx] = (rect_max.y - rect_min.y) * (rect_max.x - rect_min.x);
271}
272
273// Main rasterization method. Collaboratively works on one tile per
274// block, each thread treats one pixel. Alternates between fetching
275// and rasterizing data.
276template <uint32_t CHANNELS>
277__global__ void __launch_bounds__(BLOCK_X* BLOCK_Y)
278 renderCUDA(const uint2* __restrict__ ranges,
279 const uint32_t* __restrict__ point_list, int W, int H,
280 const float2* __restrict__ points_xy_image,
281 const float* __restrict__ features,
282 const float4* __restrict__ conic_opacity,
283 float* __restrict__ final_T, uint32_t* __restrict__ n_contrib,
284 const float* __restrict__ bg_color,
285 float* __restrict__ out_color) {
286 // Identify current tile and associated min/max pixel range.
287 auto block = cg::this_thread_block();
288 uint32_t horizontal_blocks = (W + BLOCK_X - 1) / BLOCK_X;
289 uint2 pix_min = {block.group_index().x * BLOCK_X,
290 block.group_index().y * BLOCK_Y};
291 uint2 pix_max = {min(pix_min.x + BLOCK_X, W), min(pix_min.y + BLOCK_Y, H)};
292 uint2 pix = {pix_min.x + block.thread_index().x,
293 pix_min.y + block.thread_index().y};
294 uint32_t pix_id = W * pix.y + pix.x;
295 float2 pixf = {(float)pix.x, (float)pix.y};
296
297 // Check if this thread is associated with a valid pixel or outside.
298 bool inside = pix.x < W && pix.y < H;
299 // Done threads can help with fetching, but don't rasterize
300 bool done = !inside;
301
302 // Load start/end range of IDs to process in bit sorted list.
303 uint2 range =
304 ranges[block.group_index().y * horizontal_blocks + block.group_index().x];
305 const int rounds = ((range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE);
306 int toDo = range.y - range.x;
307
308 // Allocate storage for batches of collectively fetched data.
309 __shared__ int collected_id[BLOCK_SIZE];
310 __shared__ float2 collected_xy[BLOCK_SIZE];
311 __shared__ float4 collected_conic_opacity[BLOCK_SIZE];
312
313 // Initialize helper variables
314 float T = 1.0f;202 final_t[p] = t;
315 uint32_t contributor = 0;203 n_contrib[p] = last;
316 uint32_t last_contributor = 0;
317 float C[CHANNELS] = {0};
318
319 // Iterate over batches until all done or range is complete
320 for (int i = 0; i < rounds; i++, toDo -= BLOCK_SIZE) {
321 // End if entire block votes that it is done rasterizing
322 int num_done = __syncthreads_count(done);
323 if (num_done == BLOCK_SIZE) break;
324
325 // Collectively fetch per-Gaussian data from global to shared
326 int progress = i * BLOCK_SIZE + block.thread_rank();
327 if (range.x + progress < range.y) {
328 int coll_id = point_list[range.x + progress];
329 collected_id[block.thread_rank()] = coll_id;
330 collected_xy[block.thread_rank()] = points_xy_image[coll_id];
331 collected_conic_opacity[block.thread_rank()] = conic_opacity[coll_id];
332 }
333 block.sync();
334
335 // Iterate over current batch
336 for (int j = 0; !done && j < min(BLOCK_SIZE, toDo); j++) {
337 // Keep track of current position in range
338 contributor++;
339
340 // Resample using conic matrix (cf. "Surface
341 // Splatting" by Zwicker et al., 2001)
342 float2 xy = collected_xy[j];
343 float2 d = {xy.x - pixf.x, xy.y - pixf.y};
344 float4 con_o = collected_conic_opacity[j];
345 float power = -0.5f * (con_o.x * d.x * d.x + con_o.z * d.y * d.y) -
346 con_o.y * d.x * d.y;
347 if (power > 0.0f) continue;
348
349 // Eq. (2) from 3D Gaussian splatting paper.
350 // Obtain alpha by multiplying with Gaussian opacity
351 // and its exponential falloff from mean.
352 // Avoid numerical instabilities (see paper appendix).
353 float alpha = min(0.99f, con_o.w * exp(power));
354 if (alpha < 1.0f / 255.0f) continue;
355 float test_T = T * (1 - alpha);
356 if (test_T < 0.0001f) {
357 done = true;
358 continue;
359 }
360
361 // Eq. (3) from 3D Gaussian splatting paper.
362 for (int ch = 0; ch < CHANNELS; ch++)
363 C[ch] += features[collected_id[j] * CHANNELS + ch] * alpha * T;
364
365 T = test_T;
366
367 // Keep track of last range entry to update this
368 // pixel.
369 last_contributor = contributor;
370 }204 }
371 }205}
372206
373 // All threads that treat valid pixel write out their final
374 // rendering data to the frame and auxiliary buffers.
375 if (inside) {
376 final_T[pix_id] = T;
377 n_contrib[pix_id] = last_contributor;
378 for (int ch = 0; ch < CHANNELS; ch++)
379 out_color[ch * H * W + pix_id] = C[ch] + T * bg_color[ch];
380 }
381}
207template __global__ void PreprocessKernel<Culling::kRect, float>(
208 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
209 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
210template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
211 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
212 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
213template __global__ void PreprocessKernel<Culling::kEllipse, float>(
214 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
215 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
216template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
217 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
218 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
382219
383void FORWARD::render(const dim3 grid, dim3 block, const uint2* ranges,
384 const uint32_t* point_list, int W, int H,
385 const float2* means2D, const float* colors,
386 const float4* conic_opacity, float* final_T,
387 uint32_t* n_contrib, const float* bg_color,
388 float* out_color) {
389 renderCUDA<NUM_CHANNELS><<<grid, block>>>(ranges, point_list, W, H, means2D,
390 colors, conic_opacity, final_T,
391 n_contrib, bg_color, out_color);
392}
393
394void FORWARD::preprocess(
395 int P, int D, int M, const float* means3D, const glm::vec3* scales,
396 const float scale_modifier, const glm::vec4* rotations,
397 const float* opacities, const float* shs, bool* clamped,
398 const float* cov3D_precomp, const float* colors_precomp,
399 const float* viewmatrix, const float* projmatrix, const glm::vec3* cam_pos,
400 const int W, int H, const float focal_x, float focal_y,
401 const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
402 const bool mip_splatting_enabled, int* radii, float2* means2D,
403 float* depths, float* cov3Ds, float* rgb, float4* conic_opacity,
404 const dim3 grid, uint32_t* tiles_touched, bool prefiltered, int2* rects,
405 float3 boxmin, float3 boxmax) {
406 preprocessCUDA<NUM_CHANNELS><<<(P + 255) / 256, 256>>>(
407 P, D, M, means3D, scales, scale_modifier, rotations, opacities, shs,
408 clamped, cov3D_precomp, colors_precomp, viewmatrix, projmatrix, cam_pos,
409 W, H, tan_fovx, tan_fovy, focal_x, focal_y, low_pass_2d_kernel_size,
410 mip_splatting_enabled, radii, means2D, depths, cov3Ds, rgb, conic_opacity,
411 grid, tiles_touched, prefiltered, rects, boxmin, boxmax);
412}
220} // namespace gs
↕ Expand 9 unchanged lines (lines 1–9)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #include "forward.h"
1313
14−#include <cooperative_groups.h>
14+#include "auxiliary.h"
15−#include <cooperative_groups/reduce.h>
1615
17−#include "auxiliary.h"
18−namespace cg = cooperative_groups;
16+namespace gs {
1917
20−// Forward method for converting the input spherical harmonics
21−// coefficients of each Gaussian to a simple RGB color.
22−__device__ glm::vec3 computeColorFromSH(int idx, int deg, int max_coeffs,
18+// Screen-space geometry of Gaussian `i`: everything in `Projection` except
19+// `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
20+// tile-rectangle cull).
21+template <Culling kCulling>
22+GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
23+ float3 mean, int i) {
24+ Projection out;
25+ out.radius = 0;
23− const glm::vec3* means,
26+ const float3 t = cam.ToView(mean);
24− glm::vec3 campos, const float* shs,
25− bool* clamped) {
26− // The implementation is loosely based on code for
27− // "Differentiable Point-Based Radiance Fields for
28− // Efficient View Synthesis" by Zhang et al. (2022)
29− glm::vec3 pos = means[idx];
30− glm::vec3 dir = pos - campos;
31− dir = dir / glm::length(dir);
27+ if (t.z <= kNearPlane) return out;
28+ const float opacity = p.opacities[i];
29+ float limit = 0.0f;
30+ if (!AlphaLimit(opacity, limit)) {
31+ if constexpr (kCulling == Culling::kEllipse) return out;
32+ }
33+ const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
34+ p.scales[3 * i + 2]};
35+ const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
36+ ProjectionState st;
37+ ProjectCovarianceFromView(cam, t, scale, quat, st);
38+ const float det = st.a * st.c - st.b * st.b;
39+ if (det == 0.0f) return out;
40+ const float inv = 1.0f / det;
41+ out.conic = {st.c * inv, -st.b * inv, st.a * inv};
42+ out.xy = ProjectMean(cam, st.t);
43+ out.depth = st.t.z;
44+ out.opacity = opacity;
45+ out.limit = limit;
46+ out.radius = BoundingRadius(st.a, st.b, st.c);
47+ return out;
48+}
3249
33− glm::vec3* sh = ((glm::vec3*)shs) + idx * max_coeffs;
34− glm::vec3 result = SH_C0 * sh[0];
50+// Color seen along `dir`, clamped at zero, reading from an SH row.
51+template <typename ShT>
52+GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
53+ float basis[kShCoeffs];
54+ sh::Basis(dir, basis);
55+ const float3 raw = sh::EvalRow(sh_row, basis);
56+ return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
57+}
3558
36− if (deg > 0) {
37− float x = dir.x;
38− float y = dir.y;
39− float z = dir.z;
40− result = result - SH_C1 * y * sh[1] + SH_C1 * z * sh[2] - SH_C1 * x * sh[3];
59+// ---------------------------------------------------------------------------
60+// Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
61+// evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
62+// depth sort via per-thread `atomicAdd`.
4163
64+template <Culling kCulling, typename ShT>
65+__global__ void __launch_bounds__(kGaussianThreads, 8)
66+ PreprocessKernel(GaussianParams params, Frame frame,
67+ Splat* __restrict__ splats,
68+ int32_t* __restrict__ tile_counts,
69+ uint32_t* __restrict__ depth_keys,
70+ int32_t* __restrict__ ids,
71+ int32_t* __restrict__ num_visible) {
72+ const int i = blockIdx.x * blockDim.x + threadIdx.x;
73+ if (i >= params.n) return;
74+ const Camera cam = frame.LoadCamera();
75+ // Geometry first: only Gaussians that reach a tile need their SH row.
76+ const float3 mean = LoadMean(params, i);
77+ const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
78+ Splat s;
79+ int count = 0;
80+ if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
81+ const float min_power =
82+ pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;
83+ s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
84+ s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
85+ s.p2 = {0.0f, 0.0f, 0.0f, min_power};
86+ ForEachTileRow<kCulling>(
87+ s, frame.grid, frame.width, frame.height,
88+ [&](int, int tx0, int tx1) { count += tx1 - tx0; });
89+ }
42− if (deg > 1) {
90+ if (count == 0) {
43− float xx = x * x, yy = y * y, zz = z * z;
44− float xy = x * y, yz = y * z, xz = x * z;
45− result = result + SH_C2[0] * xy * sh[4] + SH_C2[1] * yz * sh[5] +
46− SH_C2[2] * (2.0f * zz - xx - yy) * sh[6] +
47− SH_C2[3] * xz * sh[7] + SH_C2[4] * (xx - yy) * sh[8];
91+ s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
92+ splats[i] = s;
93+ tile_counts[i] = 0;
94+ return;
95+ }
96+ const int vis_idx = atomicAdd(num_visible, 1);
97+ float inv_len;
98+ const float3 rgb =
99+ ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
100+ s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
101+ splats[i] = s;
102+ tile_counts[i] = count;
103+ depth_keys[vis_idx] = AsUint(pr.depth);
104+ ids[vis_idx] = i;
105+}
48106
107+__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
108+ RasterForwardKernel(const int2* __restrict__ ranges,
109+ const int32_t* __restrict__ point_list,
110+ const Splat* __restrict__ splats, Frame frame,
111+ float* __restrict__ image, float* __restrict__ final_t,
112+ int32_t* __restrict__ n_contrib) {
113+ constexpr int kRasterThreads = RasterLayout<1>::kThreads;
114+ const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
115+ const int tid = threadIdx.x;
116+ const int warp = tid / 32;
117+ const int lane = tid % 32;
118+ const RasterPixels px = RasterPixels::For<1>(tid);
119+ const float fx = static_cast<float>(px.x);
120+ const float fy = static_cast<float>(px.y0);
121+ const int2 range = ranges[tile];
122+
123+ float t = 1.0f, r = 0.0f, g = 0.0f, b = 0.0f;
124+ int last = 0;
125+ bool done = px.x >= frame.width || px.y0 >= frame.height;
126+
127+ __shared__ float4 batch_geom0[kRasterThreads];
128+ __shared__ float2 batch_geom1[kRasterThreads];
129+ __shared__ float4 batch_rgb_o[kRasterThreads];
130+ __shared__ uint8_t batch_mask[kRasterThreads];
131+ const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
132+ for (int start = range.x; start < range.y; start += kRasterThreads) {
133+ if (__syncthreads_count(done) == kRasterThreads) break;
134+ const int idx = start + tid;
135+ if (idx < range.y) {
136+ const int32_t id = point_list[idx];
137+ const float4 p0 = __ldg(&splats[id].p0);
138+ const float4 p1 = __ldg(&splats[id].p1);
139+ const float4 p2 = __ldg(&splats[id].p2);
140+ const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
141+ batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
142+ batch_geom1[tid] = {c2.z, p2.w};
143+ batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
144+ batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
145+ } else {
146+ batch_mask[tid] = 0u;
147+ }
148+ __syncthreads();
149+ const int count = min(kRasterThreads, range.y - start);
150+ const int base = start - range.x;
151+ bool warp_done = __all_sync(warp::kFullMask, done);
152+#pragma unroll
153+ for (int c = 0; c < kRasterThreads; c += 32) {
154+ if (c < count && !warp_done) {
155+ uint32_t mask = __ballot_sync(
156+ warp::kFullMask,
157+ c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
158+ while (mask != 0u) {
159+ const int bit = __ffs(mask) - 1;
160+ mask &= mask - 1u;
161+ const int j = c + bit;
49− if (deg > 2) {
162+ if (!done) {
50− result = result + SH_C3[0] * y * (3.0f * xx - yy) * sh[9] +
51− SH_C3[1] * xy * z * sh[10] +
52− SH_C3[2] * y * (4.0f * zz - xx - yy) * sh[11] +
53− SH_C3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy) * sh[12] +
54− SH_C3[4] * x * (4.0f * zz - xx - yy) * sh[13] +
55− SH_C3[5] * z * (xx - yy) * sh[14] +
163+ const float4 g0 = batch_geom0[j];
164+ const float2 g1 = batch_geom1[j];
165+ const float dx = g0.x - fx;
166+ const float dy = g0.y - fy;
167+ const float power =
56− SH_C3[6] * x * (xx - 3.0f * yy) * sh[15];
168+ g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
169+ if (power <= 0.0f && power >= g1.y) {
170+ const float4 ro = batch_rgb_o[j];
171+ const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
172+ if (alpha >= kMinAlpha) {
173+ const float next_t = t * (1.0f - alpha);
174+ if (next_t < kMinTransmittance) {
175+ done = true;
176+ } else {
177+ const float w = alpha * t;
178+ r += w * ro.x;
179+ g += w * ro.y;
180+ b += w * ro.z;
181+ t = next_t;
182+ last = base + j + 1;
57183⇥ }
58184⇥ }
59185⇥ }
60− result += 0.5f;
61−
62− // RGB colors are clamped to positive values. If values are
63− // clamped, we need to keep track of this for the backward pass.
64− clamped[3 * idx + 0] = (result.x < 0);
65− clamped[3 * idx + 1] = (result.y < 0);
66− clamped[3 * idx + 2] = (result.z < 0);
67− return glm::max(result, 0.0f);
68186⇥ }
69−
70−// Forward version of 2D covariance matrix computation
71−__device__ float4 computeCov2D(const float3& mean, float focal_x, float focal_y,
72− float tan_fovx, float tan_fovy,
73− const float* cov3D, const float* viewmatrix,
74− float kernel_size, bool mip_splatting_enabled) {
75− // The following models the steps outlined by equations 29
76− // and 31 in "EWA Splatting" (Zwicker et al., 2002).
77− // Additionally considers aspect / scaling of viewport.
78− // Transposes used to account for row-/column-major conventions.
79− float3 t = transformPoint4x3(mean, viewmatrix);
80−
81− const float limx = 1.3f * tan_fovx;
82− const float limy = 1.3f * tan_fovy;
83− const float txtz = t.x / t.z;
84− const float tytz = t.y / t.z;
85− t.x = min(limx, max(-limx, txtz)) * t.z;
86− t.y = min(limy, max(-limy, tytz)) * t.z;
87−
88− glm::mat3 J =
89− glm::mat3(focal_x / t.z, 0.0f, -(focal_x * t.x) / (t.z * t.z), 0.0f,
90− focal_y / t.z, -(focal_y * t.y) / (t.z * t.z), 0, 0, 0);
91−
92− glm::mat3 W = glm::mat3(viewmatrix[0], viewmatrix[4], viewmatrix[8],
93− viewmatrix[1], viewmatrix[5], viewmatrix[9],
94− viewmatrix[2], viewmatrix[6], viewmatrix[10]);
95−
96− glm::mat3 T = W * J;
97−
98− glm::mat3 Vrk = glm::mat3(cov3D[0], cov3D[1], cov3D[2], cov3D[1], cov3D[3],
99− cov3D[4], cov3D[2], cov3D[4], cov3D[5]);
100−
101− glm::mat3 cov = glm::transpose(T) * glm::transpose(Vrk) * T;
102−
103− float coef;
104− if (mip_splatting_enabled) {
105− // compute the coef of alpha based on the detemintant
106− const float det_0 =
107− max(1e-6, cov[0][0] * cov[1][1] - cov[0][1] * cov[0][1]);
108− const float det_1 =
109− max(1e-6, (cov[0][0] + kernel_size) * (cov[1][1] + kernel_size) -
110− cov[0][1] * cov[0][1]);
111− coef = sqrt(det_0 / (det_1 + 1e-6) + 1e-6);
112−
113− if (det_0 <= 1e-6 || det_1 <= 1e-6) {
114− coef = 0.0f;
187+ if (__all_sync(warp::kFullMask, done)) {
188+ warp_done = true;
189+ break;
115190⇥ }
116−
117− cov[0][0] += kernel_size;
118− cov[1][1] += kernel_size;
119− } else {
120− // Apply low-pass filter: every Gaussian should be at least
121− // one pixel wide/high. Discard 3rd row and column.
122− cov[0][0] += 0.3f;
123− cov[1][1] += 0.3f;
124− coef = 1.0f;
125191⇥ }
126− return {float(cov[0][0]), float(cov[0][1]), float(cov[1][1]), coef};
192+ warp_done = __all_sync(warp::kFullMask, done);
127193⇥ }
128−
129−// Forward method for converting scale and rotation properties of each
130−// Gaussian to a 3D covariance matrix in world space. Also takes care
131−// of quaternion normalization.
132−__device__ void computeCov3D(const glm::vec3 scale, float mod,
133− const glm::vec4 rot, float* cov3D) {
134− // Create scaling matrix
135− glm::mat3 S = glm::mat3(1.0f);
136− S[0][0] = mod * scale.x;
137− S[1][1] = mod * scale.y;
138− S[2][2] = mod * scale.z;
139−
140− // Normalize quaternion to get valid rotation
141− glm::vec4 q = rot; // / glm::length(rot);
142− float r = q.x;
143− float x = q.y;
144− float y = q.z;
145− float z = q.w;
146−
147− // Compute rotation matrix from quaternion
148− glm::mat3 R = glm::mat3(1.f - 2.f * (y * y + z * z), 2.f * (x * y - r * z),
149− 2.f * (x * z + r * y), 2.f * (x * y + r * z),
150− 1.f - 2.f * (x * x + z * z), 2.f * (y * z - r * x),
151− 2.f * (x * z - r * y), 2.f * (y * z + r * x),
152− 1.f - 2.f * (x * x + y * y));
153−
154− glm::mat3 M = S * R;
155−
156− // Compute 3D world covariance matrix Sigma
157− glm::mat3 Sigma = glm::transpose(M) * M;
158−
159− // Covariance is symmetric, only store upper right
160− cov3D[0] = Sigma[0][0];
161− cov3D[1] = Sigma[0][1];
162− cov3D[2] = Sigma[0][2];
163− cov3D[3] = Sigma[1][1];
164− cov3D[4] = Sigma[1][2];
165− cov3D[5] = Sigma[2][2];
166194⇥ }
167−
168−// Perform initial steps for each Gaussian prior to rasterization.
169−template <int C>
170−__global__ void preprocessCUDA(
171− int P, int D, int M, const float* orig_points, const glm::vec3* scales,
172− const float scale_modifier, const glm::vec4* rotations,
173− const float* opacities, const float* shs, bool* clamped,
174− const float* cov3D_precomp, const float* colors_precomp,
175− const float* viewmatrix, const float* projmatrix, const glm::vec3* cam_pos,
176− const int W, int H, const float tan_fovx, float tan_fovy,
177− const float focal_x, float focal_y, const float low_pass_2d_kernel_size,
178− const bool mip_splatting_enabled, int* radii, float2* points_xy_image,
179− float* depths, float* cov3Ds, float* rgb, float4* conic_opacity,
180− const dim3 grid, uint32_t* tiles_touched, bool prefiltered, int2* rects,
181− float3 boxmin, float3 boxmax) {
182− auto idx = cg::this_grid().thread_rank();
183− if (idx >= P) return;
184−
185− // Initialize radius and touched tiles to 0. If this isn't changed,
186− // this Gaussian will not be processed further.
187− radii[idx] = 0;
188− tiles_touched[idx] = 0;
189−
190− // Perform near culling, quit if outside.
191− float3 p_view;
192− if (!in_frustum(idx, orig_points, viewmatrix, projmatrix, prefiltered,
193− p_view))
194− return;
195−
196− // Transform point by projecting
197− float3 p_orig = {orig_points[3 * idx], orig_points[3 * idx + 1],
198− orig_points[3 * idx + 2]};
199−
200− if (p_orig.x < boxmin.x || p_orig.y < boxmin.y || p_orig.z < boxmin.z ||
201− p_orig.x > boxmax.x || p_orig.y > boxmax.y || p_orig.z > boxmax.z)
202− return;
203−
204− float4 p_hom = transformPoint4x4(p_orig, projmatrix);
205− float p_w = 1.0f / (p_hom.w + 0.0000001f);
206− float3 p_proj = {p_hom.x * p_w, p_hom.y * p_w, p_hom.z * p_w};
207−
208− // If 3D covariance matrix is precomputed, use it, otherwise compute
209− // from scaling and rotation parameters.
210− const float* cov3D;
211− if (cov3D_precomp != nullptr) {
212− cov3D = cov3D_precomp + idx * 6;
213− } else {
214− computeCov3D(scales[idx], scale_modifier, rotations[idx], cov3Ds + idx * 6);
215− cov3D = cov3Ds + idx * 6;
216195 }
217−
218− // Compute 2D screen-space covariance matrix
219− float4 cov =
220− computeCov2D(p_orig, focal_x, focal_y, tan_fovx, tan_fovy, cov3D,
221− viewmatrix, low_pass_2d_kernel_size, mip_splatting_enabled);
222−
223− // Invert covariance (EWA algorithm)
224− float det = (cov.x * cov.z - cov.y * cov.y);
225− if (det == 0.0f) return;
226− float det_inv = 1.f / det;
227− float3 conic = {cov.z * det_inv, -cov.y * det_inv, cov.x * det_inv};
228−
229− // Compute extent in screen space (by finding eigenvalues of
230− // 2D covariance matrix). Use extent to compute a bounding rectangle
231− // of screen-space tiles that this Gaussian overlaps with. Quit if
232− // rectangle covers 0 tiles.
233−
234− float mid = 0.5f * (cov.x + cov.z);
235− float lambda1 = mid + sqrt(max(0.1f, mid * mid - det));
236− float lambda2 = mid - sqrt(max(0.1f, mid * mid - det));
237− float my_radius = ceil(3.f * sqrt(max(lambda1, lambda2)));
238− float2 point_image = {ndc2Pix(p_proj.x, W), ndc2Pix(p_proj.y, H)};
239− uint2 rect_min, rect_max;
240−
241− if (rects == nullptr) // More conservative
242− {
243− getRect(point_image, my_radius, rect_min, rect_max, grid);
244− } else // Slightly more aggressive, might need a math cleanup
245− {
196+ if (px.x < frame.width && px.y0 < frame.height) {
197+ const int pixels = frame.width * frame.height;
246− const int2 my_rect = {(int)ceil(3.f * sqrt(cov.x)),
198+ const int p = px.y0 * frame.width + px.x;
247− (int)ceil(3.f * sqrt(cov.z))};
248− rects[idx] = my_rect;
249− getRect(point_image, my_rect, rect_min, rect_max, grid);
250− }
251−
252− if ((rect_max.x - rect_min.x) * (rect_max.y - rect_min.y) == 0) return;
253−
254− // If colors have been precomputed, use them, otherwise convert
255− // spherical harmonics coefficients to RGB color.
256− if (colors_precomp == nullptr) {
257− glm::vec3 result = computeColorFromSH(idx, D, M, (glm::vec3*)orig_points,
258− *cam_pos, shs, clamped);
259− rgb[idx * C + 0] = result.x;
199+ image[0 * pixels + p] = r;
260− rgb[idx * C + 1] = result.y;
200+ image[1 * pixels + p] = g;
261− rgb[idx * C + 2] = result.z;
201+ image[2 * pixels + p] = b;
262− }
263−
264− // Store some useful helper data for the next steps.
265− depths[idx] = p_view.z;
266− radii[idx] = my_radius;
267− points_xy_image[idx] = point_image;
268− // Inverse 2D covariance and opacity neatly pack into one float4
269− conic_opacity[idx] = {conic.x, conic.y, conic.z, opacities[idx] * cov.w};
270− tiles_touched[idx] = (rect_max.y - rect_min.y) * (rect_max.x - rect_min.x);
271−}
272−
273−// Main rasterization method. Collaboratively works on one tile per
274−// block, each thread treats one pixel. Alternates between fetching
275−// and rasterizing data.
276−template <uint32_t CHANNELS>
277−__global__ void __launch_bounds__(BLOCK_X* BLOCK_Y)
278− renderCUDA(const uint2* __restrict__ ranges,
279− const uint32_t* __restrict__ point_list, int W, int H,
280− const float2* __restrict__ points_xy_image,
281− const float* __restrict__ features,
282− const float4* __restrict__ conic_opacity,
283− float* __restrict__ final_T, uint32_t* __restrict__ n_contrib,
284− const float* __restrict__ bg_color,
285− float* __restrict__ out_color) {
286− // Identify current tile and associated min/max pixel range.
287− auto block = cg::this_thread_block();
288− uint32_t horizontal_blocks = (W + BLOCK_X - 1) / BLOCK_X;
289− uint2 pix_min = {block.group_index().x * BLOCK_X,
290− block.group_index().y * BLOCK_Y};
291− uint2 pix_max = {min(pix_min.x + BLOCK_X, W), min(pix_min.y + BLOCK_Y, H)};
292− uint2 pix = {pix_min.x + block.thread_index().x,
293− pix_min.y + block.thread_index().y};
294− uint32_t pix_id = W * pix.y + pix.x;
295− float2 pixf = {(float)pix.x, (float)pix.y};
296−
297− // Check if this thread is associated with a valid pixel or outside.
298− bool inside = pix.x < W && pix.y < H;
299− // Done threads can help with fetching, but don't rasterize
300− bool done = !inside;
301−
302− // Load start/end range of IDs to process in bit sorted list.
303− uint2 range =
304− ranges[block.group_index().y * horizontal_blocks + block.group_index().x];
305− const int rounds = ((range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE);
306− int toDo = range.y - range.x;
307−
308− // Allocate storage for batches of collectively fetched data.
309− __shared__ int collected_id[BLOCK_SIZE];
310− __shared__ float2 collected_xy[BLOCK_SIZE];
311− __shared__ float4 collected_conic_opacity[BLOCK_SIZE];
312−
313− // Initialize helper variables
314− float T = 1.0f;
202+ final_t[p] = t;
315− uint32_t contributor = 0;
203+ n_contrib[p] = last;
316− uint32_t last_contributor = 0;
317− float C[CHANNELS] = {0};
318−
319− // Iterate over batches until all done or range is complete
320− for (int i = 0; i < rounds; i++, toDo -= BLOCK_SIZE) {
321− // End if entire block votes that it is done rasterizing
322− int num_done = __syncthreads_count(done);
323− if (num_done == BLOCK_SIZE) break;
324−
325− // Collectively fetch per-Gaussian data from global to shared
326− int progress = i * BLOCK_SIZE + block.thread_rank();
327− if (range.x + progress < range.y) {
328− int coll_id = point_list[range.x + progress];
329− collected_id[block.thread_rank()] = coll_id;
330− collected_xy[block.thread_rank()] = points_xy_image[coll_id];
331− collected_conic_opacity[block.thread_rank()] = conic_opacity[coll_id];
332− }
333− block.sync();
334−
335− // Iterate over current batch
336− for (int j = 0; !done && j < min(BLOCK_SIZE, toDo); j++) {
337− // Keep track of current position in range
338− contributor++;
339−
340− // Resample using conic matrix (cf. "Surface
341− // Splatting" by Zwicker et al., 2001)
342− float2 xy = collected_xy[j];
343− float2 d = {xy.x - pixf.x, xy.y - pixf.y};
344− float4 con_o = collected_conic_opacity[j];
345− float power = -0.5f * (con_o.x * d.x * d.x + con_o.z * d.y * d.y) -
346− con_o.y * d.x * d.y;
347− if (power > 0.0f) continue;
348−
349− // Eq. (2) from 3D Gaussian splatting paper.
350− // Obtain alpha by multiplying with Gaussian opacity
351− // and its exponential falloff from mean.
352− // Avoid numerical instabilities (see paper appendix).
353− float alpha = min(0.99f, con_o.w * exp(power));
354− if (alpha < 1.0f / 255.0f) continue;
355− float test_T = T * (1 - alpha);
356− if (test_T < 0.0001f) {
357− done = true;
358− continue;
359− }
360−
361− // Eq. (3) from 3D Gaussian splatting paper.
362− for (int ch = 0; ch < CHANNELS; ch++)
363− C[ch] += features[collected_id[j] * CHANNELS + ch] * alpha * T;
364−
365− T = test_T;
366−
367− // Keep track of last range entry to update this
368− // pixel.
369− last_contributor = contributor;
370204⇥ }
371205⇥}
372206
373− // All threads that treat valid pixel write out their final
374− // rendering data to the frame and auxiliary buffers.
375− if (inside) {
376− final_T[pix_id] = T;
377− n_contrib[pix_id] = last_contributor;
378− for (int ch = 0; ch < CHANNELS; ch++)
379− out_color[ch * H * W + pix_id] = C[ch] + T * bg_color[ch];
380− }
381−}
207+template __global__ void PreprocessKernel<Culling::kRect, float>(
208+ GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
209+ uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
210+template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
211+ GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
212+ uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
213+template __global__ void PreprocessKernel<Culling::kEllipse, float>(
214+ GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
215+ uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
216+template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
217+ GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
218+ uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
382219
383−void FORWARD::render(const dim3 grid, dim3 block, const uint2* ranges,
384− const uint32_t* point_list, int W, int H,
385− const float2* means2D, const float* colors,
386− const float4* conic_opacity, float* final_T,
387− uint32_t* n_contrib, const float* bg_color,
388− float* out_color) {
389− renderCUDA<NUM_CHANNELS><<<grid, block>>>(ranges, point_list, W, H, means2D,
390− colors, conic_opacity, final_T,
391− n_contrib, bg_color, out_color);
392−}
393−
394−void FORWARD::preprocess(
395− int P, int D, int M, const float* means3D, const glm::vec3* scales,
396− const float scale_modifier, const glm::vec4* rotations,
397− const float* opacities, const float* shs, bool* clamped,
398− const float* cov3D_precomp, const float* colors_precomp,
399− const float* viewmatrix, const float* projmatrix, const glm::vec3* cam_pos,
400− const int W, int H, const float focal_x, float focal_y,
401− const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
402− const bool mip_splatting_enabled, int* radii, float2* means2D,
403− float* depths, float* cov3Ds, float* rgb, float4* conic_opacity,
404− const dim3 grid, uint32_t* tiles_touched, bool prefiltered, int2* rects,
405− float3 boxmin, float3 boxmax) {
406− preprocessCUDA<NUM_CHANNELS><<<(P + 255) / 256, 256>>>(
407− P, D, M, means3D, scales, scale_modifier, rotations, opacities, shs,
408− clamped, cov3D_precomp, colors_precomp, viewmatrix, projmatrix, cam_pos,
409− W, H, tan_fovx, tan_fovy, focal_x, focal_y, low_pass_2d_kernel_size,
410− mip_splatting_enabled, radii, means2D, depths, cov3Ds, rgb, conic_opacity,
411− grid, tiles_touched, prefiltered, rects, boxmin, boxmax);
412−}
220+} // namespace gs
csrc/backward.h +25 −26
s00_baseline/csrc/backward.h
s01_literature/csrc/backward.h
↕ Expand 10 unchanged lines (lines 1–10)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#ifndef CUDA_RASTERIZER_BACKWARD_H_INCLUDED12#ifndef CUDA_RASTERIZER_BACKWARD_H_INCLUDED
13#define CUDA_RASTERIZER_BACKWARD_H_INCLUDED13#define CUDA_RASTERIZER_BACKWARD_H_INCLUDED
1414
15#include <cuda.h>15#include "auxiliary.h"
1616
17#include "cuda_runtime.h"
18#include "device_launch_parameters.h"
19#define GLM_FORCE_CUDA
20#include <glm/glm.hpp>
17namespace gs {
2118
22namespace BACKWARD {
23void render(const dim3 grid, dim3 block, const uint2* ranges,
19__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
20 RasterBackwardKernel(const int2* __restrict__ ranges,
24 const uint32_t* point_list, int W, int H, const float* bg_color,21 const int32_t* __restrict__ point_list,
25 const float2* means2D, const float4* conic_opacity,
22 const Splat* __restrict__ splats, Frame frame,
26 const float* colors, const float* final_Ts,23 const float* __restrict__ final_t,
27 const uint32_t* n_contrib, const float* dL_dpixels,24 const int32_t* __restrict__ n_contrib,
28 float3* dL_dmean2D, float4* dL_dconic2D, float* dL_dopacity,
25 const float* __restrict__ dl_dimage,
29 float* dL_dcolors);26 float* __restrict__ splat_grads);
3027
31void preprocess(int P, int D, int M, const float3* means, const int* radii,
32 const float* shs, const bool* clamped, const glm::vec3* scales,
33 const glm::vec4* rotations, const float scale_modifier,
34 const float* cov3Ds, const float4* conic_opacity,
28template <typename ShT>
29__global__ void __launch_bounds__(kGaussianThreads)
30 PreprocessBackwardKernel(GaussianParams params, Frame frame,
31 const Splat* __restrict__ splats,
35 const float* view, const float* proj, const float focal_x,32 const float* __restrict__ splat_grads,
33 GaussianGrads grads);
34
35template <typename ShT>
36cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
36 float focal_y, const float tan_fovx, float tan_fovy,37 const Frame& frame, const Splat* splats,
37 const float low_pass_2d_kernel_size,38 const float* splat_grads,
38 const bool mip_splatting_enabled, const glm::vec3* campos,
39 const float3* dL_dmean2D, const float* dL_dconics,
40 glm::vec3* dL_dmeans, float* dL_dcolor, float* dL_dcov3D,
41 float* dL_dsh, glm::vec3* dL_dscale, glm::vec4* dL_drot,
42 float* dL_dopacity);
39 const GaussianGrads& grads,
40 cudaStream_t stream);
41
43} // namespace BACKWARD42} // namespace gs
4443
45#endif44#endif
↕ Expand 10 unchanged lines (lines 1–10)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #ifndef CUDA_RASTERIZER_BACKWARD_H_INCLUDED
1313 #define CUDA_RASTERIZER_BACKWARD_H_INCLUDED
1414
15−#include <cuda.h>
15+#include "auxiliary.h"
1616
17−#include "cuda_runtime.h"
18−#include "device_launch_parameters.h"
19−#define GLM_FORCE_CUDA
20−#include <glm/glm.hpp>
17+namespace gs {
2118
22−namespace BACKWARD {
23−void render(const dim3 grid, dim3 block, const uint2* ranges,
19+__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
20+ RasterBackwardKernel(const int2* __restrict__ ranges,
24− const uint32_t* point_list, int W, int H, const float* bg_color,
21+ const int32_t* __restrict__ point_list,
25− const float2* means2D, const float4* conic_opacity,
22+ const Splat* __restrict__ splats, Frame frame,
26− const float* colors, const float* final_Ts,
23+ const float* __restrict__ final_t,
27− const uint32_t* n_contrib, const float* dL_dpixels,
24+ const int32_t* __restrict__ n_contrib,
28− float3* dL_dmean2D, float4* dL_dconic2D, float* dL_dopacity,
25+ const float* __restrict__ dl_dimage,
29− float* dL_dcolors);
26+ float* __restrict__ splat_grads);
3027
31−void preprocess(int P, int D, int M, const float3* means, const int* radii,
32− const float* shs, const bool* clamped, const glm::vec3* scales,
33− const glm::vec4* rotations, const float scale_modifier,
34− const float* cov3Ds, const float4* conic_opacity,
28+template <typename ShT>
29+__global__ void __launch_bounds__(kGaussianThreads)
30+ PreprocessBackwardKernel(GaussianParams params, Frame frame,
31+ const Splat* __restrict__ splats,
35− const float* view, const float* proj, const float focal_x,
32+ const float* __restrict__ splat_grads,
33+ GaussianGrads grads);
34+
35+template <typename ShT>
36+cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
36− float focal_y, const float tan_fovx, float tan_fovy,
37+ const Frame& frame, const Splat* splats,
37− const float low_pass_2d_kernel_size,
38+ const float* splat_grads,
38− const bool mip_splatting_enabled, const glm::vec3* campos,
39− const float3* dL_dmean2D, const float* dL_dconics,
40− glm::vec3* dL_dmeans, float* dL_dcolor, float* dL_dcov3D,
41− float* dL_dsh, glm::vec3* dL_dscale, glm::vec4* dL_drot,
42− float* dL_dopacity);
39+ const GaussianGrads& grads,
40+ cudaStream_t stream);
41+
43−} // namespace BACKWARD
42+} // namespace gs
4443
4544 #endif
csrc/backward.cu.cc +362 −624
s00_baseline/csrc/backward.cu.cc
s01_literature/csrc/backward.cu.cc
↕ Expand 9 unchanged lines (lines 1–9)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#include "backward.h"12#include "backward.h"
1313
14#include <cooperative_groups.h>14#include "auxiliary.h"
15#include <cooperative_groups/reduce.h>
1615
17#include "auxiliary.h"
18namespace cg = cooperative_groups;16namespace gs {
1917
20// Backward pass for conversion of spherical harmonics to RGB for
21// each Gaussian.
22__device__ void computeColorFromSH(int idx, int deg, int max_coeffs,
23 const glm::vec3* means, glm::vec3 campos,
18namespace sh {
19// Evaluates SH color VJP loading the full 16-basis row `sh_in` and storing
20// `dsh` to `sh_out`.
21template <typename ShT>
22GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,
23 ShT* __restrict__ sh_out) {
24 float b[kShCoeffs];
25 Basis(d, b);
26 float3 coeffs[kShCoeffs];
27 ShIo<ShT>::LoadRow(sh_in, coeffs);
28 float3 raw = {0.5f, 0.5f, 0.5f};
29#pragma unroll
30 for (int k = 0; k < kShCoeffs; ++k) {
31 raw += b[k] * coeffs[k];
32 }
33 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
34 raw.y < 0.0f ? 0.0f : g_rgb.y,
35 raw.z < 0.0f ? 0.0f : g_rgb.z};
36 float3 bgrad[kShCoeffs];
37 BasisGrad(d, bgrad);
38 float3 dsh[kShCoeffs];
39 float3 ddir = {0.0f, 0.0f, 0.0f};
40#pragma unroll
41 for (int k = 0; k < kShCoeffs; ++k) {
42 dsh[k] = b[k] * drgb;
24 const float* shs, const bool* clamped,43 const float s = dot(coeffs[k], drgb);
25 const glm::vec3* dL_dcolor,
26 glm::vec3* dL_dmeans, glm::vec3* dL_dshs) {
27 // Compute intermediate values, as it is done during forward
28 glm::vec3 pos = means[idx];
29 glm::vec3 dir_orig = pos - campos;
30 glm::vec3 dir = dir_orig / glm::length(dir_orig);
44 ddir += s * bgrad[k];
45 }
46 ShIo<ShT>::StoreRow(sh_out, dsh);
47 return ddir;
48}
49} // namespace sh
3150
32 glm::vec3* sh = ((glm::vec3*)shs) + idx * max_coeffs;
51// VJP of the projection for a touched Gaussian `i`; reads its SH row from
52// `sh_in` and writes dL/dsh directly to `sh_out`.
53// Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
54template <typename ShT>
55GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
56 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
57 const ProjectionGrad& g, float3& dmean_out,
58 float3& dscale_out, float4& dquat_out) {
59 const float3 mean = LoadMean(p, i);
60 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
61 p.scales[3 * i + 2]};
62 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
3363
34 // Use PyTorch rule for clamping: if clamping was applied,
35 // gradient becomes 0.
36 glm::vec3 dL_dRGB = dL_dcolor[idx];
37 dL_dRGB.x *= clamped[3 * idx + 0] ? 0 : 1;
38 dL_dRGB.y *= clamped[3 * idx + 1] ? 0 : 1;
39 dL_dRGB.z *= clamped[3 * idx + 2] ? 0 : 1;
64 // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
65 float3 dmean = {0.0f, 0.0f, 0.0f};
66 {
67 float inv_len;
68 const float3 dir = ViewDir(cam, mean, inv_len);
69 const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
70 // d(v/|v|)/dv = (I - dir dir^T) / |v|.
71 dmean += inv_len * (ddir - dot(dir, ddir) * dir);
72 }
4073
41 glm::vec3 dRGBdx(0, 0, 0);
42 glm::vec3 dRGBdy(0, 0, 0);
43 glm::vec3 dRGBdz(0, 0, 0);
44 float x = dir.x;
74 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
45 float y = dir.y;75 const float3 t = cam.ToView(mean);
46 float z = dir.z;76 const float tz = t.z;
77 const float limx = kFovClamp * cam.TanHalfFovX();
78 const float limy = kFovClamp * cam.TanHalfFovY();
79 const float txtz = t.x / tz, tytz = t.y / tz;
80 const bool x_clamped = txtz < -limx || txtz > limx;
81 const bool y_clamped = tytz < -limy || tytz > limy;
82 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
83 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
84 const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
85 const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
86 const Mat3& w = cam.rotation;
87 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
88 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
4789
48 // Target location for this Gaussian to write SH gradients to
49 glm::vec3* dL_dsh = dL_dshs + idx * max_coeffs;
90 const Mat3 r = QuatToRotation(quat);
91 const Sym3 sigma = Covariance3D(r, scale);
92 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
93 const float cov_a = dot(m0, s0) + kLowPass;
94 const float cov_b = dot(m0, s1);
95 const float cov_c = dot(m1, s1) + kLowPass;
5096
51 // No tricks here, just high school-level calculus.
52 float dRGBdsh0 = SH_C0;
53 dL_dsh[0] = dRGBdsh0 * dL_dRGB;
54 if (deg > 0) {
55 float dRGBdsh1 = -SH_C1 * y;
97 // --- Conic -> 2D covariance: G_cov = -K G_K K.
98 const float det = cov_a * cov_c - cov_b * cov_b;
56 float dRGBdsh2 = SH_C1 * z;99 const float inv = 1.0f / det;
57 float dRGBdsh3 = -SH_C1 * x;
58 dL_dsh[1] = dRGBdsh1 * dL_dRGB;
59 dL_dsh[2] = dRGBdsh2 * dL_dRGB;
60 dL_dsh[3] = dRGBdsh3 * dL_dRGB;
100 const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
101 const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
102 // P = G_K K.
103 const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
104 const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
105 // G_cov = -K P (symmetric).
106 const float ca = -(ka * p00 + kb * p10);
107 const float cb = -(ka * p01 + kb * p11);
108 const float cc = -(kb * p01 + kc * p11);
61109
62 dRGBdx = -SH_C1 * sh[3];
63 dRGBdy = -SH_C1 * sh[1];
64 dRGBdz = SH_C1 * sh[2];
110 // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
111 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
112 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
113 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
114 Sym3 dsig;
115 dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
116 dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
117 dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
118 dsig.xy =
119 ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
120 dsig.xz =
121 ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
122 dsig.yz =
123 ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
65124
66 if (deg > 1) {
67 float xx = x * x, yy = y * y, zz = z * z;
125 // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
126 const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
127 const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
68 float xy = x * y, yz = y * z, xz = x * z;128 const float tz2 = tz * tz, tz3 = tz2 * tz;
129 float3 dt;
130 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
131 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
132 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
133 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
134 // --- Screen-space mean u = fx x / z + cx.
135 dt.x += cam.fx / tz * g.xy.x;
136 dt.y += cam.fy / tz * g.xy.y;
137 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
138 dmean += MulT(w, dt);
139 dmean_out = dmean;
69140
70 float dRGBdsh4 = SH_C2[0] * xy;
71 float dRGBdsh5 = SH_C2[1] * yz;
72 float dRGBdsh6 = SH_C2[2] * (2.f * zz - xx - yy);
141 // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
142 const float sv[3] = {scale.x, scale.y, scale.z};
73 float dRGBdsh7 = SH_C2[3] * xz;143 float ds[3];
74 float dRGBdsh8 = SH_C2[4] * (xx - yy);
75 dL_dsh[4] = dRGBdsh4 * dL_dRGB;
76 dL_dsh[5] = dRGBdsh5 * dL_dRGB;
77 dL_dsh[6] = dRGBdsh6 * dL_dRGB;
78 dL_dsh[7] = dRGBdsh7 * dL_dRGB;
79 dL_dsh[8] = dRGBdsh8 * dL_dRGB;
144 Mat3 dr;
145#pragma unroll
146 for (int j = 0; j < 3; ++j) {
147 const float3 r_col = r.col(j);
148 const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
149 ds[j] = dot(da_col, r_col);
150 dr.m[0][j] = da_col.x * sv[j];
151 dr.m[1][j] = da_col.y * sv[j];
152 dr.m[2][j] = da_col.z * sv[j];
153 }
154 dscale_out = {ds[0], ds[1], ds[2]};
80155
156 // --- R(q) -> q.
157 const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
158 const auto& G = dr.m;
159 dquat_out = {
160 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
161 qy * G[2][0] + qx * G[2][1]),
162 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
81 dRGBdx += SH_C2[0] * y * sh[4] + SH_C2[2] * 2.f * -x * sh[6] +163 qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
164 2.0f *
165 (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
166 qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
167 2.0f *
168 (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
82 SH_C2[3] * z * sh[7] + SH_C2[4] * 2.f * x * sh[8];169 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
83 dRGBdy += SH_C2[0] * x * sh[4] + SH_C2[1] * z * sh[5] +
84 SH_C2[2] * 2.f * -y * sh[6] + SH_C2[4] * 2.f * -y * sh[8];
85 dRGBdz += SH_C2[1] * y * sh[5] + SH_C2[2] * 2.f * 2.f * z * sh[6] +
86 SH_C2[3] * x * sh[7];
170}
87171
88 if (deg > 2) {
89 float dRGBdsh9 = SH_C3[0] * y * (3.f * xx - yy);
90 float dRGBdsh10 = SH_C3[1] * xy * z;
91 float dRGBdsh11 = SH_C3[2] * y * (4.f * zz - xx - yy);
92 float dRGBdsh12 = SH_C3[3] * z * (2.f * zz - 3.f * xx - 3.f * yy);
93 float dRGBdsh13 = SH_C3[4] * x * (4.f * zz - xx - yy);
94 float dRGBdsh14 = SH_C3[5] * z * (xx - yy);
95 float dRGBdsh15 = SH_C3[6] * x * (xx - 3.f * yy);
96 dL_dsh[9] = dRGBdsh9 * dL_dRGB;
97 dL_dsh[10] = dRGBdsh10 * dL_dRGB;
98 dL_dsh[11] = dRGBdsh11 * dL_dRGB;
99 dL_dsh[12] = dRGBdsh12 * dL_dRGB;
100 dL_dsh[13] = dRGBdsh13 * dL_dRGB;
101 dL_dsh[14] = dRGBdsh14 * dL_dRGB;
102 dL_dsh[15] = dRGBdsh15 * dL_dRGB;
172// Reverse traversal per pixel; accumulates splat gradients with standard
173// 5-step full-warp butterfly reductions (45 shuffles) and 9 lane-0 atomics.
174__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
175 RasterBackwardKernel(const int2* __restrict__ ranges,
176 const int32_t* __restrict__ point_list,
177 const Splat* __restrict__ splats, Frame frame,
178 const float* __restrict__ final_t,
179 const int32_t* __restrict__ n_contrib,
180 const float* __restrict__ dl_dimage,
181 float* __restrict__ splat_grads) {
182 using L = RasterLayout<1>;
183 constexpr int kRasterThreads = L::kThreads;
184 constexpr int kWarps = L::kWarps;
185 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
186 const int tid = threadIdx.x;
187 const int warp = tid / 32;
188 const int lane = tid % 32;
189 const RasterPixels px = RasterPixels::For<1>(tid);
190 const float fx = static_cast<float>(px.x);
191 const float fy = static_cast<float>(px.y0);
192 const int2 range = ranges[tile];
103193
104 dRGBdx += (SH_C3[0] * sh[9] * 3.f * 2.f * xy + SH_C3[1] * sh[10] * yz +
105 SH_C3[2] * sh[11] * -2.f * xy +
106 SH_C3[3] * sh[12] * -3.f * 2.f * xz +
107 SH_C3[4] * sh[13] * (-3.f * xx + 4.f * zz - yy) +
108 SH_C3[5] * sh[14] * 2.f * xz +
109 SH_C3[6] * sh[15] * 3.f * (xx - yy));
194 const bool inside = px.x < frame.width && px.y0 < frame.height;
195 const int pixels = frame.width * frame.height;
196 const int p = px.y0 * frame.width + px.x;
197 float t = inside ? final_t[p] : 0.0f;
198 const float3 dpix =
199 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
200 dl_dimage[2 * pixels + p]}
201 : float3{0.0f, 0.0f, 0.0f};
202 float accum = 0.0f; // dot(color accumulated behind, dpix)
203 const int last = inside ? n_contrib[p] : 0;
110204
111 dRGBdy +=
112 (SH_C3[0] * sh[9] * 3.f * (xx - yy) + SH_C3[1] * sh[10] * xz +
113 SH_C3[2] * sh[11] * (-3.f * yy + 4.f * zz - xx) +
114 SH_C3[3] * sh[12] * -3.f * 2.f * yz +
115 SH_C3[4] * sh[13] * -2.f * xy + SH_C3[5] * sh[14] * -2.f * yz +
116 SH_C3[6] * sh[15] * -3.f * 2.f * xy);
205 int warp_max_last = last;
206#pragma unroll
207 for (int offset = 16; offset > 0; offset /= 2) {
208 warp_max_last = max(
209 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
210 }
211 __shared__ int shared_max_last[kWarps];
212 if (lane == 0) shared_max_last[warp] = warp_max_last;
213 __syncthreads();
214 int block_max_last = 0;
215#pragma unroll
216 for (int w = 0; w < kWarps; ++w) {
217 block_max_last = max(block_max_last, shared_max_last[w]);
218 }
219 if (block_max_last <= 0) return;
117220
118 dRGBdz += (SH_C3[1] * sh[10] * xy + SH_C3[2] * sh[11] * 4.f * 2.f * yz +
119 SH_C3[3] * sh[12] * 3.f * (2.f * zz - xx - yy) +
120 SH_C3[4] * sh[13] * 4.f * 2.f * xz +
121 SH_C3[5] * sh[14] * (xx - yy));
221 __shared__ float4 batch_geom0[kRasterThreads];
222 __shared__ float2 batch_geom1[kRasterThreads];
223 __shared__ float4 batch_rgb_o[kRasterThreads];
224 __shared__ uint32_t batch_id[kRasterThreads];
225 __shared__ uint8_t batch_mask[kRasterThreads];
226 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
227
228 auto raw_conic = [](float4 g0, float2 g1) {
229 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
230 kInvPower2Scale * g1.x};
231 };
232
233 for (int end = block_max_last; end > 0; end -= kRasterThreads) {
234 const int e = end - 1 - tid;
235 if (e >= 0) {
236 const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
237 const float4 p0 = __ldg(&splats[id].p0);
238 const float4 p1 = __ldg(&splats[id].p1);
239 const float4 p2 = __ldg(&splats[id].p2);
240 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
241 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
242 batch_geom1[tid] = {c2.z, p2.w};
243 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
244 batch_id[tid] = id;
245 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
246 } else {
247 batch_mask[tid] = 0u;
248 }
249 __syncthreads();
250 const int count = min(kRasterThreads, end);
251#pragma unroll
252 for (int c = 0; c < kRasterThreads; c += 32) {
253 if (c < count && end - 32 - c < warp_max_last) {
254 uint32_t mask = __ballot_sync(
255 warp::kFullMask,
256 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
257 while (mask != 0u) {
258 const int bit = __ffs(mask) - 1;
259 mask &= mask - 1u;
260 const int j = c + bit;
261 const int entry = end - 1 - j;
262 float gr[kSplatGradFloats] = {};
263 bool hit = false;
264 if (entry < last) {
265 const float4 g0 = batch_geom0[j];
266 const float2 g1 = batch_geom1[j];
267 const float dx = g0.x - fx;
268 const float dy = g0.y - fy;
269 const float power =
270 g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
271 if (power <= 0.0f && power >= g1.y) {
272 const float4 ro = batch_rgb_o[j];
273 const float g_val = Exp2Approx(power);
274 const float alpha = fminf(kMaxAlpha, ro.w * g_val);
275 if (alpha >= kMinAlpha) {
276 hit = true;
277 t = __fdividef(t, 1.0f - alpha);
278 const float w = alpha * t;
279 const float color_dot =
280 ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
281 const float diff = color_dot - accum;
282 const float dl_dalpha = t * diff;
283 accum += alpha * diff;
284 const float ga = g_val * dl_dalpha;
285 const float dl_dpower = ro.w * ga;
286 const float p_dx = dx * dl_dpower;
287 const float p_dy = dy * dl_dpower;
288 const float3 cn = raw_conic(g0, g1);
289 gr[0] = -(cn.x * p_dx + cn.y * p_dy);
290 gr[1] = -(cn.y * p_dx + cn.z * p_dy);
291 gr[2] = -0.5f * dx * p_dx;
292 gr[3] = -dx * p_dy;
293 gr[4] = -0.5f * dy * p_dy;
294 gr[5] = ga;
295 gr[6] = w * dpix.x;
296 gr[7] = w * dpix.y;
297 gr[8] = w * dpix.z;
122 }298 }
123 }299 }
124 }300 }
125
126 // The view direction is an input to the computation. View direction
127 // is influenced by the Gaussian's mean, so SHs gradients
128 // must propagate back into 3D position.
129 glm::vec3 dL_ddir(glm::dot(dRGBdx, dL_dRGB), glm::dot(dRGBdy, dL_dRGB),
130 glm::dot(dRGBdz, dL_dRGB));
131
132 // Account for normalization of direction
133 float3 dL_dmean = dnormvdv(float3{dir_orig.x, dir_orig.y, dir_orig.z},
134 float3{dL_ddir.x, dL_ddir.y, dL_ddir.z});
135
136 // Gradients of loss w.r.t. Gaussian means, but only the portion
137 // that is caused because the mean affects the view-dependent color.
138 // Additional mean gradient is accumulated in below methods.
139 dL_dmeans[idx] += glm::vec3(dL_dmean.x, dL_dmean.y, dL_dmean.z);
301 if (!__any_sync(warp::kFullMask, hit)) continue;
302 warp::ButterflyReduce(gr);
303 if (lane == 0) {
304 float* dst = splat_grads + kSplatGradFloats * batch_id[j];
305#pragma unroll
306 for (int q = 0; q < kSplatGradFloats; ++q) {
307 atomicAdd(dst + q, gr[q]);
140}308 }
141
142// Backward version of INVERSE 2D covariance matrix computation
143// (due to length launched as separate kernel before other
144// backward steps contained in preprocess)
145__global__ void computeCov2DCUDA(
146 int P, const float3* means, const int* radii, const float* cov3Ds,
147 const float4* __restrict__ conic_opacity, const float h_x, float h_y,
148 const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
149 const bool mip_splatting_enabled, const float* view_matrix,
150 const float* dL_dconics, float3* dL_dmeans, float* dL_dcov,
151 float* dL_dopacity) {
152 auto idx = cg::this_grid().thread_rank();
153 if (idx >= P || !(radii[idx] > 0)) return;
154
155 // Reading location of 3D covariance for this Gaussian
156 const float* cov3D = cov3Ds + 6 * idx;
157
158 // Fetch gradients, recompute 2D covariance and relevant
159 // intermediate forward results needed in the backward.
160 float3 mean = means[idx];
161 float3 dL_dconic = {dL_dconics[4 * idx], dL_dconics[4 * idx + 1],
162 dL_dconics[4 * idx + 3]};
163 float3 t = transformPoint4x3(mean, view_matrix);
164
165 const float limx = 1.3f * tan_fovx;
166 const float limy = 1.3f * tan_fovy;
167 const float txtz = t.x / t.z;
168 const float tytz = t.y / t.z;
169 t.x = min(limx, max(-limx, txtz)) * t.z;
170 t.y = min(limy, max(-limy, tytz)) * t.z;
171
172 const float x_grad_mul = txtz < -limx || txtz > limx ? 0 : 1;
173 const float y_grad_mul = tytz < -limy || tytz > limy ? 0 : 1;
174
175 glm::mat3 J = glm::mat3(h_x / t.z, 0.0f, -(h_x * t.x) / (t.z * t.z), 0.0f,
176 h_y / t.z, -(h_y * t.y) / (t.z * t.z), 0, 0, 0);
177
178 glm::mat3 W = glm::mat3(view_matrix[0], view_matrix[4], view_matrix[8],
179 view_matrix[1], view_matrix[5], view_matrix[9],
180 view_matrix[2], view_matrix[6], view_matrix[10]);
181
182 glm::mat3 Vrk = glm::mat3(cov3D[0], cov3D[1], cov3D[2], cov3D[1], cov3D[3],
183 cov3D[4], cov3D[2], cov3D[4], cov3D[5]);
184
185 glm::mat3 T = W * J;
186
187 glm::mat3 cov2D = glm::transpose(T) * glm::transpose(Vrk) * T;
188
189 float kernel_size_to_use;
190 if (mip_splatting_enabled)
191 kernel_size_to_use = low_pass_2d_kernel_size;
192 else
193 kernel_size_to_use = 0.3f;
194
195 float a = cov2D[0][0] += kernel_size_to_use;
196 float b = cov2D[0][1];
197 float c = cov2D[1][1] += kernel_size_to_use;
198
199 float denom = a * c - b * b;
200 float dL_da = 0, dL_db = 0, dL_dc = 0;
201 float denom2inv = 1.0f / ((denom * denom) + 0.0000001f);
202
203 if (denom2inv != 0) {
204 // Gradients of loss w.r.t. entries of 2D covariance matrix,
205 // given gradients of loss w.r.t. conic matrix (inverse covariance matrix).
206 // e.g., dL / da = dL / d_conic_a * d_conic_a / d_a
207 dL_da = denom2inv * (-c * c * dL_dconic.x + 2 * b * c * dL_dconic.y +
208 (denom - a * c) * dL_dconic.z);
209 dL_dc = denom2inv * (-a * a * dL_dconic.z + 2 * a * b * dL_dconic.y +
210 (denom - a * c) * dL_dconic.x);
211 dL_db = denom2inv * 2 *
212 (b * c * dL_dconic.x - (denom + 2 * b * b) * dL_dconic.y +
213 a * b * dL_dconic.z);
214
215 if (mip_splatting_enabled) {
216 const float det_0 =
217 max(1e-6, cov2D[0][0] * cov2D[1][1] - cov2D[0][1] * cov2D[0][1]);
218 const float det_1 = max(1e-6, (cov2D[0][0] + kernel_size_to_use) *
219 (cov2D[1][1] + kernel_size_to_use) -
220 cov2D[0][1] * cov2D[0][1]);
221 // sqrt here
222 const float coef = sqrt(det_0 / (det_1 + 1e-6) + 1e-6);
223
224 // update the gradient of alpha and save the gradient of dalpha_dcoef
225 // we need opacity as input
226 // new_opacity = coef * opacity
227 // if we know the new opacity, we can derive original opacity and then
228 // dalpha_dcoef = dopacity * opacity
229 const float opacity = conic_opacity[idx].w / (coef + 1e-6);
230 const float dL_dcoef = dL_dopacity[idx] * opacity;
231 const float dL_dsqrtcoef = dL_dcoef * 0.5 * 1. / (coef + 1e-6);
232 const float dL_ddet0 = dL_dsqrtcoef / (det_1 + 1e-6);
233 const float dL_ddet1 =
234 dL_dsqrtcoef * det_0 * (-1.f / (det_1 * det_1 + 1e-6));
235 // TODO gradient is zero if det_0 or det_1 < 0
236 const float dcoef_da = dL_ddet0 * cov2D[1][1] +
237 dL_ddet1 * (cov2D[1][1] + kernel_size_to_use);
238 const float dcoef_db =
239 dL_ddet0 * (-2. * cov2D[0][1]) + dL_ddet1 * (-2. * cov2D[0][1]);
240 const float dcoef_dc = dL_ddet0 * cov2D[0][0] +
241 dL_ddet1 * (cov2D[0][0] + kernel_size_to_use);
242
243 if (det_0 <= 1e-6 || det_1 <= 1e-6) {
244 dL_dopacity[idx] = 0;
245 } else {
246 // Gradiends of alpha respect to conv due to low pass filter
247 dL_da += dcoef_da;
248 dL_dc += dcoef_dc;
249 dL_db += dcoef_db;
250
251 // update dL_dopacity
252 dL_dopacity[idx] = dL_dopacity[idx] * coef;
309 }
310 }
311 }
312 }
313 __syncthreads();
253 }314 }
254 }315}
255316
256 // Gradients of loss L w.r.t. each 3D covariance matrix (Vrk) entry,
317// ---------------------------------------------------------------------------
257 // given gradients w.r.t. 2D covariance matrix (diagonal).318// Preprocess backward: splat gradients -> parameter gradients.
258 // cov2D = transpose(T) * transpose(Vrk) * T;
259 dL_dcov[6 * idx + 0] =
260 (T[0][0] * T[0][0] * dL_da + T[0][0] * T[1][0] * dL_db +
261 T[1][0] * T[1][0] * dL_dc);
262 dL_dcov[6 * idx + 3] =
263 (T[0][1] * T[0][1] * dL_da + T[0][1] * T[1][1] * dL_db +
264 T[1][1] * T[1][1] * dL_dc);
265 dL_dcov[6 * idx + 5] =
266 (T[0][2] * T[0][2] * dL_da + T[0][2] * T[1][2] * dL_db +
267 T[1][2] * T[1][2] * dL_dc);
319// Parameter gradient buffers are bulk-zeroed beforehand via `cudaMemsetAsync`;
320// culled or zero-gradient Gaussians return early.
268321
269 // Gradients of loss L w.r.t. each 3D covariance matrix (Vrk) entry,
270 // given gradients w.r.t. 2D covariance matrix (off-diagonal).
271 // Off-diagonal elements appear twice --> double the gradient.
272 // cov2D = transpose(T) * transpose(Vrk) * T;
273 dL_dcov[6 * idx + 1] = 2 * T[0][0] * T[0][1] * dL_da +
274 (T[0][0] * T[1][1] + T[0][1] * T[1][0]) * dL_db +
275 2 * T[1][0] * T[1][1] * dL_dc;
276 dL_dcov[6 * idx + 2] = 2 * T[0][0] * T[0][2] * dL_da +
277 (T[0][0] * T[1][2] + T[0][2] * T[1][0]) * dL_db +
278 2 * T[1][0] * T[1][2] * dL_dc;
279 dL_dcov[6 * idx + 4] = 2 * T[0][2] * T[0][1] * dL_da +
280 (T[0][1] * T[1][2] + T[0][2] * T[1][1]) * dL_db +
281 2 * T[1][1] * T[1][2] * dL_dc;
282 } else {
322template <typename ShT>
323__global__ void __launch_bounds__(kGaussianThreads)
324 PreprocessBackwardKernel(GaussianParams params, Frame frame,
325 const Splat* __restrict__ splats,
326 const float* __restrict__ splat_grads,
327 GaussianGrads grads) {
328 const int i = blockIdx.x * blockDim.x + threadIdx.x;
329 if (i >= params.n || splats[i].radius() == 0) return;
330 float g[kSplatGradFloats];
331 bool nonzero = false;
332#pragma unroll
283 for (int i = 0; i < 6; i++) dL_dcov[6 * idx + i] = 0;333 for (int q = 0; q < kSplatGradFloats; ++q) {
334 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
335 nonzero = nonzero || g[q] != 0.0f;
336 }
337 if (!nonzero) return;
338 ProjectionGrad pg;
339 pg.xy = {g[0], g[1]};
340 pg.conic = {g[2], g[3], g[4]};
341 pg.opacity = g[5];
342 pg.rgb = {g[6], g[7], g[8]};
343 float3 dmean, dscale;
344 float4 dquat;
345 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
346 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
347 grads.means[3 * i] = dmean.x;
348 grads.means[3 * i + 1] = dmean.y;
349 grads.means[3 * i + 2] = dmean.z;
350 grads.scales[3 * i] = dscale.x;
351 grads.scales[3 * i + 1] = dscale.y;
352 grads.scales[3 * i + 2] = dscale.z;
353 reinterpret_cast<float4*>(grads.quats)[i] = dquat;
354 grads.opacities[i] = g[5];
284 }355}
285356
286 // Gradients of loss w.r.t. upper 2x3 portion of intermediate matrix T
287 // cov2D = transpose(T) * transpose(Vrk) * T;
288 float dL_dT00 =
289 2 * (T[0][0] * Vrk[0][0] + T[0][1] * Vrk[0][1] + T[0][2] * Vrk[0][2]) *
290 dL_da +
291 (T[1][0] * Vrk[0][0] + T[1][1] * Vrk[0][1] + T[1][2] * Vrk[0][2]) * dL_db;
292 float dL_dT01 =
293 2 * (T[0][0] * Vrk[1][0] + T[0][1] * Vrk[1][1] + T[0][2] * Vrk[1][2]) *
294 dL_da +
295 (T[1][0] * Vrk[1][0] + T[1][1] * Vrk[1][1] + T[1][2] * Vrk[1][2]) * dL_db;
296 float dL_dT02 =
297 2 * (T[0][0] * Vrk[2][0] + T[0][1] * Vrk[2][1] + T[0][2] * Vrk[2][2]) *
298 dL_da +
299 (T[1][0] * Vrk[2][0] + T[1][1] * Vrk[2][1] + T[1][2] * Vrk[2][2]) * dL_db;
300 float dL_dT10 =
301 2 * (T[1][0] * Vrk[0][0] + T[1][1] * Vrk[0][1] + T[1][2] * Vrk[0][2]) *
302 dL_dc +
303 (T[0][0] * Vrk[0][0] + T[0][1] * Vrk[0][1] + T[0][2] * Vrk[0][2]) * dL_db;
304 float dL_dT11 =
305 2 * (T[1][0] * Vrk[1][0] + T[1][1] * Vrk[1][1] + T[1][2] * Vrk[1][2]) *
306 dL_dc +
307 (T[0][0] * Vrk[1][0] + T[0][1] * Vrk[1][1] + T[0][2] * Vrk[1][2]) * dL_db;
308 float dL_dT12 =
309 2 * (T[1][0] * Vrk[2][0] + T[1][1] * Vrk[2][1] + T[1][2] * Vrk[2][2]) *
310 dL_dc +
311 (T[0][0] * Vrk[2][0] + T[0][1] * Vrk[2][1] + T[0][2] * Vrk[2][2]) * dL_db;
312
313 // Gradients of loss w.r.t. upper 3x2 non-zero entries of Jacobian matrix
314 // T = W * J
315 float dL_dJ00 = W[0][0] * dL_dT00 + W[0][1] * dL_dT01 + W[0][2] * dL_dT02;
316 float dL_dJ02 = W[2][0] * dL_dT00 + W[2][1] * dL_dT01 + W[2][2] * dL_dT02;
317 float dL_dJ11 = W[1][0] * dL_dT10 + W[1][1] * dL_dT11 + W[1][2] * dL_dT12;
318 float dL_dJ12 = W[2][0] * dL_dT10 + W[2][1] * dL_dT11 + W[2][2] * dL_dT12;
319
320 float tz = 1.f / t.z;
321 float tz2 = tz * tz;
322 float tz3 = tz2 * tz;
323
324 // Gradients of loss w.r.t. transformed Gaussian mean t
325 float dL_dtx = x_grad_mul * -h_x * tz2 * dL_dJ02;
326 float dL_dty = y_grad_mul * -h_y * tz2 * dL_dJ12;
327 float dL_dtz = -h_x * tz2 * dL_dJ00 - h_y * tz2 * dL_dJ11 +
328 (2 * h_x * t.x) * tz3 * dL_dJ02 +
329 (2 * h_y * t.y) * tz3 * dL_dJ12;
330
331 // Account for transformation of mean to t
332 // t = transformPoint4x3(mean, view_matrix);
333 float3 dL_dmean =
334 transformVec4x3Transpose({dL_dtx, dL_dty, dL_dtz}, view_matrix);
335
357template <typename ShT>
358cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
359 const Frame& frame, const Splat* splats,
360 const float* splat_grads,
361 const GaussianGrads& grads,
362 cudaStream_t stream) {
363 const int n = params.n;
364 // As Inria's PyTorch wrapper (`torch::zeros`): bulk-zero every parameter
336 // Gradients of loss w.r.t. Gaussian means, but only the portion365 // gradient up front; culled Gaussians then exit early in the kernel.
337 // that is caused because the mean affects the covariance matrix.
338 // Additional mean gradient is accumulated in BACKWARD::preprocess.
339 dL_dmeans[idx] = dL_dmean;
366 if (auto e = cudaMemsetAsync(grads.means, 0, sizeof(float) * 3 * n, stream);
367 e != cudaSuccess)
368 return e;
369 if (auto e = cudaMemsetAsync(grads.scales, 0, sizeof(float) * 3 * n, stream);
370 e != cudaSuccess)
371 return e;
372 if (auto e = cudaMemsetAsync(grads.quats, 0, sizeof(float) * 4 * n, stream);
373 e != cudaSuccess)
374 return e;
375 if (auto e = cudaMemsetAsync(grads.opacities, 0, sizeof(float) * n, stream);
376 e != cudaSuccess)
377 return e;
378 if (auto e =
379 cudaMemsetAsync(grads.sh, 0, sizeof(ShT) * 3 * kShCoeffs * n, stream);
380 e != cudaSuccess)
381 return e;
382 PreprocessBackwardKernel<ShT>
383 <<<GaussianBlocks(n), kGaussianThreads, 0, stream>>>(
384 params, frame, splats, splat_grads, grads);
385 return cudaSuccess;
340}386}
341387
342// Backward pass for the conversion of scale and rotation to a
343// 3D covariance matrix for each Gaussian.
344__device__ void computeCov3D(int idx, const glm::vec3 scale, float mod,
388template __global__ void PreprocessBackwardKernel<float>(
389 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
390 GaussianGrads);
391template __global__ void PreprocessBackwardKernel<ShBf16>(
392 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
393 GaussianGrads);
394template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
345 const glm::vec4 rot, const float* dL_dcov3Ds,395 const Frame&, const Splat*,
346 glm::vec3* dL_dscales, glm::vec4* dL_drots) {
347 // Recompute (intermediate) results for the 3D covariance computation.
348 glm::vec4 q = rot; // / glm::length(rot);
349 float r = q.x;
350 float x = q.y;
351 float y = q.z;
352 float z = q.w;
396 const float*,
397 const GaussianGrads&,
398 cudaStream_t);
399template cudaError_t LaunchPreprocessBackward<ShBf16>(
400 const GaussianParams&, const Frame&, const Splat*, const float*,
401 const GaussianGrads&, cudaStream_t);
353402
354 glm::mat3 R = glm::mat3(1.f - 2.f * (y * y + z * z), 2.f * (x * y - r * z),
355 2.f * (x * z + r * y), 2.f * (x * y + r * z),
356 1.f - 2.f * (x * x + z * z), 2.f * (y * z - r * x),
357 2.f * (x * z - r * y), 2.f * (y * z + r * x),
358 1.f - 2.f * (x * x + y * y));
359
360 glm::mat3 S = glm::mat3(1.0f);
361
362 glm::vec3 s = mod * scale;
363 S[0][0] = s.x;
364 S[1][1] = s.y;
365 S[2][2] = s.z;
366
367 glm::mat3 M = S * R;
368
369 const float* dL_dcov3D = dL_dcov3Ds + 6 * idx;
370
371 glm::vec3 dunc(dL_dcov3D[0], dL_dcov3D[3], dL_dcov3D[5]);
372 glm::vec3 ounc = 0.5f * glm::vec3(dL_dcov3D[1], dL_dcov3D[2], dL_dcov3D[4]);
373
374 // Convert per-element covariance loss gradients to matrix form
375 glm::mat3 dL_dSigma =
376 glm::mat3(dL_dcov3D[0], 0.5f * dL_dcov3D[1], 0.5f * dL_dcov3D[2],
377 0.5f * dL_dcov3D[1], dL_dcov3D[3], 0.5f * dL_dcov3D[4],
378 0.5f * dL_dcov3D[2], 0.5f * dL_dcov3D[4], dL_dcov3D[5]);
379
380 // Compute loss gradient w.r.t. matrix M
381 // dSigma_dM = 2 * M
382 glm::mat3 dL_dM = 2.0f * M * dL_dSigma;
383
384 glm::mat3 Rt = glm::transpose(R);
385 glm::mat3 dL_dMt = glm::transpose(dL_dM);
386
387 // Gradients of loss w.r.t. scale
388 glm::vec3* dL_dscale = dL_dscales + idx;
389 dL_dscale->x = glm::dot(Rt[0], dL_dMt[0]);
390 dL_dscale->y = glm::dot(Rt[1], dL_dMt[1]);
391 dL_dscale->z = glm::dot(Rt[2], dL_dMt[2]);
392
393 dL_dMt[0] *= s.x;
394 dL_dMt[1] *= s.y;
395 dL_dMt[2] *= s.z;
396
397 // Gradients of loss w.r.t. normalized quaternion
398 glm::vec4 dL_dq;
399 dL_dq.x = 2 * z * (dL_dMt[0][1] - dL_dMt[1][0]) +
400 2 * y * (dL_dMt[2][0] - dL_dMt[0][2]) +
401 2 * x * (dL_dMt[1][2] - dL_dMt[2][1]);
402 dL_dq.y = 2 * y * (dL_dMt[1][0] + dL_dMt[0][1]) +
403 2 * z * (dL_dMt[2][0] + dL_dMt[0][2]) +
404 2 * r * (dL_dMt[1][2] - dL_dMt[2][1]) -
405 4 * x * (dL_dMt[2][2] + dL_dMt[1][1]);
406 dL_dq.z = 2 * x * (dL_dMt[1][0] + dL_dMt[0][1]) +
407 2 * r * (dL_dMt[2][0] - dL_dMt[0][2]) +
408 2 * z * (dL_dMt[1][2] + dL_dMt[2][1]) -
409 4 * y * (dL_dMt[2][2] + dL_dMt[0][0]);
410 dL_dq.w = 2 * r * (dL_dMt[0][1] - dL_dMt[1][0]) +
411 2 * x * (dL_dMt[2][0] + dL_dMt[0][2]) +
412 2 * y * (dL_dMt[1][2] + dL_dMt[2][1]) -
413 4 * z * (dL_dMt[1][1] + dL_dMt[0][0]);
414
415 // Gradients of loss w.r.t. unnormalized quaternion
416 float4* dL_drot = (float4*)(dL_drots + idx);
417 *dL_drot = float4{dL_dq.x, dL_dq.y, dL_dq.z,
418 dL_dq.w}; // dnormvdv(float4{ rot.x, rot.y, rot.z, rot.w },
419 // float4{ dL_dq.x, dL_dq.y, dL_dq.z, dL_dq.w });
420}
421
422// Backward pass of the preprocessing steps, except
423// for the covariance computation and inversion
424// (those are handled by a previous kernel call)
425template <int C>
426__global__ void preprocessCUDA(
427 int P, int D, int M, const float3* means, const int* radii,
428 const float* shs, const bool* clamped, const glm::vec3* scales,
429 const glm::vec4* rotations, const float scale_modifier, const float* proj,
430 const glm::vec3* campos, const float3* dL_dmean2D, glm::vec3* dL_dmeans,
431 float* dL_dcolor, float* dL_dcov3D, float* dL_dsh, glm::vec3* dL_dscale,
432 glm::vec4* dL_drot) {
433 auto idx = cg::this_grid().thread_rank();
434 if (idx >= P || !(radii[idx] > 0)) return;
435
436 float3 m = means[idx];
437
438 // Taking care of gradients from the screenspace points
439 float4 m_hom = transformPoint4x4(m, proj);
440 float m_w = 1.0f / (m_hom.w + 0.0000001f);
441
442 // Compute loss gradient w.r.t. 3D means due to gradients of 2D means
443 // from rendering procedure
444 glm::vec3 dL_dmean;
445 float mul1 =
446 (proj[0] * m.x + proj[4] * m.y + proj[8] * m.z + proj[12]) * m_w * m_w;
447 float mul2 =
448 (proj[1] * m.x + proj[5] * m.y + proj[9] * m.z + proj[13]) * m_w * m_w;
449 dL_dmean.x = (proj[0] * m_w - proj[3] * mul1) * dL_dmean2D[idx].x +
450 (proj[1] * m_w - proj[3] * mul2) * dL_dmean2D[idx].y;
451 dL_dmean.y = (proj[4] * m_w - proj[7] * mul1) * dL_dmean2D[idx].x +
452 (proj[5] * m_w - proj[7] * mul2) * dL_dmean2D[idx].y;
453 dL_dmean.z = (proj[8] * m_w - proj[11] * mul1) * dL_dmean2D[idx].x +
454 (proj[9] * m_w - proj[11] * mul2) * dL_dmean2D[idx].y;
455
456 // That's the second part of the mean gradient. Previous computation
457 // of cov2D and following SH conversion also affects it.
458 dL_dmeans[idx] += dL_dmean;
459
460 // Compute gradient updates due to computing colors from SHs
461 if (shs)
462 computeColorFromSH(idx, D, M, (glm::vec3*)means, *campos, shs, clamped,
463 (glm::vec3*)dL_dcolor, (glm::vec3*)dL_dmeans,
464 (glm::vec3*)dL_dsh);
465
466 // Compute gradient updates due to computing covariance from scale/rotation
467 if (scales)
468 computeCov3D(idx, scales[idx], scale_modifier, rotations[idx], dL_dcov3D,
469 dL_dscale, dL_drot);
470}
471
472// Backward version of the rendering procedure.
473template <uint32_t C>
474__global__ void __launch_bounds__(BLOCK_X* BLOCK_Y) renderCUDA(
475 const uint2* __restrict__ ranges, const uint32_t* __restrict__ point_list,
476 int W, int H, const float* __restrict__ bg_color,
477 const float2* __restrict__ points_xy_image,
478 const float4* __restrict__ conic_opacity, const float* __restrict__ colors,
479 const float* __restrict__ final_Ts, const uint32_t* __restrict__ n_contrib,
480 const float* __restrict__ dL_dpixels, float3* __restrict__ dL_dmean2D,
481 float4* __restrict__ dL_dconic2D, float* __restrict__ dL_dopacity,
482 float* __restrict__ dL_dcolors) {
483 // We rasterize again. Compute necessary block info.
484 auto block = cg::this_thread_block();
485 const uint32_t horizontal_blocks = (W + BLOCK_X - 1) / BLOCK_X;
486 const uint2 pix_min = {block.group_index().x * BLOCK_X,
487 block.group_index().y * BLOCK_Y};
488 const uint2 pix_max = {min(pix_min.x + BLOCK_X, W),
489 min(pix_min.y + BLOCK_Y, H)};
490 const uint2 pix = {pix_min.x + block.thread_index().x,
491 pix_min.y + block.thread_index().y};
492 const uint32_t pix_id = W * pix.y + pix.x;
493 const float2 pixf = {(float)pix.x, (float)pix.y};
494
495 const bool inside = pix.x < W && pix.y < H;
496 const uint2 range =
497 ranges[block.group_index().y * horizontal_blocks + block.group_index().x];
498
499 const int rounds = ((range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE);
500
501 bool done = !inside;
502 int toDo = range.y - range.x;
503
504 __shared__ int collected_id[BLOCK_SIZE];
505 __shared__ float2 collected_xy[BLOCK_SIZE];
506 __shared__ float4 collected_conic_opacity[BLOCK_SIZE];
507 __shared__ float collected_colors[C * BLOCK_SIZE];
508
509 // In the forward, we stored the final value for T, the
510 // product of all (1 - alpha) factors.
511 const float T_final = inside ? final_Ts[pix_id] : 0;
512 float T = T_final;
513
514 // We start from the back. The ID of the last contributing
515 // Gaussian is known from each pixel from the forward.
516 uint32_t contributor = toDo;
517 const int last_contributor = inside ? n_contrib[pix_id] : 0;
518
519 float accum_rec[C] = {0};
520 float dL_dpixel[C];
521 if (inside)
522 for (int i = 0; i < C; i++) dL_dpixel[i] = dL_dpixels[i * H * W + pix_id];
523
524 float last_alpha = 0;
525 float last_color[C] = {0};
526
527 // Gradient of pixel coordinate w.r.t. normalized
528 // screen-space viewport corrdinates (-1 to 1)
529 const float ddelx_dx = 0.5 * W;
530 const float ddely_dy = 0.5 * H;
531
532 // Traverse all Gaussians
533 for (int i = 0; i < rounds; i++, toDo -= BLOCK_SIZE) {
534 // Load auxiliary data into shared memory, start in the BACK
535 // and load them in revers order.
536 block.sync();
537 const int progress = i * BLOCK_SIZE + block.thread_rank();
538 if (range.x + progress < range.y) {
539 const int coll_id = point_list[range.y - progress - 1];
540 collected_id[block.thread_rank()] = coll_id;
541 collected_xy[block.thread_rank()] = points_xy_image[coll_id];
542 collected_conic_opacity[block.thread_rank()] = conic_opacity[coll_id];
543 for (int i = 0; i < C; i++)
544 collected_colors[i * BLOCK_SIZE + block.thread_rank()] =
545 colors[coll_id * C + i];
546 }
547 block.sync();
548
549 // Iterate over Gaussians
550 for (int j = 0; !done && j < min(BLOCK_SIZE, toDo); j++) {
551 // Keep track of current Gaussian ID. Skip, if this one
552 // is behind the last contributor for this pixel.
553 contributor--;
554 if (contributor >= last_contributor) continue;
555
556 // Compute blending values, as before.
557 const float2 xy = collected_xy[j];
558 const float2 d = {xy.x - pixf.x, xy.y - pixf.y};
559 const float4 con_o = collected_conic_opacity[j];
560 const float power = -0.5f * (con_o.x * d.x * d.x + con_o.z * d.y * d.y) -
561 con_o.y * d.x * d.y;
562 if (power > 0.0f) continue;
563
564 const float G = exp(power);
565 const float alpha = min(0.99f, con_o.w * G);
566 if (alpha < 1.0f / 255.0f) continue;
567
568 T = T / (1.f - alpha);
569 const float dchannel_dcolor = alpha * T;
570
571 // Propagate gradients to per-Gaussian colors and keep
572 // gradients w.r.t. alpha (blending factor for a Gaussian/pixel
573 // pair).
574 float dL_dalpha = 0.0f;
575 const int global_id = collected_id[j];
576 for (int ch = 0; ch < C; ch++) {
577 const float c = collected_colors[ch * BLOCK_SIZE + j];
578 // Update last color (to be used in the next iteration)
579 accum_rec[ch] =
580 last_alpha * last_color[ch] + (1.f - last_alpha) * accum_rec[ch];
581 last_color[ch] = c;
582
583 const float dL_dchannel = dL_dpixel[ch];
584 dL_dalpha += (c - accum_rec[ch]) * dL_dchannel;
585 // Update the gradients w.r.t. color of the Gaussian.
586 // Atomic, since this pixel is just one of potentially
587 // many that were affected by this Gaussian.
588 atomicAdd(&(dL_dcolors[global_id * C + ch]),
589 dchannel_dcolor * dL_dchannel);
590 }
591 dL_dalpha *= T;
592 // Update last alpha (to be used in the next iteration)
593 last_alpha = alpha;
594
595 // Account for fact that alpha also influences how much of
596 // the background color is added if nothing left to blend
597 float bg_dot_dpixel = 0;
598 for (int i = 0; i < C; i++) bg_dot_dpixel += bg_color[i] * dL_dpixel[i];
599 dL_dalpha += (-T_final / (1.f - alpha)) * bg_dot_dpixel;
600
601 // Helpful reusable temporary variables
602 const float dL_dG = con_o.w * dL_dalpha;
603 const float gdx = G * d.x;
604 const float gdy = G * d.y;
605 const float dG_ddelx = -gdx * con_o.x - gdy * con_o.y;
606 const float dG_ddely = -gdy * con_o.z - gdx * con_o.y;
607
608 // Update gradients w.r.t. 2D mean position of the Gaussian
609 atomicAdd(&dL_dmean2D[global_id].x, dL_dG * dG_ddelx * ddelx_dx);
610 atomicAdd(&dL_dmean2D[global_id].y, dL_dG * dG_ddely * ddely_dy);
611
612 // Update gradients w.r.t. 2D covariance (2x2 matrix, symmetric)
613 atomicAdd(&dL_dconic2D[global_id].x, -0.5f * gdx * d.x * dL_dG);
614 atomicAdd(&dL_dconic2D[global_id].y, -0.5f * gdx * d.y * dL_dG);
615 atomicAdd(&dL_dconic2D[global_id].w, -0.5f * gdy * d.y * dL_dG);
616
617 // Update gradients w.r.t. opacity of the Gaussian
618 atomicAdd(&(dL_dopacity[global_id]), G * dL_dalpha);
619 }
620 }
621}
622
623void BACKWARD::preprocess(
624 int P, int D, int M, const float3* means3D, const int* radii,
625 const float* shs, const bool* clamped, const glm::vec3* scales,
626 const glm::vec4* rotations, const float scale_modifier, const float* cov3Ds,
627 const float4* conic_opacity, const float* viewmatrix,
628 const float* projmatrix, const float focal_x, float focal_y,
629 const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
630 const bool mip_splatting_enabled, const glm::vec3* campos,
631 const float3* dL_dmean2D, const float* dL_dconic, glm::vec3* dL_dmean3D,
632 float* dL_dcolor, float* dL_dcov3D, float* dL_dsh, glm::vec3* dL_dscale,
633 glm::vec4* dL_drot, float* dL_dopacity) {
634 // Propagate gradients for the path of 2D conic matrix computation.
635 // Somewhat long, thus it is its own kernel rather than being part of
636 // "preprocess". When done, loss gradient w.r.t. 3D means has been
637 // modified and gradient w.r.t. 3D covariance matrix has been computed.
638 computeCov2DCUDA<<<(P + 255) / 256, 256>>>(
639 P, means3D, radii, cov3Ds, conic_opacity, focal_x, focal_y, tan_fovx,
640 tan_fovy, low_pass_2d_kernel_size, mip_splatting_enabled, viewmatrix,
641 dL_dconic, (float3*)dL_dmean3D, dL_dcov3D, dL_dopacity);
642
643 // Propagate gradients for remaining steps: finish 3D mean gradients,
644 // propagate color gradients to SH (if desireD), propagate 3D covariance
645 // matrix gradients to scale and rotation.
646 preprocessCUDA<NUM_CHANNELS><<<(P + 255) / 256, 256>>>(
647 P, D, M, (float3*)means3D, radii, shs, clamped, (glm::vec3*)scales,
648 (glm::vec4*)rotations, scale_modifier, projmatrix, campos,
649 (float3*)dL_dmean2D, (glm::vec3*)dL_dmean3D, dL_dcolor, dL_dcov3D, dL_dsh,
650 dL_dscale, dL_drot);
651}
652
653void BACKWARD::render(const dim3 grid, const dim3 block, const uint2* ranges,
654 const uint32_t* point_list, int W, int H,
655 const float* bg_color, const float2* means2D,
656 const float4* conic_opacity, const float* colors,
657 const float* final_Ts, const uint32_t* n_contrib,
658 const float* dL_dpixels, float3* dL_dmean2D,
659 float4* dL_dconic2D, float* dL_dopacity,
660 float* dL_dcolors) {
661 renderCUDA<NUM_CHANNELS>
662 <<<grid, block>>>(ranges, point_list, W, H, bg_color, means2D,
663 conic_opacity, colors, final_Ts, n_contrib, dL_dpixels,
664 dL_dmean2D, dL_dconic2D, dL_dopacity, dL_dcolors);
665}
403} // namespace gs
↕ Expand 9 unchanged lines (lines 1–9)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #include "backward.h"
1313
14−#include <cooperative_groups.h>
14+#include "auxiliary.h"
15−#include <cooperative_groups/reduce.h>
1615
17−#include "auxiliary.h"
18−namespace cg = cooperative_groups;
16+namespace gs {
1917
20−// Backward pass for conversion of spherical harmonics to RGB for
21−// each Gaussian.
22−__device__ void computeColorFromSH(int idx, int deg, int max_coeffs,
23− const glm::vec3* means, glm::vec3 campos,
18+namespace sh {
19+// Evaluates SH color VJP loading the full 16-basis row `sh_in` and storing
20+// `dsh` to `sh_out`.
21+template <typename ShT>
22+GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,
23+ ShT* __restrict__ sh_out) {
24+ float b[kShCoeffs];
25+ Basis(d, b);
26+ float3 coeffs[kShCoeffs];
27+ ShIo<ShT>::LoadRow(sh_in, coeffs);
28+ float3 raw = {0.5f, 0.5f, 0.5f};
29+#pragma unroll
30+ for (int k = 0; k < kShCoeffs; ++k) {
31+ raw += b[k] * coeffs[k];
32+ }
33+ const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
34+ raw.y < 0.0f ? 0.0f : g_rgb.y,
35+ raw.z < 0.0f ? 0.0f : g_rgb.z};
36+ float3 bgrad[kShCoeffs];
37+ BasisGrad(d, bgrad);
38+ float3 dsh[kShCoeffs];
39+ float3 ddir = {0.0f, 0.0f, 0.0f};
40+#pragma unroll
41+ for (int k = 0; k < kShCoeffs; ++k) {
42+ dsh[k] = b[k] * drgb;
24− const float* shs, const bool* clamped,
43+ const float s = dot(coeffs[k], drgb);
25− const glm::vec3* dL_dcolor,
26− glm::vec3* dL_dmeans, glm::vec3* dL_dshs) {
27− // Compute intermediate values, as it is done during forward
28− glm::vec3 pos = means[idx];
29− glm::vec3 dir_orig = pos - campos;
30− glm::vec3 dir = dir_orig / glm::length(dir_orig);
44+ ddir += s * bgrad[k];
45+ }
46+ ShIo<ShT>::StoreRow(sh_out, dsh);
47+ return ddir;
48+}
49+} // namespace sh
3150
32− glm::vec3* sh = ((glm::vec3*)shs) + idx * max_coeffs;
51+// VJP of the projection for a touched Gaussian `i`; reads its SH row from
52+// `sh_in` and writes dL/dsh directly to `sh_out`.
53+// Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
54+template <typename ShT>
55+GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
56+ const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
57+ const ProjectionGrad& g, float3& dmean_out,
58+ float3& dscale_out, float4& dquat_out) {
59+ const float3 mean = LoadMean(p, i);
60+ const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
61+ p.scales[3 * i + 2]};
62+ const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
3363
34− // Use PyTorch rule for clamping: if clamping was applied,
35− // gradient becomes 0.
36− glm::vec3 dL_dRGB = dL_dcolor[idx];
37− dL_dRGB.x *= clamped[3 * idx + 0] ? 0 : 1;
38− dL_dRGB.y *= clamped[3 * idx + 1] ? 0 : 1;
39− dL_dRGB.z *= clamped[3 * idx + 2] ? 0 : 1;
64+ // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
65+ float3 dmean = {0.0f, 0.0f, 0.0f};
66+ {
67+ float inv_len;
68+ const float3 dir = ViewDir(cam, mean, inv_len);
69+ const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
70+ // d(v/|v|)/dv = (I - dir dir^T) / |v|.
71+ dmean += inv_len * (ddir - dot(dir, ddir) * dir);
72+ }
4073
41− glm::vec3 dRGBdx(0, 0, 0);
42− glm::vec3 dRGBdy(0, 0, 0);
43− glm::vec3 dRGBdz(0, 0, 0);
44− float x = dir.x;
74+ // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
45− float y = dir.y;
75+ const float3 t = cam.ToView(mean);
46− float z = dir.z;
76+ const float tz = t.z;
77+ const float limx = kFovClamp * cam.TanHalfFovX();
78+ const float limy = kFovClamp * cam.TanHalfFovY();
79+ const float txtz = t.x / tz, tytz = t.y / tz;
80+ const bool x_clamped = txtz < -limx || txtz > limx;
81+ const bool y_clamped = tytz < -limy || tytz > limy;
82+ const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
83+ const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
84+ const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
85+ const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
86+ const Mat3& w = cam.rotation;
87+ const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
88+ const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
4789
48− // Target location for this Gaussian to write SH gradients to
49− glm::vec3* dL_dsh = dL_dshs + idx * max_coeffs;
90+ const Mat3 r = QuatToRotation(quat);
91+ const Sym3 sigma = Covariance3D(r, scale);
92+ const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
93+ const float cov_a = dot(m0, s0) + kLowPass;
94+ const float cov_b = dot(m0, s1);
95+ const float cov_c = dot(m1, s1) + kLowPass;
5096
51− // No tricks here, just high school-level calculus.
52− float dRGBdsh0 = SH_C0;
53− dL_dsh[0] = dRGBdsh0 * dL_dRGB;
54− if (deg > 0) {
55− float dRGBdsh1 = -SH_C1 * y;
97+ // --- Conic -> 2D covariance: G_cov = -K G_K K.
98+ const float det = cov_a * cov_c - cov_b * cov_b;
56− float dRGBdsh2 = SH_C1 * z;
99+ const float inv = 1.0f / det;
57− float dRGBdsh3 = -SH_C1 * x;
58− dL_dsh[1] = dRGBdsh1 * dL_dRGB;
59− dL_dsh[2] = dRGBdsh2 * dL_dRGB;
60− dL_dsh[3] = dRGBdsh3 * dL_dRGB;
100+ const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
101+ const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
102+ // P = G_K K.
103+ const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
104+ const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
105+ // G_cov = -K P (symmetric).
106+ const float ca = -(ka * p00 + kb * p10);
107+ const float cb = -(ka * p01 + kb * p11);
108+ const float cc = -(kb * p01 + kc * p11);
61109
62− dRGBdx = -SH_C1 * sh[3];
63− dRGBdy = -SH_C1 * sh[1];
64− dRGBdz = SH_C1 * sh[2];
110+ // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
111+ const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
112+ const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
113+ // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
114+ Sym3 dsig;
115+ dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
116+ dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
117+ dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
118+ dsig.xy =
119+ ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
120+ dsig.xz =
121+ ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
122+ dsig.yz =
123+ ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
65124
66− if (deg > 1) {
67− float xx = x * x, yy = y * y, zz = z * z;
125+ // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
126+ const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
127+ const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
68− float xy = x * y, yz = y * z, xz = x * z;
128+ const float tz2 = tz * tz, tz3 = tz2 * tz;
129+ float3 dt;
130+ dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
131+ dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
132+ dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
133+ 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
134+ // --- Screen-space mean u = fx x / z + cx.
135+ dt.x += cam.fx / tz * g.xy.x;
136+ dt.y += cam.fy / tz * g.xy.y;
137+ dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
138+ dmean += MulT(w, dt);
139+ dmean_out = dmean;
69140
70− float dRGBdsh4 = SH_C2[0] * xy;
71− float dRGBdsh5 = SH_C2[1] * yz;
72− float dRGBdsh6 = SH_C2[2] * (2.f * zz - xx - yy);
141+ // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
142+ const float sv[3] = {scale.x, scale.y, scale.z};
73− float dRGBdsh7 = SH_C2[3] * xz;
143+ float ds[3];
74− float dRGBdsh8 = SH_C2[4] * (xx - yy);
75− dL_dsh[4] = dRGBdsh4 * dL_dRGB;
76− dL_dsh[5] = dRGBdsh5 * dL_dRGB;
77− dL_dsh[6] = dRGBdsh6 * dL_dRGB;
78− dL_dsh[7] = dRGBdsh7 * dL_dRGB;
79− dL_dsh[8] = dRGBdsh8 * dL_dRGB;
144+ Mat3 dr;
145+#pragma unroll
146+ for (int j = 0; j < 3; ++j) {
147+ const float3 r_col = r.col(j);
148+ const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
149+ ds[j] = dot(da_col, r_col);
150+ dr.m[0][j] = da_col.x * sv[j];
151+ dr.m[1][j] = da_col.y * sv[j];
152+ dr.m[2][j] = da_col.z * sv[j];
153+ }
154+ dscale_out = {ds[0], ds[1], ds[2]};
80155
156+ // --- R(q) -> q.
157+ const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
158+ const auto& G = dr.m;
159+ dquat_out = {
160+ 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
161+ qy * G[2][0] + qx * G[2][1]),
162+ 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
81− dRGBdx += SH_C2[0] * y * sh[4] + SH_C2[2] * 2.f * -x * sh[6] +
163+ qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
164+ 2.0f *
165+ (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
166+ qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
167+ 2.0f *
168+ (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
82− SH_C2[3] * z * sh[7] + SH_C2[4] * 2.f * x * sh[8];
169+ 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
83− dRGBdy += SH_C2[0] * x * sh[4] + SH_C2[1] * z * sh[5] +
84− SH_C2[2] * 2.f * -y * sh[6] + SH_C2[4] * 2.f * -y * sh[8];
85− dRGBdz += SH_C2[1] * y * sh[5] + SH_C2[2] * 2.f * 2.f * z * sh[6] +
86− SH_C2[3] * x * sh[7];
170+}
87171
88− if (deg > 2) {
89− float dRGBdsh9 = SH_C3[0] * y * (3.f * xx - yy);
90− float dRGBdsh10 = SH_C3[1] * xy * z;
91− float dRGBdsh11 = SH_C3[2] * y * (4.f * zz - xx - yy);
92− float dRGBdsh12 = SH_C3[3] * z * (2.f * zz - 3.f * xx - 3.f * yy);
93− float dRGBdsh13 = SH_C3[4] * x * (4.f * zz - xx - yy);
94− float dRGBdsh14 = SH_C3[5] * z * (xx - yy);
95− float dRGBdsh15 = SH_C3[6] * x * (xx - 3.f * yy);
96− dL_dsh[9] = dRGBdsh9 * dL_dRGB;
97− dL_dsh[10] = dRGBdsh10 * dL_dRGB;
98− dL_dsh[11] = dRGBdsh11 * dL_dRGB;
99− dL_dsh[12] = dRGBdsh12 * dL_dRGB;
100− dL_dsh[13] = dRGBdsh13 * dL_dRGB;
101− dL_dsh[14] = dRGBdsh14 * dL_dRGB;
102− dL_dsh[15] = dRGBdsh15 * dL_dRGB;
172+// Reverse traversal per pixel; accumulates splat gradients with standard
173+// 5-step full-warp butterfly reductions (45 shuffles) and 9 lane-0 atomics.
174+__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
175+ RasterBackwardKernel(const int2* __restrict__ ranges,
176+ const int32_t* __restrict__ point_list,
177+ const Splat* __restrict__ splats, Frame frame,
178+ const float* __restrict__ final_t,
179+ const int32_t* __restrict__ n_contrib,
180+ const float* __restrict__ dl_dimage,
181+ float* __restrict__ splat_grads) {
182+ using L = RasterLayout<1>;
183+ constexpr int kRasterThreads = L::kThreads;
184+ constexpr int kWarps = L::kWarps;
185+ const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
186+ const int tid = threadIdx.x;
187+ const int warp = tid / 32;
188+ const int lane = tid % 32;
189+ const RasterPixels px = RasterPixels::For<1>(tid);
190+ const float fx = static_cast<float>(px.x);
191+ const float fy = static_cast<float>(px.y0);
192+ const int2 range = ranges[tile];
103193
104− dRGBdx += (SH_C3[0] * sh[9] * 3.f * 2.f * xy + SH_C3[1] * sh[10] * yz +
105− SH_C3[2] * sh[11] * -2.f * xy +
106− SH_C3[3] * sh[12] * -3.f * 2.f * xz +
107− SH_C3[4] * sh[13] * (-3.f * xx + 4.f * zz - yy) +
108− SH_C3[5] * sh[14] * 2.f * xz +
109− SH_C3[6] * sh[15] * 3.f * (xx - yy));
194+ const bool inside = px.x < frame.width && px.y0 < frame.height;
195+ const int pixels = frame.width * frame.height;
196+ const int p = px.y0 * frame.width + px.x;
197+ float t = inside ? final_t[p] : 0.0f;
198+ const float3 dpix =
199+ inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
200+ dl_dimage[2 * pixels + p]}
201+ : float3{0.0f, 0.0f, 0.0f};
202+ float accum = 0.0f; // dot(color accumulated behind, dpix)
203+ const int last = inside ? n_contrib[p] : 0;
110204
111− dRGBdy +=
112− (SH_C3[0] * sh[9] * 3.f * (xx - yy) + SH_C3[1] * sh[10] * xz +
113− SH_C3[2] * sh[11] * (-3.f * yy + 4.f * zz - xx) +
114− SH_C3[3] * sh[12] * -3.f * 2.f * yz +
115− SH_C3[4] * sh[13] * -2.f * xy + SH_C3[5] * sh[14] * -2.f * yz +
116− SH_C3[6] * sh[15] * -3.f * 2.f * xy);
205+ int warp_max_last = last;
206+#pragma unroll
207+ for (int offset = 16; offset > 0; offset /= 2) {
208+ warp_max_last = max(
209+ warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
210+ }
211+ __shared__ int shared_max_last[kWarps];
212+ if (lane == 0) shared_max_last[warp] = warp_max_last;
213+ __syncthreads();
214+ int block_max_last = 0;
215+#pragma unroll
216+ for (int w = 0; w < kWarps; ++w) {
217+ block_max_last = max(block_max_last, shared_max_last[w]);
218+ }
219+ if (block_max_last <= 0) return;
117220
118− dRGBdz += (SH_C3[1] * sh[10] * xy + SH_C3[2] * sh[11] * 4.f * 2.f * yz +
119− SH_C3[3] * sh[12] * 3.f * (2.f * zz - xx - yy) +
120− SH_C3[4] * sh[13] * 4.f * 2.f * xz +
121− SH_C3[5] * sh[14] * (xx - yy));
221+ __shared__ float4 batch_geom0[kRasterThreads];
222+ __shared__ float2 batch_geom1[kRasterThreads];
223+ __shared__ float4 batch_rgb_o[kRasterThreads];
224+ __shared__ uint32_t batch_id[kRasterThreads];
225+ __shared__ uint8_t batch_mask[kRasterThreads];
226+ const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
227+
228+ auto raw_conic = [](float4 g0, float2 g1) {
229+ return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
230+ kInvPower2Scale * g1.x};
231+ };
232+
233+ for (int end = block_max_last; end > 0; end -= kRasterThreads) {
234+ const int e = end - 1 - tid;
235+ if (e >= 0) {
236+ const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
237+ const float4 p0 = __ldg(&splats[id].p0);
238+ const float4 p1 = __ldg(&splats[id].p1);
239+ const float4 p2 = __ldg(&splats[id].p2);
240+ const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
241+ batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
242+ batch_geom1[tid] = {c2.z, p2.w};
243+ batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
244+ batch_id[tid] = id;
245+ batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
246+ } else {
247+ batch_mask[tid] = 0u;
248+ }
249+ __syncthreads();
250+ const int count = min(kRasterThreads, end);
251+#pragma unroll
252+ for (int c = 0; c < kRasterThreads; c += 32) {
253+ if (c < count && end - 32 - c < warp_max_last) {
254+ uint32_t mask = __ballot_sync(
255+ warp::kFullMask,
256+ c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
257+ while (mask != 0u) {
258+ const int bit = __ffs(mask) - 1;
259+ mask &= mask - 1u;
260+ const int j = c + bit;
261+ const int entry = end - 1 - j;
262+ float gr[kSplatGradFloats] = {};
263+ bool hit = false;
264+ if (entry < last) {
265+ const float4 g0 = batch_geom0[j];
266+ const float2 g1 = batch_geom1[j];
267+ const float dx = g0.x - fx;
268+ const float dy = g0.y - fy;
269+ const float power =
270+ g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
271+ if (power <= 0.0f && power >= g1.y) {
272+ const float4 ro = batch_rgb_o[j];
273+ const float g_val = Exp2Approx(power);
274+ const float alpha = fminf(kMaxAlpha, ro.w * g_val);
275+ if (alpha >= kMinAlpha) {
276+ hit = true;
277+ t = __fdividef(t, 1.0f - alpha);
278+ const float w = alpha * t;
279+ const float color_dot =
280+ ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
281+ const float diff = color_dot - accum;
282+ const float dl_dalpha = t * diff;
283+ accum += alpha * diff;
284+ const float ga = g_val * dl_dalpha;
285+ const float dl_dpower = ro.w * ga;
286+ const float p_dx = dx * dl_dpower;
287+ const float p_dy = dy * dl_dpower;
288+ const float3 cn = raw_conic(g0, g1);
289+ gr[0] = -(cn.x * p_dx + cn.y * p_dy);
290+ gr[1] = -(cn.y * p_dx + cn.z * p_dy);
291+ gr[2] = -0.5f * dx * p_dx;
292+ gr[3] = -dx * p_dy;
293+ gr[4] = -0.5f * dy * p_dy;
294+ gr[5] = ga;
295+ gr[6] = w * dpix.x;
296+ gr[7] = w * dpix.y;
297+ gr[8] = w * dpix.z;
122298⇥ }
123299⇥ }
124300⇥ }
125−
126− // The view direction is an input to the computation. View direction
127− // is influenced by the Gaussian's mean, so SHs gradients
128− // must propagate back into 3D position.
129− glm::vec3 dL_ddir(glm::dot(dRGBdx, dL_dRGB), glm::dot(dRGBdy, dL_dRGB),
130− glm::dot(dRGBdz, dL_dRGB));
131−
132− // Account for normalization of direction
133− float3 dL_dmean = dnormvdv(float3{dir_orig.x, dir_orig.y, dir_orig.z},
134− float3{dL_ddir.x, dL_ddir.y, dL_ddir.z});
135−
136− // Gradients of loss w.r.t. Gaussian means, but only the portion
137− // that is caused because the mean affects the view-dependent color.
138− // Additional mean gradient is accumulated in below methods.
139− dL_dmeans[idx] += glm::vec3(dL_dmean.x, dL_dmean.y, dL_dmean.z);
301+ if (!__any_sync(warp::kFullMask, hit)) continue;
302+ warp::ButterflyReduce(gr);
303+ if (lane == 0) {
304+ float* dst = splat_grads + kSplatGradFloats * batch_id[j];
305+#pragma unroll
306+ for (int q = 0; q < kSplatGradFloats; ++q) {
307+ atomicAdd(dst + q, gr[q]);
140308⇥ }
141−
142−// Backward version of INVERSE 2D covariance matrix computation
143−// (due to length launched as separate kernel before other
144−// backward steps contained in preprocess)
145−__global__ void computeCov2DCUDA(
146− int P, const float3* means, const int* radii, const float* cov3Ds,
147− const float4* __restrict__ conic_opacity, const float h_x, float h_y,
148− const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
149− const bool mip_splatting_enabled, const float* view_matrix,
150− const float* dL_dconics, float3* dL_dmeans, float* dL_dcov,
151− float* dL_dopacity) {
152− auto idx = cg::this_grid().thread_rank();
153− if (idx >= P || !(radii[idx] > 0)) return;
154−
155− // Reading location of 3D covariance for this Gaussian
156− const float* cov3D = cov3Ds + 6 * idx;
157−
158− // Fetch gradients, recompute 2D covariance and relevant
159− // intermediate forward results needed in the backward.
160− float3 mean = means[idx];
161− float3 dL_dconic = {dL_dconics[4 * idx], dL_dconics[4 * idx + 1],
162− dL_dconics[4 * idx + 3]};
163− float3 t = transformPoint4x3(mean, view_matrix);
164−
165− const float limx = 1.3f * tan_fovx;
166− const float limy = 1.3f * tan_fovy;
167− const float txtz = t.x / t.z;
168− const float tytz = t.y / t.z;
169− t.x = min(limx, max(-limx, txtz)) * t.z;
170− t.y = min(limy, max(-limy, tytz)) * t.z;
171−
172− const float x_grad_mul = txtz < -limx || txtz > limx ? 0 : 1;
173− const float y_grad_mul = tytz < -limy || tytz > limy ? 0 : 1;
174−
175− glm::mat3 J = glm::mat3(h_x / t.z, 0.0f, -(h_x * t.x) / (t.z * t.z), 0.0f,
176− h_y / t.z, -(h_y * t.y) / (t.z * t.z), 0, 0, 0);
177−
178− glm::mat3 W = glm::mat3(view_matrix[0], view_matrix[4], view_matrix[8],
179− view_matrix[1], view_matrix[5], view_matrix[9],
180− view_matrix[2], view_matrix[6], view_matrix[10]);
181−
182− glm::mat3 Vrk = glm::mat3(cov3D[0], cov3D[1], cov3D[2], cov3D[1], cov3D[3],
183− cov3D[4], cov3D[2], cov3D[4], cov3D[5]);
184−
185− glm::mat3 T = W * J;
186−
187− glm::mat3 cov2D = glm::transpose(T) * glm::transpose(Vrk) * T;
188−
189− float kernel_size_to_use;
190− if (mip_splatting_enabled)
191− kernel_size_to_use = low_pass_2d_kernel_size;
192− else
193− kernel_size_to_use = 0.3f;
194−
195− float a = cov2D[0][0] += kernel_size_to_use;
196− float b = cov2D[0][1];
197− float c = cov2D[1][1] += kernel_size_to_use;
198−
199− float denom = a * c - b * b;
200− float dL_da = 0, dL_db = 0, dL_dc = 0;
201− float denom2inv = 1.0f / ((denom * denom) + 0.0000001f);
202−
203− if (denom2inv != 0) {
204− // Gradients of loss w.r.t. entries of 2D covariance matrix,
205− // given gradients of loss w.r.t. conic matrix (inverse covariance matrix).
206− // e.g., dL / da = dL / d_conic_a * d_conic_a / d_a
207− dL_da = denom2inv * (-c * c * dL_dconic.x + 2 * b * c * dL_dconic.y +
208− (denom - a * c) * dL_dconic.z);
209− dL_dc = denom2inv * (-a * a * dL_dconic.z + 2 * a * b * dL_dconic.y +
210− (denom - a * c) * dL_dconic.x);
211− dL_db = denom2inv * 2 *
212− (b * c * dL_dconic.x - (denom + 2 * b * b) * dL_dconic.y +
213− a * b * dL_dconic.z);
214−
215− if (mip_splatting_enabled) {
216− const float det_0 =
217− max(1e-6, cov2D[0][0] * cov2D[1][1] - cov2D[0][1] * cov2D[0][1]);
218− const float det_1 = max(1e-6, (cov2D[0][0] + kernel_size_to_use) *
219− (cov2D[1][1] + kernel_size_to_use) -
220− cov2D[0][1] * cov2D[0][1]);
221− // sqrt here
222− const float coef = sqrt(det_0 / (det_1 + 1e-6) + 1e-6);
223−
224− // update the gradient of alpha and save the gradient of dalpha_dcoef
225− // we need opacity as input
226− // new_opacity = coef * opacity
227− // if we know the new opacity, we can derive original opacity and then
228− // dalpha_dcoef = dopacity * opacity
229− const float opacity = conic_opacity[idx].w / (coef + 1e-6);
230− const float dL_dcoef = dL_dopacity[idx] * opacity;
231− const float dL_dsqrtcoef = dL_dcoef * 0.5 * 1. / (coef + 1e-6);
232− const float dL_ddet0 = dL_dsqrtcoef / (det_1 + 1e-6);
233− const float dL_ddet1 =
234− dL_dsqrtcoef * det_0 * (-1.f / (det_1 * det_1 + 1e-6));
235− // TODO gradient is zero if det_0 or det_1 < 0
236− const float dcoef_da = dL_ddet0 * cov2D[1][1] +
237− dL_ddet1 * (cov2D[1][1] + kernel_size_to_use);
238− const float dcoef_db =
239− dL_ddet0 * (-2. * cov2D[0][1]) + dL_ddet1 * (-2. * cov2D[0][1]);
240− const float dcoef_dc = dL_ddet0 * cov2D[0][0] +
241− dL_ddet1 * (cov2D[0][0] + kernel_size_to_use);
242−
243− if (det_0 <= 1e-6 || det_1 <= 1e-6) {
244− dL_dopacity[idx] = 0;
245− } else {
246− // Gradiends of alpha respect to conv due to low pass filter
247− dL_da += dcoef_da;
248− dL_dc += dcoef_dc;
249− dL_db += dcoef_db;
250−
251− // update dL_dopacity
252− dL_dopacity[idx] = dL_dopacity[idx] * coef;
309+ }
310+ }
311+ }
312+ }
313+ __syncthreads();
253314⇥ }
254315⇥}
255316
256− // Gradients of loss L w.r.t. each 3D covariance matrix (Vrk) entry,
317+// ---------------------------------------------------------------------------
257− // given gradients w.r.t. 2D covariance matrix (diagonal).
318+// Preprocess backward: splat gradients -> parameter gradients.
258− // cov2D = transpose(T) * transpose(Vrk) * T;
259− dL_dcov[6 * idx + 0] =
260− (T[0][0] * T[0][0] * dL_da + T[0][0] * T[1][0] * dL_db +
261− T[1][0] * T[1][0] * dL_dc);
262− dL_dcov[6 * idx + 3] =
263− (T[0][1] * T[0][1] * dL_da + T[0][1] * T[1][1] * dL_db +
264− T[1][1] * T[1][1] * dL_dc);
265− dL_dcov[6 * idx + 5] =
266− (T[0][2] * T[0][2] * dL_da + T[0][2] * T[1][2] * dL_db +
267− T[1][2] * T[1][2] * dL_dc);
319+// Parameter gradient buffers are bulk-zeroed beforehand via `cudaMemsetAsync`;
320+// culled or zero-gradient Gaussians return early.
268321
269− // Gradients of loss L w.r.t. each 3D covariance matrix (Vrk) entry,
270− // given gradients w.r.t. 2D covariance matrix (off-diagonal).
271− // Off-diagonal elements appear twice --> double the gradient.
272− // cov2D = transpose(T) * transpose(Vrk) * T;
273− dL_dcov[6 * idx + 1] = 2 * T[0][0] * T[0][1] * dL_da +
274− (T[0][0] * T[1][1] + T[0][1] * T[1][0]) * dL_db +
275− 2 * T[1][0] * T[1][1] * dL_dc;
276− dL_dcov[6 * idx + 2] = 2 * T[0][0] * T[0][2] * dL_da +
277− (T[0][0] * T[1][2] + T[0][2] * T[1][0]) * dL_db +
278− 2 * T[1][0] * T[1][2] * dL_dc;
279− dL_dcov[6 * idx + 4] = 2 * T[0][2] * T[0][1] * dL_da +
280− (T[0][1] * T[1][2] + T[0][2] * T[1][1]) * dL_db +
281− 2 * T[1][1] * T[1][2] * dL_dc;
282− } else {
322+template <typename ShT>
323+__global__ void __launch_bounds__(kGaussianThreads)
324+ PreprocessBackwardKernel(GaussianParams params, Frame frame,
325+ const Splat* __restrict__ splats,
326+ const float* __restrict__ splat_grads,
327+ GaussianGrads grads) {
328+ const int i = blockIdx.x * blockDim.x + threadIdx.x;
329+ if (i >= params.n || splats[i].radius() == 0) return;
330+ float g[kSplatGradFloats];
331+ bool nonzero = false;
332+#pragma unroll
283− for (int i = 0; i < 6; i++) dL_dcov[6 * idx + i] = 0;
333+ for (int q = 0; q < kSplatGradFloats; ++q) {
334+ g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
335+ nonzero = nonzero || g[q] != 0.0f;
336+ }
337+ if (!nonzero) return;
338+ ProjectionGrad pg;
339+ pg.xy = {g[0], g[1]};
340+ pg.conic = {g[2], g[3], g[4]};
341+ pg.opacity = g[5];
342+ pg.rgb = {g[6], g[7], g[8]};
343+ float3 dmean, dscale;
344+ float4 dquat;
345+ ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
346+ ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
347+ grads.means[3 * i] = dmean.x;
348+ grads.means[3 * i + 1] = dmean.y;
349+ grads.means[3 * i + 2] = dmean.z;
350+ grads.scales[3 * i] = dscale.x;
351+ grads.scales[3 * i + 1] = dscale.y;
352+ grads.scales[3 * i + 2] = dscale.z;
353+ reinterpret_cast<float4*>(grads.quats)[i] = dquat;
354+ grads.opacities[i] = g[5];
284355⇥}
285356
286− // Gradients of loss w.r.t. upper 2x3 portion of intermediate matrix T
287− // cov2D = transpose(T) * transpose(Vrk) * T;
288− float dL_dT00 =
289− 2 * (T[0][0] * Vrk[0][0] + T[0][1] * Vrk[0][1] + T[0][2] * Vrk[0][2]) *
290− dL_da +
291− (T[1][0] * Vrk[0][0] + T[1][1] * Vrk[0][1] + T[1][2] * Vrk[0][2]) * dL_db;
292− float dL_dT01 =
293− 2 * (T[0][0] * Vrk[1][0] + T[0][1] * Vrk[1][1] + T[0][2] * Vrk[1][2]) *
294− dL_da +
295− (T[1][0] * Vrk[1][0] + T[1][1] * Vrk[1][1] + T[1][2] * Vrk[1][2]) * dL_db;
296− float dL_dT02 =
297− 2 * (T[0][0] * Vrk[2][0] + T[0][1] * Vrk[2][1] + T[0][2] * Vrk[2][2]) *
298− dL_da +
299− (T[1][0] * Vrk[2][0] + T[1][1] * Vrk[2][1] + T[1][2] * Vrk[2][2]) * dL_db;
300− float dL_dT10 =
301− 2 * (T[1][0] * Vrk[0][0] + T[1][1] * Vrk[0][1] + T[1][2] * Vrk[0][2]) *
302− dL_dc +
303− (T[0][0] * Vrk[0][0] + T[0][1] * Vrk[0][1] + T[0][2] * Vrk[0][2]) * dL_db;
304− float dL_dT11 =
305− 2 * (T[1][0] * Vrk[1][0] + T[1][1] * Vrk[1][1] + T[1][2] * Vrk[1][2]) *
306− dL_dc +
307− (T[0][0] * Vrk[1][0] + T[0][1] * Vrk[1][1] + T[0][2] * Vrk[1][2]) * dL_db;
308− float dL_dT12 =
309− 2 * (T[1][0] * Vrk[2][0] + T[1][1] * Vrk[2][1] + T[1][2] * Vrk[2][2]) *
310− dL_dc +
311− (T[0][0] * Vrk[2][0] + T[0][1] * Vrk[2][1] + T[0][2] * Vrk[2][2]) * dL_db;
312−
313− // Gradients of loss w.r.t. upper 3x2 non-zero entries of Jacobian matrix
314− // T = W * J
315− float dL_dJ00 = W[0][0] * dL_dT00 + W[0][1] * dL_dT01 + W[0][2] * dL_dT02;
316− float dL_dJ02 = W[2][0] * dL_dT00 + W[2][1] * dL_dT01 + W[2][2] * dL_dT02;
317− float dL_dJ11 = W[1][0] * dL_dT10 + W[1][1] * dL_dT11 + W[1][2] * dL_dT12;
318− float dL_dJ12 = W[2][0] * dL_dT10 + W[2][1] * dL_dT11 + W[2][2] * dL_dT12;
319−
320− float tz = 1.f / t.z;
321− float tz2 = tz * tz;
322− float tz3 = tz2 * tz;
323−
324− // Gradients of loss w.r.t. transformed Gaussian mean t
325− float dL_dtx = x_grad_mul * -h_x * tz2 * dL_dJ02;
326− float dL_dty = y_grad_mul * -h_y * tz2 * dL_dJ12;
327− float dL_dtz = -h_x * tz2 * dL_dJ00 - h_y * tz2 * dL_dJ11 +
328− (2 * h_x * t.x) * tz3 * dL_dJ02 +
329− (2 * h_y * t.y) * tz3 * dL_dJ12;
330−
331− // Account for transformation of mean to t
332− // t = transformPoint4x3(mean, view_matrix);
333− float3 dL_dmean =
334− transformVec4x3Transpose({dL_dtx, dL_dty, dL_dtz}, view_matrix);
335−
357+template <typename ShT>
358+cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
359+ const Frame& frame, const Splat* splats,
360+ const float* splat_grads,
361+ const GaussianGrads& grads,
362+ cudaStream_t stream) {
363+ const int n = params.n;
364+ // As Inria's PyTorch wrapper (`torch::zeros`): bulk-zero every parameter
336− // Gradients of loss w.r.t. Gaussian means, but only the portion
365+ // gradient up front; culled Gaussians then exit early in the kernel.
337− // that is caused because the mean affects the covariance matrix.
338− // Additional mean gradient is accumulated in BACKWARD::preprocess.
339− dL_dmeans[idx] = dL_dmean;
366+ if (auto e = cudaMemsetAsync(grads.means, 0, sizeof(float) * 3 * n, stream);
367+ e != cudaSuccess)
368+ return e;
369+ if (auto e = cudaMemsetAsync(grads.scales, 0, sizeof(float) * 3 * n, stream);
370+ e != cudaSuccess)
371+ return e;
372+ if (auto e = cudaMemsetAsync(grads.quats, 0, sizeof(float) * 4 * n, stream);
373+ e != cudaSuccess)
374+ return e;
375+ if (auto e = cudaMemsetAsync(grads.opacities, 0, sizeof(float) * n, stream);
376+ e != cudaSuccess)
377+ return e;
378+ if (auto e =
379+ cudaMemsetAsync(grads.sh, 0, sizeof(ShT) * 3 * kShCoeffs * n, stream);
380+ e != cudaSuccess)
381+ return e;
382+ PreprocessBackwardKernel<ShT>
383+ <<<GaussianBlocks(n), kGaussianThreads, 0, stream>>>(
384+ params, frame, splats, splat_grads, grads);
385+ return cudaSuccess;
340386 }
341387
342−// Backward pass for the conversion of scale and rotation to a
343−// 3D covariance matrix for each Gaussian.
344−__device__ void computeCov3D(int idx, const glm::vec3 scale, float mod,
388+template __global__ void PreprocessBackwardKernel<float>(
389+ GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
390+ GaussianGrads);
391+template __global__ void PreprocessBackwardKernel<ShBf16>(
392+ GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
393+ GaussianGrads);
394+template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
345− const glm::vec4 rot, const float* dL_dcov3Ds,
395+ const Frame&, const Splat*,
346− glm::vec3* dL_dscales, glm::vec4* dL_drots) {
347− // Recompute (intermediate) results for the 3D covariance computation.
348− glm::vec4 q = rot; // / glm::length(rot);
349− float r = q.x;
350− float x = q.y;
351− float y = q.z;
352− float z = q.w;
396+ const float*,
397+ const GaussianGrads&,
398+ cudaStream_t);
399+template cudaError_t LaunchPreprocessBackward<ShBf16>(
400+ const GaussianParams&, const Frame&, const Splat*, const float*,
401+ const GaussianGrads&, cudaStream_t);
353402
354− glm::mat3 R = glm::mat3(1.f - 2.f * (y * y + z * z), 2.f * (x * y - r * z),
355− 2.f * (x * z + r * y), 2.f * (x * y + r * z),
356− 1.f - 2.f * (x * x + z * z), 2.f * (y * z - r * x),
357− 2.f * (x * z - r * y), 2.f * (y * z + r * x),
358− 1.f - 2.f * (x * x + y * y));
359−
360− glm::mat3 S = glm::mat3(1.0f);
361−
362− glm::vec3 s = mod * scale;
363− S[0][0] = s.x;
364− S[1][1] = s.y;
365− S[2][2] = s.z;
366−
367− glm::mat3 M = S * R;
368−
369− const float* dL_dcov3D = dL_dcov3Ds + 6 * idx;
370−
371− glm::vec3 dunc(dL_dcov3D[0], dL_dcov3D[3], dL_dcov3D[5]);
372− glm::vec3 ounc = 0.5f * glm::vec3(dL_dcov3D[1], dL_dcov3D[2], dL_dcov3D[4]);
373−
374− // Convert per-element covariance loss gradients to matrix form
375− glm::mat3 dL_dSigma =
376− glm::mat3(dL_dcov3D[0], 0.5f * dL_dcov3D[1], 0.5f * dL_dcov3D[2],
377− 0.5f * dL_dcov3D[1], dL_dcov3D[3], 0.5f * dL_dcov3D[4],
378− 0.5f * dL_dcov3D[2], 0.5f * dL_dcov3D[4], dL_dcov3D[5]);
379−
380− // Compute loss gradient w.r.t. matrix M
381− // dSigma_dM = 2 * M
382− glm::mat3 dL_dM = 2.0f * M * dL_dSigma;
383−
384− glm::mat3 Rt = glm::transpose(R);
385− glm::mat3 dL_dMt = glm::transpose(dL_dM);
386−
387− // Gradients of loss w.r.t. scale
388− glm::vec3* dL_dscale = dL_dscales + idx;
389− dL_dscale->x = glm::dot(Rt[0], dL_dMt[0]);
390− dL_dscale->y = glm::dot(Rt[1], dL_dMt[1]);
391− dL_dscale->z = glm::dot(Rt[2], dL_dMt[2]);
392−
393− dL_dMt[0] *= s.x;
394− dL_dMt[1] *= s.y;
395− dL_dMt[2] *= s.z;
396−
397− // Gradients of loss w.r.t. normalized quaternion
398− glm::vec4 dL_dq;
399− dL_dq.x = 2 * z * (dL_dMt[0][1] - dL_dMt[1][0]) +
400− 2 * y * (dL_dMt[2][0] - dL_dMt[0][2]) +
401− 2 * x * (dL_dMt[1][2] - dL_dMt[2][1]);
402− dL_dq.y = 2 * y * (dL_dMt[1][0] + dL_dMt[0][1]) +
403− 2 * z * (dL_dMt[2][0] + dL_dMt[0][2]) +
404− 2 * r * (dL_dMt[1][2] - dL_dMt[2][1]) -
405− 4 * x * (dL_dMt[2][2] + dL_dMt[1][1]);
406− dL_dq.z = 2 * x * (dL_dMt[1][0] + dL_dMt[0][1]) +
407− 2 * r * (dL_dMt[2][0] - dL_dMt[0][2]) +
408− 2 * z * (dL_dMt[1][2] + dL_dMt[2][1]) -
409− 4 * y * (dL_dMt[2][2] + dL_dMt[0][0]);
410− dL_dq.w = 2 * r * (dL_dMt[0][1] - dL_dMt[1][0]) +
411− 2 * x * (dL_dMt[2][0] + dL_dMt[0][2]) +
412− 2 * y * (dL_dMt[1][2] + dL_dMt[2][1]) -
413− 4 * z * (dL_dMt[1][1] + dL_dMt[0][0]);
414−
415− // Gradients of loss w.r.t. unnormalized quaternion
416− float4* dL_drot = (float4*)(dL_drots + idx);
417− *dL_drot = float4{dL_dq.x, dL_dq.y, dL_dq.z,
418− dL_dq.w}; // dnormvdv(float4{ rot.x, rot.y, rot.z, rot.w },
419− // float4{ dL_dq.x, dL_dq.y, dL_dq.z, dL_dq.w });
420−}
421−
422−// Backward pass of the preprocessing steps, except
423−// for the covariance computation and inversion
424−// (those are handled by a previous kernel call)
425−template <int C>
426−__global__ void preprocessCUDA(
427− int P, int D, int M, const float3* means, const int* radii,
428− const float* shs, const bool* clamped, const glm::vec3* scales,
429− const glm::vec4* rotations, const float scale_modifier, const float* proj,
430− const glm::vec3* campos, const float3* dL_dmean2D, glm::vec3* dL_dmeans,
431− float* dL_dcolor, float* dL_dcov3D, float* dL_dsh, glm::vec3* dL_dscale,
432− glm::vec4* dL_drot) {
433− auto idx = cg::this_grid().thread_rank();
434− if (idx >= P || !(radii[idx] > 0)) return;
435−
436− float3 m = means[idx];
437−
438− // Taking care of gradients from the screenspace points
439− float4 m_hom = transformPoint4x4(m, proj);
440− float m_w = 1.0f / (m_hom.w + 0.0000001f);
441−
442− // Compute loss gradient w.r.t. 3D means due to gradients of 2D means
443− // from rendering procedure
444− glm::vec3 dL_dmean;
445− float mul1 =
446− (proj[0] * m.x + proj[4] * m.y + proj[8] * m.z + proj[12]) * m_w * m_w;
447− float mul2 =
448− (proj[1] * m.x + proj[5] * m.y + proj[9] * m.z + proj[13]) * m_w * m_w;
449− dL_dmean.x = (proj[0] * m_w - proj[3] * mul1) * dL_dmean2D[idx].x +
450− (proj[1] * m_w - proj[3] * mul2) * dL_dmean2D[idx].y;
451− dL_dmean.y = (proj[4] * m_w - proj[7] * mul1) * dL_dmean2D[idx].x +
452− (proj[5] * m_w - proj[7] * mul2) * dL_dmean2D[idx].y;
453− dL_dmean.z = (proj[8] * m_w - proj[11] * mul1) * dL_dmean2D[idx].x +
454− (proj[9] * m_w - proj[11] * mul2) * dL_dmean2D[idx].y;
455−
456− // That's the second part of the mean gradient. Previous computation
457− // of cov2D and following SH conversion also affects it.
458− dL_dmeans[idx] += dL_dmean;
459−
460− // Compute gradient updates due to computing colors from SHs
461− if (shs)
462− computeColorFromSH(idx, D, M, (glm::vec3*)means, *campos, shs, clamped,
463− (glm::vec3*)dL_dcolor, (glm::vec3*)dL_dmeans,
464− (glm::vec3*)dL_dsh);
465−
466− // Compute gradient updates due to computing covariance from scale/rotation
467− if (scales)
468− computeCov3D(idx, scales[idx], scale_modifier, rotations[idx], dL_dcov3D,
469− dL_dscale, dL_drot);
470−}
471−
472−// Backward version of the rendering procedure.
473−template <uint32_t C>
474−__global__ void __launch_bounds__(BLOCK_X* BLOCK_Y) renderCUDA(
475− const uint2* __restrict__ ranges, const uint32_t* __restrict__ point_list,
476− int W, int H, const float* __restrict__ bg_color,
477− const float2* __restrict__ points_xy_image,
478− const float4* __restrict__ conic_opacity, const float* __restrict__ colors,
479− const float* __restrict__ final_Ts, const uint32_t* __restrict__ n_contrib,
480− const float* __restrict__ dL_dpixels, float3* __restrict__ dL_dmean2D,
481− float4* __restrict__ dL_dconic2D, float* __restrict__ dL_dopacity,
482− float* __restrict__ dL_dcolors) {
483− // We rasterize again. Compute necessary block info.
484− auto block = cg::this_thread_block();
485− const uint32_t horizontal_blocks = (W + BLOCK_X - 1) / BLOCK_X;
486− const uint2 pix_min = {block.group_index().x * BLOCK_X,
487− block.group_index().y * BLOCK_Y};
488− const uint2 pix_max = {min(pix_min.x + BLOCK_X, W),
489− min(pix_min.y + BLOCK_Y, H)};
490− const uint2 pix = {pix_min.x + block.thread_index().x,
491− pix_min.y + block.thread_index().y};
492− const uint32_t pix_id = W * pix.y + pix.x;
493− const float2 pixf = {(float)pix.x, (float)pix.y};
494−
495− const bool inside = pix.x < W && pix.y < H;
496− const uint2 range =
497− ranges[block.group_index().y * horizontal_blocks + block.group_index().x];
498−
499− const int rounds = ((range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE);
500−
501− bool done = !inside;
502− int toDo = range.y - range.x;
503−
504− __shared__ int collected_id[BLOCK_SIZE];
505− __shared__ float2 collected_xy[BLOCK_SIZE];
506− __shared__ float4 collected_conic_opacity[BLOCK_SIZE];
507− __shared__ float collected_colors[C * BLOCK_SIZE];
508−
509− // In the forward, we stored the final value for T, the
510− // product of all (1 - alpha) factors.
511− const float T_final = inside ? final_Ts[pix_id] : 0;
512− float T = T_final;
513−
514− // We start from the back. The ID of the last contributing
515− // Gaussian is known from each pixel from the forward.
516− uint32_t contributor = toDo;
517− const int last_contributor = inside ? n_contrib[pix_id] : 0;
518−
519− float accum_rec[C] = {0};
520− float dL_dpixel[C];
521− if (inside)
522− for (int i = 0; i < C; i++) dL_dpixel[i] = dL_dpixels[i * H * W + pix_id];
523−
524− float last_alpha = 0;
525− float last_color[C] = {0};
526−
527− // Gradient of pixel coordinate w.r.t. normalized
528− // screen-space viewport corrdinates (-1 to 1)
529− const float ddelx_dx = 0.5 * W;
530− const float ddely_dy = 0.5 * H;
531−
532− // Traverse all Gaussians
533− for (int i = 0; i < rounds; i++, toDo -= BLOCK_SIZE) {
534− // Load auxiliary data into shared memory, start in the BACK
535− // and load them in revers order.
536− block.sync();
537− const int progress = i * BLOCK_SIZE + block.thread_rank();
538− if (range.x + progress < range.y) {
539− const int coll_id = point_list[range.y - progress - 1];
540− collected_id[block.thread_rank()] = coll_id;
541− collected_xy[block.thread_rank()] = points_xy_image[coll_id];
542− collected_conic_opacity[block.thread_rank()] = conic_opacity[coll_id];
543− for (int i = 0; i < C; i++)
544− collected_colors[i * BLOCK_SIZE + block.thread_rank()] =
545− colors[coll_id * C + i];
546− }
547− block.sync();
548−
549− // Iterate over Gaussians
550− for (int j = 0; !done && j < min(BLOCK_SIZE, toDo); j++) {
551− // Keep track of current Gaussian ID. Skip, if this one
552− // is behind the last contributor for this pixel.
553− contributor--;
554− if (contributor >= last_contributor) continue;
555−
556− // Compute blending values, as before.
557− const float2 xy = collected_xy[j];
558− const float2 d = {xy.x - pixf.x, xy.y - pixf.y};
559− const float4 con_o = collected_conic_opacity[j];
560− const float power = -0.5f * (con_o.x * d.x * d.x + con_o.z * d.y * d.y) -
561− con_o.y * d.x * d.y;
562− if (power > 0.0f) continue;
563−
564− const float G = exp(power);
565− const float alpha = min(0.99f, con_o.w * G);
566− if (alpha < 1.0f / 255.0f) continue;
567−
568− T = T / (1.f - alpha);
569− const float dchannel_dcolor = alpha * T;
570−
571− // Propagate gradients to per-Gaussian colors and keep
572− // gradients w.r.t. alpha (blending factor for a Gaussian/pixel
573− // pair).
574− float dL_dalpha = 0.0f;
575− const int global_id = collected_id[j];
576− for (int ch = 0; ch < C; ch++) {
577− const float c = collected_colors[ch * BLOCK_SIZE + j];
578− // Update last color (to be used in the next iteration)
579− accum_rec[ch] =
580− last_alpha * last_color[ch] + (1.f - last_alpha) * accum_rec[ch];
581− last_color[ch] = c;
582−
583− const float dL_dchannel = dL_dpixel[ch];
584− dL_dalpha += (c - accum_rec[ch]) * dL_dchannel;
585− // Update the gradients w.r.t. color of the Gaussian.
586− // Atomic, since this pixel is just one of potentially
587− // many that were affected by this Gaussian.
588− atomicAdd(&(dL_dcolors[global_id * C + ch]),
589− dchannel_dcolor * dL_dchannel);
590− }
591− dL_dalpha *= T;
592− // Update last alpha (to be used in the next iteration)
593− last_alpha = alpha;
594−
595− // Account for fact that alpha also influences how much of
596− // the background color is added if nothing left to blend
597− float bg_dot_dpixel = 0;
598− for (int i = 0; i < C; i++) bg_dot_dpixel += bg_color[i] * dL_dpixel[i];
599− dL_dalpha += (-T_final / (1.f - alpha)) * bg_dot_dpixel;
600−
601− // Helpful reusable temporary variables
602− const float dL_dG = con_o.w * dL_dalpha;
603− const float gdx = G * d.x;
604− const float gdy = G * d.y;
605− const float dG_ddelx = -gdx * con_o.x - gdy * con_o.y;
606− const float dG_ddely = -gdy * con_o.z - gdx * con_o.y;
607−
608− // Update gradients w.r.t. 2D mean position of the Gaussian
609− atomicAdd(&dL_dmean2D[global_id].x, dL_dG * dG_ddelx * ddelx_dx);
610− atomicAdd(&dL_dmean2D[global_id].y, dL_dG * dG_ddely * ddely_dy);
611−
612− // Update gradients w.r.t. 2D covariance (2x2 matrix, symmetric)
613− atomicAdd(&dL_dconic2D[global_id].x, -0.5f * gdx * d.x * dL_dG);
614− atomicAdd(&dL_dconic2D[global_id].y, -0.5f * gdx * d.y * dL_dG);
615− atomicAdd(&dL_dconic2D[global_id].w, -0.5f * gdy * d.y * dL_dG);
616−
617− // Update gradients w.r.t. opacity of the Gaussian
618− atomicAdd(&(dL_dopacity[global_id]), G * dL_dalpha);
619− }
620− }
621−}
622−
623−void BACKWARD::preprocess(
624− int P, int D, int M, const float3* means3D, const int* radii,
625− const float* shs, const bool* clamped, const glm::vec3* scales,
626− const glm::vec4* rotations, const float scale_modifier, const float* cov3Ds,
627− const float4* conic_opacity, const float* viewmatrix,
628− const float* projmatrix, const float focal_x, float focal_y,
629− const float tan_fovx, float tan_fovy, const float low_pass_2d_kernel_size,
630− const bool mip_splatting_enabled, const glm::vec3* campos,
631− const float3* dL_dmean2D, const float* dL_dconic, glm::vec3* dL_dmean3D,
632− float* dL_dcolor, float* dL_dcov3D, float* dL_dsh, glm::vec3* dL_dscale,
633− glm::vec4* dL_drot, float* dL_dopacity) {
634− // Propagate gradients for the path of 2D conic matrix computation.
635− // Somewhat long, thus it is its own kernel rather than being part of
636− // "preprocess". When done, loss gradient w.r.t. 3D means has been
637− // modified and gradient w.r.t. 3D covariance matrix has been computed.
638− computeCov2DCUDA<<<(P + 255) / 256, 256>>>(
639− P, means3D, radii, cov3Ds, conic_opacity, focal_x, focal_y, tan_fovx,
640− tan_fovy, low_pass_2d_kernel_size, mip_splatting_enabled, viewmatrix,
641− dL_dconic, (float3*)dL_dmean3D, dL_dcov3D, dL_dopacity);
642−
643− // Propagate gradients for remaining steps: finish 3D mean gradients,
644− // propagate color gradients to SH (if desireD), propagate 3D covariance
645− // matrix gradients to scale and rotation.
646− preprocessCUDA<NUM_CHANNELS><<<(P + 255) / 256, 256>>>(
647− P, D, M, (float3*)means3D, radii, shs, clamped, (glm::vec3*)scales,
648− (glm::vec4*)rotations, scale_modifier, projmatrix, campos,
649− (float3*)dL_dmean2D, (glm::vec3*)dL_dmean3D, dL_dcolor, dL_dcov3D, dL_dsh,
650− dL_dscale, dL_drot);
651−}
652−
653−void BACKWARD::render(const dim3 grid, const dim3 block, const uint2* ranges,
654− const uint32_t* point_list, int W, int H,
655− const float* bg_color, const float2* means2D,
656− const float4* conic_opacity, const float* colors,
657− const float* final_Ts, const uint32_t* n_contrib,
658− const float* dL_dpixels, float3* dL_dmean2D,
659− float4* dL_dconic2D, float* dL_dopacity,
660− float* dL_dcolors) {
661− renderCUDA<NUM_CHANNELS>
662− <<<grid, block>>>(ranges, point_list, W, H, bg_color, means2D,
663− conic_opacity, colors, final_Ts, n_contrib, dL_dpixels,
664− dL_dmean2D, dL_dconic2D, dL_dopacity, dL_dcolors);
665−}
403+} // namespace gs
csrc/rasterizer_impl.h +43 −51
s00_baseline/csrc/rasterizer_impl.h
s01_literature/csrc/rasterizer_impl.h
↕ Expand 9 unchanged lines (lines 1–9)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#pragma once12#pragma once
1313
14#include <cstddef>
14#include <cuda_runtime_api.h>15#include <cstdint>
16#include <iterator>
1517
16#include <iostream>18#include "auxiliary.h"
17#include <vector>
1819
19#include "rasterizer.h"
20namespace gs {
2021
21namespace CudaRasterizer {
22template <typename T>
23static void obtain(char*& chunk, T*& ptr, std::size_t count,
24 std::size_t alignment) {
25 std::size_t offset =
26 (reinterpret_cast<std::uintptr_t>(chunk) + alignment - 1) &
22// Number of bits needed to represent tile ids in [0, num_tiles).
23inline int TileBits(int num_tiles) {
27 ~(alignment - 1);24 int bits = 1;
28 ptr = reinterpret_cast<T*>(offset);
29 chunk = reinterpret_cast<char*>(ptr + count);
25 while ((1 << bits) < num_tiles) ++bits;
26 return bits;
30}27}
3128
32struct GeometryState {29struct GatherCountsIter {
33 size_t scan_size;
34 float* depths;
35 char* scanning_space;
36 bool* clamped;
37 int* internal_radii;
38 float2* means2D;
39 float* cov3D;
40 float4* conic_opacity;
41 float* rgb;
42 uint32_t* point_offsets;30 const int32_t* counts;
43 uint32_t* tiles_touched;31 const int32_t* order;
44
45 static GeometryState fromChunk(char*& chunk, size_t P);
32 __host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
33 return counts[order[idx]];
34 }
35 __host__ __device__ __forceinline__ GatherCountsIter
36 operator+(ptrdiff_t offset) const {
37 return {counts, order + offset};
38 }
39 __host__ __device__ __forceinline__ int32_t operator*() const {
40 return counts[*order];
41 }
46};42};
4743
48struct ImageState {
49 uint2* ranges;
44template <Culling kCulling>
45__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
46 const int32_t* __restrict__ sorted_ids,
47 const int32_t* __restrict__ offsets, int n_vis,
48 Frame frame, uint16_t* __restrict__ keys,
50 uint32_t* n_contrib;49 int32_t* __restrict__ values);
51 float* accum_alpha;
5250
53 static ImageState fromChunk(char*& chunk, size_t N);
51__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
52 int num_pairs, int2* __restrict__ ranges);
53
54} // namespace gs
55
56namespace std {
57template <>
58struct iterator_traits<gs::GatherCountsIter> {
59 using difference_type = ptrdiff_t;
60 using value_type = int32_t;
61 using pointer = const int32_t*;
62 using reference = int32_t;
63 using iterator_category = random_access_iterator_tag;
54};64};
55
56struct BinningState {
57 size_t sorting_size;
58 uint64_t* point_list_keys_unsorted;
59 uint64_t* point_list_keys;
60 uint32_t* point_list_unsorted;
61 uint32_t* point_list;
62 char* list_sorting_space;
63
64 static BinningState fromChunk(char*& chunk, size_t P);
65};
66
67template <typename T>
68size_t required(size_t P) {
69 char* size = nullptr;
70 T::fromChunk(size, P);
71 return ((size_t)size) + 128;
72}
73}; // namespace CudaRasterizer65} // namespace std
↕ Expand 9 unchanged lines (lines 1–9)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #pragma once
1313
14+#include <cstddef>
14−#include <cuda_runtime_api.h>
15+#include <cstdint>
16+#include <iterator>
1517
16−#include <iostream>
18+#include "auxiliary.h"
17−#include <vector>
1819
19−#include "rasterizer.h"
20+namespace gs {
2021
21−namespace CudaRasterizer {
22−template <typename T>
23−static void obtain(char*& chunk, T*& ptr, std::size_t count,
24− std::size_t alignment) {
25− std::size_t offset =
26− (reinterpret_cast<std::uintptr_t>(chunk) + alignment - 1) &
22+// Number of bits needed to represent tile ids in [0, num_tiles).
23+inline int TileBits(int num_tiles) {
27− ~(alignment - 1);
24+ int bits = 1;
28− ptr = reinterpret_cast<T*>(offset);
29− chunk = reinterpret_cast<char*>(ptr + count);
25+ while ((1 << bits) < num_tiles) ++bits;
26+ return bits;
3027 }
3128
32−struct GeometryState {
29+struct GatherCountsIter {
33− size_t scan_size;
34− float* depths;
35− char* scanning_space;
36− bool* clamped;
37− int* internal_radii;
38− float2* means2D;
39− float* cov3D;
40− float4* conic_opacity;
41− float* rgb;
42− uint32_t* point_offsets;
30+ const int32_t* counts;
43− uint32_t* tiles_touched;
31+ const int32_t* order;
44−
45− static GeometryState fromChunk(char*& chunk, size_t P);
32+ __host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
33+ return counts[order[idx]];
34+ }
35+ __host__ __device__ __forceinline__ GatherCountsIter
36+ operator+(ptrdiff_t offset) const {
37+ return {counts, order + offset};
38+ }
39+ __host__ __device__ __forceinline__ int32_t operator*() const {
40+ return counts[*order];
41+ }
4642 };
4743
48−struct ImageState {
49− uint2* ranges;
44+template <Culling kCulling>
45+__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
46+ const int32_t* __restrict__ sorted_ids,
47+ const int32_t* __restrict__ offsets, int n_vis,
48+ Frame frame, uint16_t* __restrict__ keys,
50− uint32_t* n_contrib;
49+ int32_t* __restrict__ values);
51− float* accum_alpha;
5250
53− static ImageState fromChunk(char*& chunk, size_t N);
51+__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
52+ int num_pairs, int2* __restrict__ ranges);
53+
54+} // namespace gs
55+
56+namespace std {
57+template <>
58+struct iterator_traits<gs::GatherCountsIter> {
59+ using difference_type = ptrdiff_t;
60+ using value_type = int32_t;
61+ using pointer = const int32_t*;
62+ using reference = int32_t;
63+ using iterator_category = random_access_iterator_tag;
5464 };
55−
56−struct BinningState {
57− size_t sorting_size;
58− uint64_t* point_list_keys_unsorted;
59− uint64_t* point_list_keys;
60− uint32_t* point_list_unsorted;
61− uint32_t* point_list;
62− char* list_sorting_space;
63−
64− static BinningState fromChunk(char*& chunk, size_t P);
65−};
66−
67−template <typename T>
68−size_t required(size_t P) {
69− char* size = nullptr;
70− T::fromChunk(size, P);
71− return ((size_t)size) + 128;
72−}
73−}; // namespace CudaRasterizer
65+} // namespace std
csrc/rasterizer_impl.cu.cc +47 −316
s00_baseline/csrc/rasterizer_impl.cu.cc
s01_literature/csrc/rasterizer_impl.cu.cc
↕ Expand 9 unchanged lines (lines 1–9)
1/*1/*
2 * Copyright (C) 2023, Inria2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.4 * All rights reserved.
5 *5 *
6 * This software is free for non-commercial, research and evaluation use6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.7 * under the terms of the LICENSE.md file.
8 *8 *
9 * For inquiries contact george.drettakis@inria.fr9 * For inquiries contact george.drettakis@inria.fr
10 */10 */
1111
12#include "rasterizer_impl.h"12#include "rasterizer_impl.h"
1313
14#include <cuda.h>
15
16#include <algorithm>
17#include <cub/cub.cuh>
18#include <cub/device/device_radix_sort.cuh>
19#include <fstream>
20#include <iostream>
21#include <numeric>
22
23#include "cuda_runtime.h"
24#include "device_launch_parameters.h"
25#define GLM_FORCE_CUDA
26#include <cooperative_groups.h>
27#include <cooperative_groups/reduce.h>
28
29#include <glm/glm.hpp>
30namespace cg = cooperative_groups;
31
32#include "auxiliary.h"14#include "auxiliary.h"
33#include "backward.h"15#include "backward.h"
34#include "forward.h"16#include "forward.h"
3517
36// Helper function to find the next-highest bit of the MSB
37// on the CPU.
18namespace gs {
19
20// ---------------------------------------------------------------------------
21// Binning. Visible Gaussians are sorted by depth; pairs are then emitted in
22// depth order and stably sorted by 16-bit tile id alone.
23
24template <Culling kCulling>
25__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
26 const int32_t* __restrict__ order,
27 const int32_t* __restrict__ inclusive_offsets,
28 int n_vis, Frame frame,
29 uint16_t* __restrict__ keys,
38uint32_t getHigherMsb(uint32_t n) {30 int32_t* __restrict__ values) {
39 uint32_t msb = sizeof(n) * 4;
31 const int k = blockIdx.x * blockDim.x + threadIdx.x;
32 if (k >= n_vis) return;
33 int out = k == 0 ? 0 : inclusive_offsets[k - 1];
34 if (inclusive_offsets[k] == out) return;
35 const int i = order[k];
36 const int tiles_x = frame.grid.tiles_x;
40 uint32_t step = msb;37 const Splat s = splats[i];
41 while (step > 1) {
38 ForEachTileRow<kCulling>(s, frame.grid, frame.width, frame.height,
39 [&](int ty, int tx0, int tx1) {
40 const int row_base = ty * tiles_x;
41 for (int tx = tx0; tx < tx1; ++tx) {
42 keys[out] = static_cast<uint16_t>(row_base + tx);
42 step /= 2;43 values[out] = i;
43 if (n >> msb)
44 msb += step;
45 else
46 msb -= step;
44 ++out;
47 }45 }
48 if (n >> msb) msb++;
49 return msb;
46 });
50}47}
5148
52// Wrapper method to call auxiliary coarse frustum containment test.
53// Mark all Gaussians that pass it.
54__global__ void checkFrustum(int P, const float* orig_points,
55 const float* viewmatrix, const float* projmatrix,
56 bool* present) {
57 auto idx = cg::this_grid().thread_rank();
58 if (idx >= P) return;
59
60 float3 p_view;
61 present[idx] =
62 in_frustum(idx, orig_points, viewmatrix, projmatrix, false, p_view);
63}
64
65// Generates one key/value pair for all Gaussian / tile overlaps.
66// Run once per Gaussian (1:N mapping).
49// Start/end of every tile's run in the sorted pair list. `ranges` must be
50// zeroed beforehand (empty tiles stay [0, 0)).
67__global__ void duplicateWithKeys(int P, const float2* points_xy,51__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
68 const float* depths, const uint32_t* offsets,
69 uint64_t* gaussian_keys_unsorted,
70 uint32_t* gaussian_values_unsorted,
71 int* radii, dim3 grid, int2* rects) {52 int num_pairs, int2* __restrict__ ranges) {
72 auto idx = cg::this_grid().thread_rank();
53 const int i = blockIdx.x * blockDim.x + threadIdx.x;
73 if (idx >= P) return;54 if (i >= num_pairs) return;
74
75 // Generate no key/value pair for invisible Gaussians
76 if (radii[idx] > 0) {
77 // Find this Gaussian's offset in buffer for writing keys/values.
78 uint32_t off = (idx == 0) ? 0 : offsets[idx - 1];
79 uint2 rect_min, rect_max;
80
81 if (rects == nullptr)
82 getRect(points_xy[idx], radii[idx], rect_min, rect_max, grid);
83 else
84 getRect(points_xy[idx], rects[idx], rect_min, rect_max, grid);
85
86 // For each tile that the bounding rect overlaps, emit a
87 // key/value pair. The key is | tile ID | depth |,
88 // and the value is the ID of the Gaussian. Sorting the values
89 // with this key yields Gaussian IDs in a list, such that they
90 // are first sorted by tile and then by depth.
91 for (int y = rect_min.y; y < rect_max.y; y++) {
92 for (int x = rect_min.x; x < rect_max.x; x++) {
93 uint64_t key = y * grid.x + x;55 const uint16_t tile = sorted_tiles[i];
94 key <<= 32;
95 key |= *((uint32_t*)&depths[idx]);
96 gaussian_keys_unsorted[off] = key;
97 gaussian_values_unsorted[off] = idx;
98 off++;
99 }
100 }
56 if (i == 0 || sorted_tiles[i - 1] != tile) ranges[tile].x = i;
57 if (i == num_pairs - 1 || sorted_tiles[i + 1] != tile) {
58 ranges[tile].y = i + 1;
101 }59 }
102}60}
10361
104// Check keys to see if it is at the start/end of one tile's range in
105// the full sorted list. If yes, write start/end of this tile.
106// Run once per instanced (duplicated) Gaussian ID.
107__global__ void identifyTileRanges(int L, uint64_t* point_list_keys,
62template __global__ void EmitPairsKernel<Culling::kRect>(
63 const Splat* __restrict__, const int32_t* __restrict__,
64 const int32_t* __restrict__, int, Frame, uint16_t* __restrict__,
108 uint2* ranges) {65 int32_t* __restrict__);
109 auto idx = cg::this_grid().thread_rank();
110 if (idx >= L) return;
66template __global__ void EmitPairsKernel<Culling::kEllipse>(
67 const Splat* __restrict__, const int32_t* __restrict__,
68 const int32_t* __restrict__, int, Frame, uint16_t* __restrict__,
69 int32_t* __restrict__);
11170
112 // Read tile ID from key. Update start/end of tile range if at limit.
113 uint64_t key = point_list_keys[idx];
114 uint32_t currtile = key >> 32;
115 if (idx == 0)
116 ranges[currtile].x = 0;
117 else {
118 uint32_t prevtile = point_list_keys[idx - 1] >> 32;
119 if (currtile != prevtile) {
120 ranges[prevtile].y = idx;
121 ranges[currtile].x = idx;
122 }
123 if (idx == L - 1) ranges[currtile].y = L;
124 }
125}
126
127// Mark Gaussians as visible/invisible, based on view frustum testing
128void CudaRasterizer::Rasterizer::markVisible(int P, float* means3D,
129 float* viewmatrix,
130 float* projmatrix, bool* present) {
131 checkFrustum<<<(P + 255) / 256, 256>>>(P, means3D, viewmatrix, projmatrix,
132 present);
133}
134
135CudaRasterizer::GeometryState CudaRasterizer::GeometryState::fromChunk(
136 char*& chunk, size_t P) {
137 GeometryState geom;
138 obtain(chunk, geom.depths, P, 128);
139 obtain(chunk, geom.clamped, P * 3, 128);
140 obtain(chunk, geom.internal_radii, P, 128);
141 obtain(chunk, geom.means2D, P, 128);
142 obtain(chunk, geom.cov3D, P * 6, 128);
143 obtain(chunk, geom.conic_opacity, P, 128);
144 obtain(chunk, geom.rgb, P * 3, 128);
145 obtain(chunk, geom.tiles_touched, P, 128);
146 cub::DeviceScan::InclusiveSum(nullptr, geom.scan_size, geom.tiles_touched,
147 geom.tiles_touched, P);
148 obtain(chunk, geom.scanning_space, geom.scan_size, 128);
149 obtain(chunk, geom.point_offsets, P, 128);
150 return geom;
151}
152
153CudaRasterizer::ImageState CudaRasterizer::ImageState::fromChunk(char*& chunk,
154 size_t N) {
155 ImageState img;
156 obtain(chunk, img.accum_alpha, N, 128);
157 obtain(chunk, img.n_contrib, N, 128);
158 obtain(chunk, img.ranges, N, 128);
159 return img;
160}
161
162CudaRasterizer::BinningState CudaRasterizer::BinningState::fromChunk(
163 char*& chunk, size_t P) {
164 BinningState binning;
165 obtain(chunk, binning.point_list, P, 128);
166 obtain(chunk, binning.point_list_unsorted, P, 128);
167 obtain(chunk, binning.point_list_keys, P, 128);
168 obtain(chunk, binning.point_list_keys_unsorted, P, 128);
169 cub::DeviceRadixSort::SortPairs(
170 nullptr, binning.sorting_size, binning.point_list_keys_unsorted,
171 binning.point_list_keys, binning.point_list_unsorted, binning.point_list,
172 P);
173 obtain(chunk, binning.list_sorting_space, binning.sorting_size, 128);
174 return binning;
175}
176
177// Forward rendering procedure for differentiable rasterization
178// of Gaussians.
179int CudaRasterizer::Rasterizer::forward(
180 std::function<char*(size_t)> geometryBuffer,
181 std::function<char*(size_t)> binningBuffer,
182 std::function<char*(size_t)> imageBuffer, const int P, int D, int M,
183 const float* background, const int width, int height, const float* means3D,
184 const float* shs, const float* colors_precomp, const float* opacities,
185 const float* scales, const float scale_modifier, const float* rotations,
186 const float* cov3D_precomp, const float* viewmatrix,
187 const float* projmatrix, const float* cam_pos, const float tan_fovx,
188 float tan_fovy, const float low_pass_2d_kernel_size,
189 const bool mip_splatting_enabled, const bool prefiltered, float* out_color,
190 int* radii, int* rects, float* boxmin, float* boxmax) {
191 const float focal_y = height / (2.0f * tan_fovy);
192 const float focal_x = width / (2.0f * tan_fovx);
193
194 size_t chunk_size = required<GeometryState>(P);
195 char* chunkptr = geometryBuffer(chunk_size);
196 GeometryState geomState = GeometryState::fromChunk(chunkptr, P);
197
198 if (radii == nullptr) {
199 radii = geomState.internal_radii;
200 }
201
202 dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
203 (height + BLOCK_Y - 1) / BLOCK_Y, 1);
204 dim3 block(BLOCK_X, BLOCK_Y, 1);
205
206 // Dynamically resize image-based auxiliary buffers during training
207 int img_chunk_size = required<ImageState>(width * height);
208 char* img_chunkptr = imageBuffer(img_chunk_size);
209 ImageState imgState = ImageState::fromChunk(img_chunkptr, width * height);
210
211 if (NUM_CHANNELS != 3 && colors_precomp == nullptr) {
212 throw std::runtime_error(
213 "For non-RGB, provide precomputed Gaussian colors!");
214 }
215
216 float3 minn = {-FLT_MAX, -FLT_MAX, -FLT_MAX};
217 float3 maxx = {FLT_MAX, FLT_MAX, FLT_MAX};
218 if (boxmin != nullptr) {
219 minn = *((float3*)boxmin);
220 maxx = *((float3*)boxmax);
221 }
222
223 // Run preprocessing per-Gaussian (transformation, bounding, conversion of SHs
224 // to RGB)
225 FORWARD::preprocess(
226 P, D, M, means3D, (glm::vec3*)scales, scale_modifier,
227 (glm::vec4*)rotations, opacities, shs, geomState.clamped, cov3D_precomp,
228 colors_precomp, viewmatrix, projmatrix, (glm::vec3*)cam_pos, width,
229 height, focal_x, focal_y, tan_fovx, tan_fovy, low_pass_2d_kernel_size,
230 mip_splatting_enabled, radii, geomState.means2D, geomState.depths,
231 geomState.cov3D, geomState.rgb, geomState.conic_opacity, tile_grid,
232 geomState.tiles_touched, prefiltered, (int2*)rects, minn, maxx);
233
234 // Compute prefix sum over full list of touched tile counts by Gaussians
235 // E.g., [2, 3, 0, 2, 1] -> [2, 5, 5, 7, 8]
236 cub::DeviceScan::InclusiveSum(geomState.scanning_space, geomState.scan_size,
237 geomState.tiles_touched,
238 geomState.point_offsets, P);
239
240 // Retrieve total number of Gaussian instances to launch and resize aux
241 // buffers
242 int num_rendered;
243 cudaMemcpy(&num_rendered, geomState.point_offsets + P - 1, sizeof(int),
244 cudaMemcpyDeviceToHost);
245
246 if (num_rendered == 0) return 0;
247
248 int binning_chunk_size = required<BinningState>(num_rendered);
249 char* binning_chunkptr = binningBuffer(binning_chunk_size);
250 BinningState binningState =
251 BinningState::fromChunk(binning_chunkptr, num_rendered);
252
253 // For each instance to be rendered, produce adequate [ tile | depth ] key
254 // and corresponding dublicated Gaussian indices to be sorted
255 duplicateWithKeys<<<(P + 255) / 256, 256>>>(
256 P, geomState.means2D, geomState.depths, geomState.point_offsets,
257 binningState.point_list_keys_unsorted, binningState.point_list_unsorted,
258 radii, tile_grid, (int2*)rects);
259
260 int bit = getHigherMsb(tile_grid.x * tile_grid.y);
261
262 // Sort complete list of (duplicated) Gaussian indices by keys
263 cub::DeviceRadixSort::SortPairs(
264 binningState.list_sorting_space, binningState.sorting_size,
265 binningState.point_list_keys_unsorted, binningState.point_list_keys,
266 binningState.point_list_unsorted, binningState.point_list, num_rendered,
267 0, 32 + bit);
268
269 cudaMemset(imgState.ranges, 0, tile_grid.x * tile_grid.y * sizeof(uint2));
270
271 // Identify start and end of per-tile workloads in sorted list
272 identifyTileRanges<<<(num_rendered + 255) / 256, 256>>>(
273 num_rendered, binningState.point_list_keys, imgState.ranges);
274
275 // Let each tile blend its range of Gaussians independently in parallel
276 const float* feature_ptr =
277 colors_precomp != nullptr ? colors_precomp : geomState.rgb;
278 FORWARD::render(tile_grid, block, imgState.ranges, binningState.point_list,
279 width, height, geomState.means2D, feature_ptr,
280 geomState.conic_opacity, imgState.accum_alpha,
281 imgState.n_contrib, background, out_color);
282
283 return num_rendered;
284}
285
286// Produce necessary gradients for optimization, corresponding
287// to forward render pass
288void CudaRasterizer::Rasterizer::backward(
289 const int P, int D, int M, int R, const float* background, const int width,
290 int height, const float* means3D, const float* shs,
291 const float* colors_precomp, const float* scales,
292 const float scale_modifier, const float* rotations,
293 const float* cov3D_precomp, const float* viewmatrix,
294 const float* projmatrix, const float* campos, const float tan_fovx,
295 float tan_fovy, const float low_pass_2d_kernel_size,
296 const bool mip_splatting_enabled, const int* radii, char* geom_buffer,
297 char* binning_buffer, char* img_buffer, const float* dL_dpix,
298 float* dL_dmean2D, float* dL_dconic, float* dL_dopacity, float* dL_dcolor,
299 float* dL_dmean3D, float* dL_dcov3D, float* dL_dsh, float* dL_dscale,
300 float* dL_drot) {
301 GeometryState geomState = GeometryState::fromChunk(geom_buffer, P);
302 BinningState binningState = BinningState::fromChunk(binning_buffer, R);
303 ImageState imgState = ImageState::fromChunk(img_buffer, width * height);
304
305 if (radii == nullptr) {
306 radii = geomState.internal_radii;
307 }
308
309 const float focal_y = height / (2.0f * tan_fovy);
310 const float focal_x = width / (2.0f * tan_fovx);
311
312 const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
313 (height + BLOCK_Y - 1) / BLOCK_Y, 1);
314 const dim3 block(BLOCK_X, BLOCK_Y, 1);
315
316 // Compute loss gradients w.r.t. 2D mean position, conic matrix,
317 // opacity and RGB of Gaussians from per-pixel loss gradients.
318 // If we were given precomputed colors and not SHs, use them.
319 const float* color_ptr =
320 (colors_precomp != nullptr) ? colors_precomp : geomState.rgb;
321 BACKWARD::render(tile_grid, block, imgState.ranges, binningState.point_list,
322 width, height, background, geomState.means2D,
323 geomState.conic_opacity, color_ptr, imgState.accum_alpha,
324 imgState.n_contrib, dL_dpix, (float3*)dL_dmean2D,
325 (float4*)dL_dconic, dL_dopacity, dL_dcolor);
326
327 // Take care of the rest of preprocessing. Was the precomputed covariance
328 // given to us or a scales/rot pair? If precomputed, pass that. If not,
329 // use the one we computed ourselves.
330 const float* cov3D_ptr =
331 (cov3D_precomp != nullptr) ? cov3D_precomp : geomState.cov3D;
332 BACKWARD::preprocess(P, D, M, (float3*)means3D, radii, shs, geomState.clamped,
333 (glm::vec3*)scales, (glm::vec4*)rotations,
334 scale_modifier, cov3D_ptr, geomState.conic_opacity,
335 viewmatrix, projmatrix, focal_x, focal_y, tan_fovx,
336 tan_fovy, low_pass_2d_kernel_size, mip_splatting_enabled,
337 (glm::vec3*)campos, (float3*)dL_dmean2D, dL_dconic,
338 (glm::vec3*)dL_dmean3D, dL_dcolor, dL_dcov3D, dL_dsh,
339 (glm::vec3*)dL_dscale, (glm::vec4*)dL_drot, dL_dopacity);
340}
71} // namespace gs
↕ Expand 9 unchanged lines (lines 1–9)
11 /*
22 * Copyright (C) 2023, Inria
33 * GRAPHDECO research group, https://team.inria.fr/graphdeco
44 * All rights reserved.
55 *
66 * This software is free for non-commercial, research and evaluation use
77 * under the terms of the LICENSE.md file.
88 *
99 * For inquiries contact george.drettakis@inria.fr
1010 */
1111
1212 #include "rasterizer_impl.h"
1313
14−#include <cuda.h>
15−
16−#include <algorithm>
17−#include <cub/cub.cuh>
18−#include <cub/device/device_radix_sort.cuh>
19−#include <fstream>
20−#include <iostream>
21−#include <numeric>
22−
23−#include "cuda_runtime.h"
24−#include "device_launch_parameters.h"
25−#define GLM_FORCE_CUDA
26−#include <cooperative_groups.h>
27−#include <cooperative_groups/reduce.h>
28−
29−#include <glm/glm.hpp>
30−namespace cg = cooperative_groups;
31−
3214 #include "auxiliary.h"
3315 #include "backward.h"
3416 #include "forward.h"
3517
36−// Helper function to find the next-highest bit of the MSB
37−// on the CPU.
18+namespace gs {
19+
20+// ---------------------------------------------------------------------------
21+// Binning. Visible Gaussians are sorted by depth; pairs are then emitted in
22+// depth order and stably sorted by 16-bit tile id alone.
23+
24+template <Culling kCulling>
25+__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
26+ const int32_t* __restrict__ order,
27+ const int32_t* __restrict__ inclusive_offsets,
28+ int n_vis, Frame frame,
29+ uint16_t* __restrict__ keys,
38−uint32_t getHigherMsb(uint32_t n) {
30+ int32_t* __restrict__ values) {
39− uint32_t msb = sizeof(n) * 4;
31+ const int k = blockIdx.x * blockDim.x + threadIdx.x;
32+ if (k >= n_vis) return;
33+ int out = k == 0 ? 0 : inclusive_offsets[k - 1];
34+ if (inclusive_offsets[k] == out) return;
35+ const int i = order[k];
36+ const int tiles_x = frame.grid.tiles_x;
40− uint32_t step = msb;
37+ const Splat s = splats[i];
41− while (step > 1) {
38+ ForEachTileRow<kCulling>(s, frame.grid, frame.width, frame.height,
39+ [&](int ty, int tx0, int tx1) {
40+ const int row_base = ty * tiles_x;
41+ for (int tx = tx0; tx < tx1; ++tx) {
42+ keys[out] = static_cast<uint16_t>(row_base + tx);
42− step /= 2;
43+ values[out] = i;
43− if (n >> msb)
44− msb += step;
45− else
46− msb -= step;
44+ ++out;
4745⇥ }
48− if (n >> msb) msb++;
49− return msb;
46+ });
5047 }
5148
52−// Wrapper method to call auxiliary coarse frustum containment test.
53−// Mark all Gaussians that pass it.
54−__global__ void checkFrustum(int P, const float* orig_points,
55− const float* viewmatrix, const float* projmatrix,
56− bool* present) {
57− auto idx = cg::this_grid().thread_rank();
58− if (idx >= P) return;
59−
60− float3 p_view;
61− present[idx] =
62− in_frustum(idx, orig_points, viewmatrix, projmatrix, false, p_view);
63−}
64−
65−// Generates one key/value pair for all Gaussian / tile overlaps.
66−// Run once per Gaussian (1:N mapping).
49+// Start/end of every tile's run in the sorted pair list. `ranges` must be
50+// zeroed beforehand (empty tiles stay [0, 0)).
67−__global__ void duplicateWithKeys(int P, const float2* points_xy,
51+__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
68− const float* depths, const uint32_t* offsets,
69− uint64_t* gaussian_keys_unsorted,
70− uint32_t* gaussian_values_unsorted,
71− int* radii, dim3 grid, int2* rects) {
52+ int num_pairs, int2* __restrict__ ranges) {
72− auto idx = cg::this_grid().thread_rank();
53+ const int i = blockIdx.x * blockDim.x + threadIdx.x;
73− if (idx >= P) return;
54+ if (i >= num_pairs) return;
74−
75− // Generate no key/value pair for invisible Gaussians
76− if (radii[idx] > 0) {
77− // Find this Gaussian's offset in buffer for writing keys/values.
78− uint32_t off = (idx == 0) ? 0 : offsets[idx - 1];
79− uint2 rect_min, rect_max;
80−
81− if (rects == nullptr)
82− getRect(points_xy[idx], radii[idx], rect_min, rect_max, grid);
83− else
84− getRect(points_xy[idx], rects[idx], rect_min, rect_max, grid);
85−
86− // For each tile that the bounding rect overlaps, emit a
87− // key/value pair. The key is | tile ID | depth |,
88− // and the value is the ID of the Gaussian. Sorting the values
89− // with this key yields Gaussian IDs in a list, such that they
90− // are first sorted by tile and then by depth.
91− for (int y = rect_min.y; y < rect_max.y; y++) {
92− for (int x = rect_min.x; x < rect_max.x; x++) {
93− uint64_t key = y * grid.x + x;
55+ const uint16_t tile = sorted_tiles[i];
94− key <<= 32;
95− key |= *((uint32_t*)&depths[idx]);
96− gaussian_keys_unsorted[off] = key;
97− gaussian_values_unsorted[off] = idx;
98− off++;
99− }
100− }
56+ if (i == 0 || sorted_tiles[i - 1] != tile) ranges[tile].x = i;
57+ if (i == num_pairs - 1 || sorted_tiles[i + 1] != tile) {
58+ ranges[tile].y = i + 1;
10159 }
10260 }
10361
104−// Check keys to see if it is at the start/end of one tile's range in
105−// the full sorted list. If yes, write start/end of this tile.
106−// Run once per instanced (duplicated) Gaussian ID.
107−__global__ void identifyTileRanges(int L, uint64_t* point_list_keys,
62+template __global__ void EmitPairsKernel<Culling::kRect>(
63+ const Splat* __restrict__, const int32_t* __restrict__,
64+ const int32_t* __restrict__, int, Frame, uint16_t* __restrict__,
108− uint2* ranges) {
65+ int32_t* __restrict__);
109− auto idx = cg::this_grid().thread_rank();
110− if (idx >= L) return;
66+template __global__ void EmitPairsKernel<Culling::kEllipse>(
67+ const Splat* __restrict__, const int32_t* __restrict__,
68+ const int32_t* __restrict__, int, Frame, uint16_t* __restrict__,
69+ int32_t* __restrict__);
11170
112− // Read tile ID from key. Update start/end of tile range if at limit.
113− uint64_t key = point_list_keys[idx];
114− uint32_t currtile = key >> 32;
115− if (idx == 0)
116− ranges[currtile].x = 0;
117− else {
118− uint32_t prevtile = point_list_keys[idx - 1] >> 32;
119− if (currtile != prevtile) {
120− ranges[prevtile].y = idx;
121− ranges[currtile].x = idx;
122− }
123− if (idx == L - 1) ranges[currtile].y = L;
124− }
125−}
126−
127−// Mark Gaussians as visible/invisible, based on view frustum testing
128−void CudaRasterizer::Rasterizer::markVisible(int P, float* means3D,
129− float* viewmatrix,
130− float* projmatrix, bool* present) {
131− checkFrustum<<<(P + 255) / 256, 256>>>(P, means3D, viewmatrix, projmatrix,
132− present);
133−}
134−
135−CudaRasterizer::GeometryState CudaRasterizer::GeometryState::fromChunk(
136− char*& chunk, size_t P) {
137− GeometryState geom;
138− obtain(chunk, geom.depths, P, 128);
139− obtain(chunk, geom.clamped, P * 3, 128);
140− obtain(chunk, geom.internal_radii, P, 128);
141− obtain(chunk, geom.means2D, P, 128);
142− obtain(chunk, geom.cov3D, P * 6, 128);
143− obtain(chunk, geom.conic_opacity, P, 128);
144− obtain(chunk, geom.rgb, P * 3, 128);
145− obtain(chunk, geom.tiles_touched, P, 128);
146− cub::DeviceScan::InclusiveSum(nullptr, geom.scan_size, geom.tiles_touched,
147− geom.tiles_touched, P);
148− obtain(chunk, geom.scanning_space, geom.scan_size, 128);
149− obtain(chunk, geom.point_offsets, P, 128);
150− return geom;
151−}
152−
153−CudaRasterizer::ImageState CudaRasterizer::ImageState::fromChunk(char*& chunk,
154− size_t N) {
155− ImageState img;
156− obtain(chunk, img.accum_alpha, N, 128);
157− obtain(chunk, img.n_contrib, N, 128);
158− obtain(chunk, img.ranges, N, 128);
159− return img;
160−}
161−
162−CudaRasterizer::BinningState CudaRasterizer::BinningState::fromChunk(
163− char*& chunk, size_t P) {
164− BinningState binning;
165− obtain(chunk, binning.point_list, P, 128);
166− obtain(chunk, binning.point_list_unsorted, P, 128);
167− obtain(chunk, binning.point_list_keys, P, 128);
168− obtain(chunk, binning.point_list_keys_unsorted, P, 128);
169− cub::DeviceRadixSort::SortPairs(
170− nullptr, binning.sorting_size, binning.point_list_keys_unsorted,
171− binning.point_list_keys, binning.point_list_unsorted, binning.point_list,
172− P);
173− obtain(chunk, binning.list_sorting_space, binning.sorting_size, 128);
174− return binning;
175−}
176−
177−// Forward rendering procedure for differentiable rasterization
178−// of Gaussians.
179−int CudaRasterizer::Rasterizer::forward(
180− std::function<char*(size_t)> geometryBuffer,
181− std::function<char*(size_t)> binningBuffer,
182− std::function<char*(size_t)> imageBuffer, const int P, int D, int M,
183− const float* background, const int width, int height, const float* means3D,
184− const float* shs, const float* colors_precomp, const float* opacities,
185− const float* scales, const float scale_modifier, const float* rotations,
186− const float* cov3D_precomp, const float* viewmatrix,
187− const float* projmatrix, const float* cam_pos, const float tan_fovx,
188− float tan_fovy, const float low_pass_2d_kernel_size,
189− const bool mip_splatting_enabled, const bool prefiltered, float* out_color,
190− int* radii, int* rects, float* boxmin, float* boxmax) {
191− const float focal_y = height / (2.0f * tan_fovy);
192− const float focal_x = width / (2.0f * tan_fovx);
193−
194− size_t chunk_size = required<GeometryState>(P);
195− char* chunkptr = geometryBuffer(chunk_size);
196− GeometryState geomState = GeometryState::fromChunk(chunkptr, P);
197−
198− if (radii == nullptr) {
199− radii = geomState.internal_radii;
200− }
201−
202− dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
203− (height + BLOCK_Y - 1) / BLOCK_Y, 1);
204− dim3 block(BLOCK_X, BLOCK_Y, 1);
205−
206− // Dynamically resize image-based auxiliary buffers during training
207− int img_chunk_size = required<ImageState>(width * height);
208− char* img_chunkptr = imageBuffer(img_chunk_size);
209− ImageState imgState = ImageState::fromChunk(img_chunkptr, width * height);
210−
211− if (NUM_CHANNELS != 3 && colors_precomp == nullptr) {
212− throw std::runtime_error(
213− "For non-RGB, provide precomputed Gaussian colors!");
214− }
215−
216− float3 minn = {-FLT_MAX, -FLT_MAX, -FLT_MAX};
217− float3 maxx = {FLT_MAX, FLT_MAX, FLT_MAX};
218− if (boxmin != nullptr) {
219− minn = *((float3*)boxmin);
220− maxx = *((float3*)boxmax);
221− }
222−
223− // Run preprocessing per-Gaussian (transformation, bounding, conversion of SHs
224− // to RGB)
225− FORWARD::preprocess(
226− P, D, M, means3D, (glm::vec3*)scales, scale_modifier,
227− (glm::vec4*)rotations, opacities, shs, geomState.clamped, cov3D_precomp,
228− colors_precomp, viewmatrix, projmatrix, (glm::vec3*)cam_pos, width,
229− height, focal_x, focal_y, tan_fovx, tan_fovy, low_pass_2d_kernel_size,
230− mip_splatting_enabled, radii, geomState.means2D, geomState.depths,
231− geomState.cov3D, geomState.rgb, geomState.conic_opacity, tile_grid,
232− geomState.tiles_touched, prefiltered, (int2*)rects, minn, maxx);
233−
234− // Compute prefix sum over full list of touched tile counts by Gaussians
235− // E.g., [2, 3, 0, 2, 1] -> [2, 5, 5, 7, 8]
236− cub::DeviceScan::InclusiveSum(geomState.scanning_space, geomState.scan_size,
237− geomState.tiles_touched,
238− geomState.point_offsets, P);
239−
240− // Retrieve total number of Gaussian instances to launch and resize aux
241− // buffers
242− int num_rendered;
243− cudaMemcpy(&num_rendered, geomState.point_offsets + P - 1, sizeof(int),
244− cudaMemcpyDeviceToHost);
245−
246− if (num_rendered == 0) return 0;
247−
248− int binning_chunk_size = required<BinningState>(num_rendered);
249− char* binning_chunkptr = binningBuffer(binning_chunk_size);
250− BinningState binningState =
251− BinningState::fromChunk(binning_chunkptr, num_rendered);
252−
253− // For each instance to be rendered, produce adequate [ tile | depth ] key
254− // and corresponding dublicated Gaussian indices to be sorted
255− duplicateWithKeys<<<(P + 255) / 256, 256>>>(
256− P, geomState.means2D, geomState.depths, geomState.point_offsets,
257− binningState.point_list_keys_unsorted, binningState.point_list_unsorted,
258− radii, tile_grid, (int2*)rects);
259−
260− int bit = getHigherMsb(tile_grid.x * tile_grid.y);
261−
262− // Sort complete list of (duplicated) Gaussian indices by keys
263− cub::DeviceRadixSort::SortPairs(
264− binningState.list_sorting_space, binningState.sorting_size,
265− binningState.point_list_keys_unsorted, binningState.point_list_keys,
266− binningState.point_list_unsorted, binningState.point_list, num_rendered,
267− 0, 32 + bit);
268−
269− cudaMemset(imgState.ranges, 0, tile_grid.x * tile_grid.y * sizeof(uint2));
270−
271− // Identify start and end of per-tile workloads in sorted list
272− identifyTileRanges<<<(num_rendered + 255) / 256, 256>>>(
273− num_rendered, binningState.point_list_keys, imgState.ranges);
274−
275− // Let each tile blend its range of Gaussians independently in parallel
276− const float* feature_ptr =
277− colors_precomp != nullptr ? colors_precomp : geomState.rgb;
278− FORWARD::render(tile_grid, block, imgState.ranges, binningState.point_list,
279− width, height, geomState.means2D, feature_ptr,
280− geomState.conic_opacity, imgState.accum_alpha,
281− imgState.n_contrib, background, out_color);
282−
283− return num_rendered;
284−}
285−
286−// Produce necessary gradients for optimization, corresponding
287−// to forward render pass
288−void CudaRasterizer::Rasterizer::backward(
289− const int P, int D, int M, int R, const float* background, const int width,
290− int height, const float* means3D, const float* shs,
291− const float* colors_precomp, const float* scales,
292− const float scale_modifier, const float* rotations,
293− const float* cov3D_precomp, const float* viewmatrix,
294− const float* projmatrix, const float* campos, const float tan_fovx,
295− float tan_fovy, const float low_pass_2d_kernel_size,
296− const bool mip_splatting_enabled, const int* radii, char* geom_buffer,
297− char* binning_buffer, char* img_buffer, const float* dL_dpix,
298− float* dL_dmean2D, float* dL_dconic, float* dL_dopacity, float* dL_dcolor,
299− float* dL_dmean3D, float* dL_dcov3D, float* dL_dsh, float* dL_dscale,
300− float* dL_drot) {
301− GeometryState geomState = GeometryState::fromChunk(geom_buffer, P);
302− BinningState binningState = BinningState::fromChunk(binning_buffer, R);
303− ImageState imgState = ImageState::fromChunk(img_buffer, width * height);
304−
305− if (radii == nullptr) {
306− radii = geomState.internal_radii;
307− }
308−
309− const float focal_y = height / (2.0f * tan_fovy);
310− const float focal_x = width / (2.0f * tan_fovx);
311−
312− const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
313− (height + BLOCK_Y - 1) / BLOCK_Y, 1);
314− const dim3 block(BLOCK_X, BLOCK_Y, 1);
315−
316− // Compute loss gradients w.r.t. 2D mean position, conic matrix,
317− // opacity and RGB of Gaussians from per-pixel loss gradients.
318− // If we were given precomputed colors and not SHs, use them.
319− const float* color_ptr =
320− (colors_precomp != nullptr) ? colors_precomp : geomState.rgb;
321− BACKWARD::render(tile_grid, block, imgState.ranges, binningState.point_list,
322− width, height, background, geomState.means2D,
323− geomState.conic_opacity, color_ptr, imgState.accum_alpha,
324− imgState.n_contrib, dL_dpix, (float3*)dL_dmean2D,
325− (float4*)dL_dconic, dL_dopacity, dL_dcolor);
326−
327− // Take care of the rest of preprocessing. Was the precomputed covariance
328− // given to us or a scales/rot pair? If precomputed, pass that. If not,
329− // use the one we computed ourselves.
330− const float* cov3D_ptr =
331− (cov3D_precomp != nullptr) ? cov3D_precomp : geomState.cov3D;
332− BACKWARD::preprocess(P, D, M, (float3*)means3D, radii, shs, geomState.clamped,
333− (glm::vec3*)scales, (glm::vec4*)rotations,
334− scale_modifier, cov3D_ptr, geomState.conic_opacity,
335− viewmatrix, projmatrix, focal_x, focal_y, tan_fovx,
336− tan_fovy, low_pass_2d_kernel_size, mip_splatting_enabled,
337− (glm::vec3*)campos, (float3*)dL_dmean2D, dL_dconic,
338− (glm::vec3*)dL_dmean3D, dL_dcolor, dL_dcov3D, dL_dsh,
339− (glm::vec3*)dL_dscale, (glm::vec4*)dL_drot, dL_dopacity);
340−}
71+} // namespace gs
Unchanged files in s01_literature (identical to s00_baseline): csrc/config.h