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

s03_sh_streaming → s04_raster_pruning

Step 4: In-Place Forward Contributor Compaction & splat_touched Gating  ·  +262 / −111 lines across 6 changed files

Core files: csrc/auxiliary.h +12/−3csrc/forward.h +3/−2csrc/forward.cu.cc +119/−37csrc/backward.h +7/−4csrc/backward.cu.cc +118/−65csrc/rasterizer_impl.h +3/−0
csrc/auxiliary.h +12 −3
s03_sh_streaming/csrc/auxiliary.h
s04_raster_pruning/csrc/auxiliary.h
↕ Expand 689 unchanged lines (lines 1–689) · in Exp2Approx
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
1414
15#include <cuda_runtime.h>15#include <cuda_runtime.h>
1616
17#include <cmath>17#include <cmath>
18#include <cstdint>18#include <cstdint>
1919
20#include "config.h"20#include "config.h"
21#include "stdio.h"21#include "stdio.h"
2222
23#define BLOCK_SIZE (BLOCK_X * BLOCK_Y)23#define BLOCK_SIZE (BLOCK_X * BLOCK_Y)
24#define NUM_WARPS (BLOCK_SIZE / 32)24#define NUM_WARPS (BLOCK_SIZE / 32)
2525
26// Spherical harmonics coefficients26// Spherical harmonics coefficients
27__device__ const float SH_C0 = 0.28209479177387814f;27__device__ const float SH_C0 = 0.28209479177387814f;
28__device__ const float SH_C1 = 0.4886025119029199f;28__device__ const float SH_C1 = 0.4886025119029199f;
29__device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,29__device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,
30 0.31539156525252005f, -1.0925484305920792f,30 0.31539156525252005f, -1.0925484305920792f,
31 0.5462742152960396f};31 0.5462742152960396f};
32__device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,32__device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,
33 -0.4570457994644658f, 0.3731763325901154f,33 -0.4570457994644658f, 0.3731763325901154f,
34 -0.4570457994644658f, 1.445305721320277f,34 -0.4570457994644658f, 1.445305721320277f,
35 -0.5900435899266435f};35 -0.5900435899266435f};
3636
37namespace gs {37namespace gs {
3838
39#define GS_HD __host__ __device__ __forceinline__39#define GS_HD __host__ __device__ __forceinline__
40#define GS_D __device__ __forceinline__40#define GS_D __device__ __forceinline__
4141
42GS_HD float3 operator+(float3 a, float3 b) {42GS_HD float3 operator+(float3 a, float3 b) {
43 return {a.x + b.x, a.y + b.y, a.z + b.z};43 return {a.x + b.x, a.y + b.y, a.z + b.z};
44}44}
45GS_HD float3 operator-(float3 a, float3 b) {45GS_HD float3 operator-(float3 a, float3 b) {
46 return {a.x - b.x, a.y - b.y, a.z - b.z};46 return {a.x - b.x, a.y - b.y, a.z - b.z};
47}47}
48GS_HD float3 operator*(float s, float3 a) {48GS_HD float3 operator*(float s, float3 a) {
49 return {s * a.x, s * a.y, s * a.z};49 return {s * a.x, s * a.y, s * a.z};
50}50}
51GS_HD float3& operator+=(float3& a, float3 b) {51GS_HD float3& operator+=(float3& a, float3 b) {
52 a.x += b.x;52 a.x += b.x;
53 a.y += b.y;53 a.y += b.y;
54 a.z += b.z;54 a.z += b.z;
55 return a;55 return a;
56}56}
57GS_HD float dot(float3 a, float3 b) {57GS_HD float dot(float3 a, float3 b) {
58 return a.x * b.x + a.y * b.y + a.z * b.z;58 return a.x * b.x + a.y * b.y + a.z * b.z;
59}59}
6060
61// Row-major 3x3 matrix.61// Row-major 3x3 matrix.
62struct Mat3 {62struct Mat3 {
63 float m[3][3];63 float m[3][3];
64 GS_HD float3 row(int i) const { return {m[i][0], m[i][1], m[i][2]}; }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]}; }65 GS_HD float3 col(int j) const { return {m[0][j], m[1][j], m[2][j]}; }
66};66};
6767
68GS_HD float3 Mul(const Mat3& a, float3 v) {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)};69 return {dot(a.row(0), v), dot(a.row(1), v), dot(a.row(2), v)};
70}70}
71// a^T v.71// a^T v.
72GS_HD float3 MulT(const Mat3& a, float3 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)};73 return {dot(a.col(0), v), dot(a.col(1), v), dot(a.col(2), v)};
74}74}
7575
76// Reinterprets bits; used to pack small integers into float records.76// Reinterprets bits; used to pack small integers into float records.
77GS_HD float AsFloat(uint32_t u) {77GS_HD float AsFloat(uint32_t u) {
78 union {78 union {
79 uint32_t u;79 uint32_t u;
80 float f;80 float f;
81 } c{u};81 } c{u};
82 return c.f;82 return c.f;
83}83}
84GS_HD uint32_t AsUint(float f) {84GS_HD uint32_t AsUint(float f) {
85 union {85 union {
86 float f;86 float f;
87 uint32_t u;87 uint32_t u;
88 } c{f};88 } c{f};
89 return c.u;89 return c.u;
90}90}
9191
92// Number of floats in the packed camera buffer built by `camera.py`: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].93// [R (9, row-major world-to-camera), t (3), fx, fy, cx, cy, position (3), 0].
94inline constexpr int kCameraFloats = 20;94inline constexpr int kCameraFloats = 20;
9595
96// Pinhole camera with OpenCV axes (+x right, +y down, +z forward). Pixel (i, j)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.97// is sampled at coordinate (j, i), as in the Inria rasterizer.
98struct Camera {98struct Camera {
99 Mat3 rotation; // World-to-camera rotation.99 Mat3 rotation; // World-to-camera rotation.
100 float3 translation;100 float3 translation;
101 float fx, fy, cx, cy;101 float fx, fy, cx, cy;
102 float3 position; // Camera center in world space.102 float3 position; // Camera center in world space.
103 int width, height;103 int width, height;
104104
105 GS_HD static Camera Load(const float* p, int width, int height) {105 GS_HD static Camera Load(const float* p, int width, int height) {
106 Camera c;106 Camera c;
107 for (int i = 0; i < 3; ++i)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];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]};109 c.translation = {p[9], p[10], p[11]};
110 c.fx = p[12];110 c.fx = p[12];
111 c.fy = p[13];111 c.fy = p[13];
112 c.cx = p[14];112 c.cx = p[14];
113 c.cy = p[15];113 c.cy = p[15];
114 c.position = {p[16], p[17], p[18]};114 c.position = {p[16], p[17], p[18]};
115 c.width = width;115 c.width = width;
116 c.height = height;116 c.height = height;
117 return c;117 return c;
118 }118 }
119119
120 GS_HD float3 ToView(float3 p) const { return Mul(rotation, p) + translation; }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.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); }122 GS_HD float TanHalfFovX() const { return width / (2.0f * fx); }
123 GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }123 GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }
124};124};
125125
126// Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block126// Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block
127// rasterizes one tile.127// rasterizes one tile.
128inline constexpr int kTileSize = 16;128inline constexpr int kTileSize = 16;
129inline constexpr int kTilePixels = kTileSize * kTileSize;129inline constexpr int kTilePixels = kTileSize * kTileSize;
130130
131struct TileGrid {131struct TileGrid {
132 int tiles_x, tiles_y;132 int tiles_x, tiles_y;
133 GS_HD static TileGrid For(int width, int height) {133 GS_HD static TileGrid For(int width, int height) {
134 return {(width + kTileSize - 1) / kTileSize,134 return {(width + kTileSize - 1) / kTileSize,
135 (height + kTileSize - 1) / kTileSize};135 (height + kTileSize - 1) / kTileSize};
136 }136 }
137 GS_HD int num_tiles() const { return tiles_x * tiles_y; }137 GS_HD int num_tiles() const { return tiles_x * tiles_y; }
138};138};
139139
140// Smallest alpha that contributes to a pixel; smaller ones are skipped.140// Smallest alpha that contributes to a pixel; smaller ones are skipped.
141inline constexpr float kMinAlpha = 1.0f / 255.0f;141inline constexpr float kMinAlpha = 1.0f / 255.0f;
142inline constexpr float kMaxAlpha = 0.99f;142inline constexpr float kMaxAlpha = 0.99f;
143// Transmittance below which a pixel stops accumulating.143// Transmittance below which a pixel stops accumulating.
144inline constexpr float kMinTransmittance = 1e-4f;144inline constexpr float kMinTransmittance = 1e-4f;
145145
146// One splat as stored in memory: three float4 so the rasterizer issues three146// One splat as stored in memory: three float4 so the rasterizer issues three
147// 16-byte loads. Layout:147// 16-byte loads. Layout:
148// p0 = (x, y, conic.a, conic.b)148// p0 = (x, y, conic.a, conic.b)
149// p1 = (conic.c, opacity, limit, radius) radius 0 => culled149// p1 = (conic.c, opacity, limit, radius) radius 0 => culled
150// p2 = (r, g, b, min_power)150// p2 = (r, g, b, min_power)
151inline constexpr int kSplatFloats = 12;151inline constexpr int kSplatFloats = 12;
152struct alignas(16) Splat {152struct alignas(16) Splat {
153 float4 p0, p1, p2;153 float4 p0, p1, p2;
154 GS_HD float2 xy() const { return {p0.x, p0.y}; }154 GS_HD float2 xy() const { return {p0.x, p0.y}; }
155 GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }155 GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }
156 GS_HD float opacity() const { return p1.y; }156 GS_HD float opacity() const { return p1.y; }
157 GS_HD float limit() const { return p1.z; }157 GS_HD float limit() const { return p1.z; }
158 GS_HD int radius() const { return static_cast<int>(p1.w); }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}; }159 GS_HD float3 rgb() const { return {p2.x, p2.y, p2.z}; }
160};160};
161161
162// Gradient w.r.t. a splat's rasterized quantities, accumulated by the raster162// 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).163// backward pass: (dx, dy, da, db, dc, dopacity, dr, dg, db).
164inline constexpr int kSplatGradFloats = 9;164inline constexpr int kSplatGradFloats = 9;
165165
166// Half-open tile rectangle [x0, x1) x [y0, y1).166// Half-open tile rectangle [x0, x1) x [y0, y1).
167struct TileRect {167struct TileRect {
168 int x0, y0, x1, y1;168 int x0, y0, x1, y1;
169 GS_HD int area() const { return (x1 - x0) * (y1 - y0); }169 GS_HD int area() const { return (x1 - x0) * (y1 - y0); }
170};170};
171171
172// Tile rectangle of a splat of integer bounding `radius`, identical to the172// Tile rectangle of a splat of integer bounding `radius`, identical to the
173// reference rasterizer.173// reference rasterizer.
174GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {174GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {
175 const float r = static_cast<float>(radius);175 const float r = static_cast<float>(radius);
176 auto lo = [](float v, int n) {176 auto lo = [](float v, int n) {
177 return min(n, max(0, static_cast<int>(v / kTileSize)));177 return min(n, max(0, static_cast<int>(v / kTileSize)));
178 };178 };
179 auto hi = [](float v, int n) {179 auto hi = [](float v, int n) {
180 return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));180 return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));
181 };181 };
182 return {lo(xy.x - r, grid.tiles_x), lo(xy.y - r, grid.tiles_y),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)};183 hi(xy.x + r, grid.tiles_x), hi(xy.y + r, grid.tiles_y)};
184}184}
185185
186// How splats are assigned to tiles.186// How splats are assigned to tiles.
187enum class Culling : int {187enum class Culling : int {
188 // Every tile of the bounding rectangle (reference behaviour).188 // Every tile of the bounding rectangle (reference behaviour).
189 kRect = 0,189 kRect = 0,
190 // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives190 // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives
191 // the same image: skipped (splat, tile) pairs never contribute.191 // the same image: skipped (splat, tile) pairs never contribute.
192 kEllipse = 1,192 kEllipse = 1,
193};193};
194194
195// alpha = o exp(-q/2) >= 1/255 <=> q <= 2 ln(255 o). Writes the bound on q195// 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 the196// with margins that keep tests conservative under float rounding of the
197// rasterizer's power. Returns false if no pixel can reach 1/255.197// rasterizer's power. Returns false if no pixel can reach 1/255.
198GS_HD bool AlphaLimit(float opacity, float& limit) {198GS_HD bool AlphaLimit(float opacity, float& limit) {
199 const float o255 = opacity * 255.0f;199 const float o255 = opacity * 255.0f;
200 if (!(o255 >= 0.999f)) return false;200 if (!(o255 >= 0.999f)) return false;
201 limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;201 limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;
202 return true;202 return true;
203}203}
204204
205// Pad (pixels) on ellipse extents, against rounding.205// Pad (pixels) on ellipse extents, against rounding.
206inline constexpr float kSpanPadPx = 0.01f;206inline constexpr float kSpanPadPx = 0.01f;
207207
208// Precomputed splat-level ellipse constants for fast band intersection tests208// 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.209// of {q(d) <= limit}, where q(d) = a dx^2 + 2 b dx dy + c dy^2.
210struct EllipseBounds {210struct EllipseBounds {
211 float a, b, c;211 float a, b, c;
212 float det;212 float det;
213 float inv_a;213 float inv_a;
214 float a_limit;214 float a_limit;
215 float y_ext;215 float y_ext;
216 float dy_peak;216 float dy_peak;
217217
218 GS_D static bool Init(float3 conic, float opacity, float cached_limit,218 GS_D static bool Init(float3 conic, float opacity, float cached_limit,
219 EllipseBounds& out) {219 EllipseBounds& out) {
220 float limit = cached_limit;220 float limit = cached_limit;
221 if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;221 if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;
222 out.a = conic.x;222 out.a = conic.x;
223 out.b = conic.y;223 out.b = conic.y;
224 out.c = conic.z;224 out.c = conic.z;
225 out.det = out.a * out.c - out.b * out.b;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;226 if (!(out.det > 0.0f && out.a > 0.0f && out.c > 0.0f)) return false;
227 out.inv_a = 1.0f / out.a;227 out.inv_a = 1.0f / out.a;
228 out.a_limit = out.a * limit;228 out.a_limit = out.a * limit;
229 const float limit_over_det = limit / out.det;229 const float limit_over_det = limit / out.det;
230 out.y_ext = sqrtf(out.a * limit_over_det);230 out.y_ext = sqrtf(out.a * limit_over_det);
231 out.dy_peak = out.b * sqrtf(limit_over_det / out.c);231 out.dy_peak = out.b * sqrtf(limit_over_det / out.c);
232 return true;232 return true;
233 }233 }
234234
235 // Span [x0_px, x1_px] of integer pixel columns hit across integer pixel rows235 // 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.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,237 GS_D bool SpanForPixelRows(float2 xy, int y0_px, int y1_px, int& x0_px,
238 int& x1_px) const {238 int& x1_px) const {
239 const int ylo_i =239 const int ylo_i =
240 max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));240 max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));
241 const int yhi_i =241 const int yhi_i =
242 min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));242 min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));
243 if (ylo_i > yhi_i) return false;243 if (ylo_i > yhi_i) return false;
244 const float ylo = static_cast<float>(ylo_i) - xy.y;244 const float ylo = static_cast<float>(ylo_i) - xy.y;
245 const float yhi = static_cast<float>(yhi_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 right246 // 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.247 // (left) branch is concave (convex) and peaks at dy = -+dy_peak.
248 const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));248 const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));
249 const float dyl = 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;250 const float dyr2 = dyr * dyr, dyl2 = dyl * dyl;
251 const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));251 const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));
252 const float rl =252 const float rl =
253 (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));253 (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));
254 const float hi = (-b * dyr + rr) * inv_a;254 const float hi = (-b * dyr + rr) * inv_a;
255 const float lo = (-b * dyl - rl) * inv_a;255 const float lo = (-b * dyl - rl) * inv_a;
256 x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));256 x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));
257 x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));257 x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));
258 return x0_px <= x1_px;258 return x0_px <= x1_px;
259 }259 }
260};260};
261261
262// Visits every tile row `ty` with the half-open span [tx0, tx1) of tiles that262// 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).263// splat `s` (not culled) is binned to. Rows may be empty (tx0 >= tx1).
264template <Culling kCulling, typename Fn>264template <Culling kCulling, typename Fn>
265GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,265GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,
266 int height, Fn&& fn) {266 int height, Fn&& fn) {
267 const float2 xy = s.xy();267 const float2 xy = s.xy();
268 const TileRect rect = RectFor(xy, s.radius(), grid);268 const TileRect rect = RectFor(xy, s.radius(), grid);
269 if constexpr (kCulling == Culling::kRect) {269 if constexpr (kCulling == Culling::kRect) {
270 for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);270 for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);
271 } else {271 } else {
272 EllipseBounds eb;272 EllipseBounds eb;
273 if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;273 if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;
274 const int min_y =274 const int min_y =
275 max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));275 max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));
276 const int max_y =276 const int max_y =
277 min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));277 min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));
278 if (min_y > max_y) return;278 if (min_y > max_y) return;
279 const int ty0 = max(rect.y0, min_y / kTileSize);279 const int ty0 = max(rect.y0, min_y / kTileSize);
280 const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);280 const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);
281 for (int ty = ty0; ty < ty1; ++ty) {281 for (int ty = ty0; ty < ty1; ++ty) {
282 int x0_px, x1_px;282 int x0_px, x1_px;
283 if (!eb.SpanForPixelRows(xy, ty * kTileSize,283 if (!eb.SpanForPixelRows(xy, ty * kTileSize,
284 min((ty + 1) * kTileSize, height) - 1, x0_px,284 min((ty + 1) * kTileSize, height) - 1, x0_px,
285 x1_px)) {285 x1_px)) {
286 continue;286 continue;
287 }287 }
288 x0_px = max(0, x0_px);288 x0_px = max(0, x0_px);
289 x1_px = min(width - 1, x1_px);289 x1_px = min(width - 1, x1_px);
290 if (x0_px > x1_px) continue;290 if (x0_px > x1_px) continue;
291 const int tx0 = max(rect.x0, x0_px / kTileSize);291 const int tx0 = max(rect.x0, x0_px / kTileSize);
292 const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);292 const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);
293 if (tx0 < tx1) fn(ty, tx0, tx1);293 if (tx0 < tx1) fn(ty, tx0, tx1);
294 }294 }
295 }295 }
296}296}
297297
298namespace warp {298namespace warp {
299299
300inline constexpr unsigned kFullMask = 0xffffffffu;300inline constexpr unsigned kFullMask = 0xffffffffu;
301301
302// Reduce-scatter of an array of `N` <= 32 per-lane values over the warp.302// Reduce-scatter of an array of `N` <= 32 per-lane values over the warp.
303//303//
304// A butterfly where each step halves the payload: at offset `o`, lanes with304// A butterfly where each step halves the payload: at offset `o`, lanes with
305// bit `o` clear keep the low half of the array and send the high half, the305// bit `o` clear keep the low half of the array and send the high half, the
306// others do the opposite. After log2(32) steps every lane holds a single306// others do the opposite. After log2(32) steps every lane holds a single
307// value, v[0], which is the warp total of element `ScatterIndex(N, lane)`307// value, v[0], which is the warp total of element `ScatterIndex(N, lane)`
308// (or junk if that is -1). For N = 9 this takes 12 shuffles instead of the308// (or junk if that is -1). For N = 9 this takes 12 shuffles instead of the
309// 45 of nine independent butterflies, and the totals end in distinct lanes,309// 45 of nine independent butterflies, and the totals end in distinct lanes,
310// so writing them out is one instruction.310// so writing them out is one instruction.
311311
312namespace internal {312namespace internal {
313template <int N, int kOffset>313template <int N, int kOffset>
314GS_D void HalvingStep(float* v, int lane) {314GS_D void HalvingStep(float* v, int lane) {
315 constexpr int kHalf = (N + 1) / 2;315 constexpr int kHalf = (N + 1) / 2;
316 const bool upper = lane & kOffset;316 const bool upper = lane & kOffset;
317#pragma unroll317#pragma unroll
318 for (int i = 0; i < kHalf; ++i) {318 for (int i = 0; i < kHalf; ++i) {
319 const float lo = v[i];319 const float lo = v[i];
320 const float hi = i + kHalf < N ? v[i + kHalf] : 0.0f;320 const float hi = i + kHalf < N ? v[i + kHalf] : 0.0f;
321 const float send = upper ? lo : hi;321 const float send = upper ? lo : hi;
322 const float keep = upper ? hi : lo;322 const float keep = upper ? hi : lo;
323 v[i] = keep + __shfl_xor_sync(kFullMask, send, kOffset);323 v[i] = keep + __shfl_xor_sync(kFullMask, send, kOffset);
324 }324 }
325}325}
326326
327template <int N, int kOffset>327template <int N, int kOffset>
328GS_D void ReduceScatterImpl(float* v, int lane) {328GS_D void ReduceScatterImpl(float* v, int lane) {
329 if constexpr (kOffset > 0) {329 if constexpr (kOffset > 0) {
330 HalvingStep<N, kOffset>(v, lane);330 HalvingStep<N, kOffset>(v, lane);
331 ReduceScatterImpl<(N + 1) / 2, kOffset / 2>(v, lane);331 ReduceScatterImpl<(N + 1) / 2, kOffset / 2>(v, lane);
332 }332 }
333}333}
334} // namespace internal334} // namespace internal
335335
336template <int N>336template <int N>
337GS_D void ReduceScatter(float (&v)[N], int lane) {337GS_D void ReduceScatter(float (&v)[N], int lane) {
338 static_assert(N >= 1 && N <= 32);338 static_assert(N >= 1 && N <= 32);
339 internal::ReduceScatterImpl<N, 16>(v, lane);339 internal::ReduceScatterImpl<N, 16>(v, lane);
340}340}
341341
342template <int N>342template <int N>
343GS_D void ReduceScatter16(float (&v)[N], int lane) {343GS_D void ReduceScatter16(float (&v)[N], int lane) {
344 static_assert(N >= 1 && N <= 16);344 static_assert(N >= 1 && N <= 16);
345 internal::ReduceScatterImpl<N, 8>(v, lane);345 internal::ReduceScatterImpl<N, 8>(v, lane);
346}346}
347347
348// Element whose warp total lane `lane` holds after `ReduceScatter<N>`, or -1.348// Element whose warp total lane `lane` holds after `ReduceScatter<N>`, or -1.
349GS_HD constexpr int ScatterIndex(int n, int lane) {349GS_HD constexpr int ScatterIndex(int n, int lane) {
350 // Payload sizes and halves per step, offset 16 first.350 // Payload sizes and halves per step, offset 16 first.
351 int sizes[5] = {};351 int sizes[5] = {};
352 int size = n;352 int size = n;
353 for (int s = 0; s < 5; ++s) {353 for (int s = 0; s < 5; ++s) {
354 sizes[s] = size;354 sizes[s] = size;
355 size = (size + 1) / 2;355 size = (size + 1) / 2;
356 }356 }
357 int index = 0; // Slot in the final (size 1) array.357 int index = 0; // Slot in the final (size 1) array.
358 for (int s = 4; s >= 0; --s) {358 for (int s = 4; s >= 0; --s) {
359 const int half = (sizes[s] + 1) / 2;359 const int half = (sizes[s] + 1) / 2;
360 if (lane & (16 >> s)) index += half;360 if (lane & (16 >> s)) index += half;
361 if (index >= sizes[s]) return -1;361 if (index >= sizes[s]) return -1;
362 }362 }
363 return index;363 return index;
364}364}
365365
366// Element whose 16-lane half-warp total lane `lane` holds after366// Element whose 16-lane half-warp total lane `lane` holds after
367// `ReduceScatter16<N>`, or -1.367// `ReduceScatter16<N>`, or -1.
368GS_HD constexpr int ScatterIndex16(int n, int lane) {368GS_HD constexpr int ScatterIndex16(int n, int lane) {
369 const int half_lane = lane & 15;369 const int half_lane = lane & 15;
370 int sizes[4] = {};370 int sizes[4] = {};
371 int size = n;371 int size = n;
372 for (int s = 0; s < 4; ++s) {372 for (int s = 0; s < 4; ++s) {
373 sizes[s] = size;373 sizes[s] = size;
374 size = (size + 1) / 2;374 size = (size + 1) / 2;
375 }375 }
376 int index = 0;376 int index = 0;
377 for (int s = 3; s >= 0; --s) {377 for (int s = 3; s >= 0; --s) {
378 const int half = (sizes[s] + 1) / 2;378 const int half = (sizes[s] + 1) / 2;
379 if (half_lane & (8 >> s)) index += half;379 if (half_lane & (8 >> s)) index += half;
380 if (index >= sizes[s]) return -1;380 if (index >= sizes[s]) return -1;
381 }381 }
382 return index;382 return index;
383}383}
384384
385} // namespace warp385} // namespace warp
386386
387inline constexpr int kShCoeffs = 16; // Degree 3.387inline constexpr int kShCoeffs = 16; // Degree 3.
388inline constexpr float kNearPlane = 0.2f;388inline constexpr float kNearPlane = 0.2f;
389inline constexpr float kLowPass = 0.3f;389inline constexpr float kLowPass = 0.3f;
390inline constexpr float kFovClamp = 1.3f;390inline constexpr float kFovClamp = 1.3f;
391391
392// Pointers to the trainable parameters, all activated, row-major.392// Pointers to the trainable parameters, all activated, row-major.
393struct GaussianParams {393struct GaussianParams {
394 const float* means; // (N, 3)394 const float* means; // (N, 3)
395 const float* scales; // (N, 3)395 const float* scales; // (N, 3)
396 const float* quats; // (N, 4) as (w, x, y, z), not normalized.396 const float* quats; // (N, 4) as (w, x, y, z), not normalized.
397 const float* opacities; // (N,)397 const float* opacities; // (N,)
398 const void* sh; // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).398 const void* sh; // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).
399 int n;399 int n;
400};400};
401401
402// Gradients w.r.t. `GaussianParams`, same layout and SH storage type.402// Gradients w.r.t. `GaussianParams`, same layout and SH storage type.
403struct GaussianGrads {403struct GaussianGrads {
404 float* means;404 float* means;
405 float* scales;405 float* scales;
406 float* quats;406 float* quats;
407 float* opacities;407 float* opacities;
408 void* sh;408 void* sh;
409};409};
410410
411// Screen-space quantities needed by rasterization. `conic` is the inverse 2D411// Screen-space quantities needed by rasterization. `conic` is the inverse 2D
412// covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.412// covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.
413struct Projection {413struct Projection {
414 float2 xy;414 float2 xy;
415 float depth;415 float depth;
416 float3 conic;416 float3 conic;
417 float opacity;417 float opacity;
418 float limit;418 float limit;
419 float3 rgb;419 float3 rgb;
420 int radius; // 0 means culled.420 int radius; // 0 means culled.
421};421};
422422
423// ----------------------------------------------------------------------------423// ----------------------------------------------------------------------------
424// SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian424// SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian
425// and is accessed in 4 groups of 4 coefficients (12 values). The coefficients425// and is accessed in 4 groups of 4 coefficients (12 values). The coefficients
426// may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.426// may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.
427427
428using ShBf16 = uint16_t; // bfloat16 bits.428using ShBf16 = uint16_t; // bfloat16 bits.
429429
430template <typename T>430template <typename T>
431struct ShIo;431struct ShIo;
432432
433template <>433template <>
434struct ShIo<float> {434struct ShIo<float> {
435 static constexpr int kRowVec16 = 12; // 16-byte vectors per row.435 static constexpr int kRowVec16 = 12; // 16-byte vectors per row.
436 GS_D static void LoadGroup(const float* row, int grp, float v[12]) {436 GS_D static void LoadGroup(const float* row, int grp, float v[12]) {
437 const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;437 const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;
438 const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);438 const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);
439 v[0] = a.x, v[1] = a.y, v[2] = a.z, v[3] = a.w;439 v[0] = a.x, v[1] = a.y, v[2] = a.z, v[3] = a.w;
440 v[4] = b.x, v[5] = b.y, v[6] = b.z, v[7] = b.w;440 v[4] = b.x, v[5] = b.y, v[6] = b.z, v[7] = b.w;
441 v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;441 v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;
442 }442 }
443 GS_D static void StoreGroup(float* row, int grp, const float v[12]) {443 GS_D static void StoreGroup(float* row, int grp, const float v[12]) {
444 float4* p = reinterpret_cast<float4*>(row) + 3 * grp;444 float4* p = reinterpret_cast<float4*>(row) + 3 * grp;
445 p[0] = {v[0], v[1], v[2], v[3]};445 p[0] = {v[0], v[1], v[2], v[3]};
446 p[1] = {v[4], v[5], v[6], v[7]};446 p[1] = {v[4], v[5], v[6], v[7]};
447 p[2] = {v[8], v[9], v[10], v[11]};447 p[2] = {v[8], v[9], v[10], v[11]};
448 }448 }
449};449};
450450
451template <>451template <>
452struct ShIo<ShBf16> {452struct ShIo<ShBf16> {
453 static constexpr int kRowVec16 = 6;453 static constexpr int kRowVec16 = 6;
454 GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }454 GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
455 GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }455 GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
456 // Round to nearest even (inputs are finite).456 // Round to nearest even (inputs are finite).
457 GS_D static uint32_t Round(float f) {457 GS_D static uint32_t Round(float f) {
458 const uint32_t u = __float_as_uint(f);458 const uint32_t u = __float_as_uint(f);
459 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;459 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
460 }460 }
461 GS_D static uint32_t Pack(float lo, float hi) {461 GS_D static uint32_t Pack(float lo, float hi) {
462 return Round(lo) | (Round(hi) << 16);462 return Round(lo) | (Round(hi) << 16);
463 }463 }
464 GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {464 GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {
465 // 24-byte groups: three 8-byte loads.465 // 24-byte groups: three 8-byte loads.
466 const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;466 const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;
467#pragma unroll467#pragma unroll
468 for (int k = 0; k < 3; ++k) {468 for (int k = 0; k < 3; ++k) {
469 const uint2 w = __ldg(p + k);469 const uint2 w = __ldg(p + k);
470 v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);470 v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);
471 v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);471 v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);
472 }472 }
473 }473 }
474 GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {474 GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {
475 uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;475 uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;
476#pragma unroll476#pragma unroll
477 for (int k = 0; k < 3; ++k) {477 for (int k = 0; k < 3; ++k) {
478 p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};478 p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};
479 }479 }
480 }480 }
481};481};
482482
483template <typename ShT>483template <typename ShT>
484GS_D const ShT* ShRow(const void* sh, int i) {484GS_D const ShT* ShRow(const void* sh, int i) {
485 return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;485 return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
486}486}
487template <typename ShT>487template <typename ShT>
488GS_D ShT* ShRow(void* sh, int i) {488GS_D ShT* ShRow(void* sh, int i) {
489 return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;489 return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
490}490}
491491
492// ----------------------------------------------------------------------------492// ----------------------------------------------------------------------------
493// Spherical harmonics (degree 3, Inria sign convention).493// Spherical harmonics (degree 3, Inria sign convention).
494494
495namespace sh {495namespace sh {
496inline constexpr float kC0 = 0.28209479177387814f;496inline constexpr float kC0 = 0.28209479177387814f;
497inline constexpr float kC1 = 0.4886025119029199f;497inline constexpr float kC1 = 0.4886025119029199f;
498inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,498inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
499 0.31539156525252005f, -1.0925484305920792f,499 0.31539156525252005f, -1.0925484305920792f,
500 0.5462742152960396f};500 0.5462742152960396f};
501inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,501inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
502 -0.4570457994644658f, 0.3731763325901154f,502 -0.4570457994644658f, 0.3731763325901154f,
503 -0.4570457994644658f, 1.445305721320277f,503 -0.4570457994644658f, 1.445305721320277f,
504 -0.5900435899266435f};504 -0.5900435899266435f};
505505
506// Basis values for a unit direction.506// Basis values for a unit direction.
507GS_D void Basis(float3 d, float b[kShCoeffs]) {507GS_D void Basis(float3 d, float b[kShCoeffs]) {
508 const float x = d.x, y = d.y, z = d.z;508 const float x = d.x, y = d.y, z = d.z;
509 const float xx = x * x, yy = y * y, zz = z * z;509 const float xx = x * x, yy = y * y, zz = z * z;
510 b[0] = kC0;510 b[0] = kC0;
511 b[1] = -kC1 * y;511 b[1] = -kC1 * y;
512 b[2] = kC1 * z;512 b[2] = kC1 * z;
513 b[3] = -kC1 * x;513 b[3] = -kC1 * x;
514 b[4] = kC2[0] * x * y;514 b[4] = kC2[0] * x * y;
515 b[5] = kC2[1] * y * z;515 b[5] = kC2[1] * y * z;
516 b[6] = kC2[2] * (2.0f * zz - xx - yy);516 b[6] = kC2[2] * (2.0f * zz - xx - yy);
517 b[7] = kC2[3] * x * z;517 b[7] = kC2[3] * x * z;
518 b[8] = kC2[4] * (xx - yy);518 b[8] = kC2[4] * (xx - yy);
519 b[9] = kC3[0] * y * (3.0f * xx - yy);519 b[9] = kC3[0] * y * (3.0f * xx - yy);
520 b[10] = kC3[1] * x * y * z;520 b[10] = kC3[1] * x * y * z;
521 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);521 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
522 b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);522 b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
523 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);523 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
524 b[14] = kC3[5] * z * (xx - yy);524 b[14] = kC3[5] * z * (xx - yy);
525 b[15] = kC3[6] * x * (xx - 3.0f * yy);525 b[15] = kC3[6] * x * (xx - 3.0f * yy);
526}526}
527527
528// Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.528// Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.
529template <typename ShT>529template <typename ShT>
530GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {530GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
531 float3 c = {0.5f, 0.5f, 0.5f};531 float3 c = {0.5f, 0.5f, 0.5f};
532#pragma unroll532#pragma unroll
533 for (int grp = 0; grp < 4; ++grp) {533 for (int grp = 0; grp < 4; ++grp) {
534 const int k0 = 4 * grp;534 const int k0 = 4 * grp;
535 float v[12];535 float v[12];
536 ShIo<ShT>::LoadGroup(sh_row, grp, v);536 ShIo<ShT>::LoadGroup(sh_row, grp, v);
537 c += b[k0] * float3{v[0], v[1], v[2]};537 c += b[k0] * float3{v[0], v[1], v[2]};
538 c += b[k0 + 1] * float3{v[3], v[4], v[5]};538 c += b[k0 + 1] * float3{v[3], v[4], v[5]};
539 c += b[k0 + 2] * float3{v[6], v[7], v[8]};539 c += b[k0 + 2] * float3{v[6], v[7], v[8]};
540 c += b[k0 + 3] * float3{v[9], v[10], v[11]};540 c += b[k0 + 3] * float3{v[9], v[10], v[11]};
541 }541 }
542 return c;542 return c;
543}543}
544} // namespace sh544} // namespace sh
545545
546// ----------------------------------------------------------------------------546// ----------------------------------------------------------------------------
547// Covariance helpers.547// Covariance helpers.
548548
549// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).549// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
550GS_D Mat3 QuatToRotation(float4 q) {550GS_D Mat3 QuatToRotation(float4 q) {
551 const float w = q.x, x = q.y, y = q.z, z = q.w;551 const float w = q.x, x = q.y, y = q.z, z = q.w;
552 Mat3 r;552 Mat3 r;
553 r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);553 r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
554 r.m[0][1] = 2.0f * (x * y - w * z);554 r.m[0][1] = 2.0f * (x * y - w * z);
555 r.m[0][2] = 2.0f * (x * z + w * y);555 r.m[0][2] = 2.0f * (x * z + w * y);
556 r.m[1][0] = 2.0f * (x * y + w * z);556 r.m[1][0] = 2.0f * (x * y + w * z);
557 r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);557 r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
558 r.m[1][2] = 2.0f * (y * z - w * x);558 r.m[1][2] = 2.0f * (y * z - w * x);
559 r.m[2][0] = 2.0f * (x * z - w * y);559 r.m[2][0] = 2.0f * (x * z - w * y);
560 r.m[2][1] = 2.0f * (y * z + w * x);560 r.m[2][1] = 2.0f * (y * z + w * x);
561 r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);561 r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
562 return r;562 return r;
563}563}
564564
565// Symmetric 3x3 matrix.565// Symmetric 3x3 matrix.
566struct Sym3 {566struct Sym3 {
567 float xx, xy, xz, yy, yz, zz;567 float xx, xy, xz, yy, yz, zz;
568 GS_D float3 Mul(float3 v) const {568 GS_D float3 Mul(float3 v) const {
569 return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,569 return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
570 xz * v.x + yz * v.y + zz * v.z};570 xz * v.x + yz * v.y + zz * v.z};
571 }571 }
572};572};
573573
574// Sigma = A A^T with A = R diag(s).574// Sigma = A A^T with A = R diag(s).
575GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {575GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
576 Mat3 a;576 Mat3 a;
577 const float sv[3] = {s.x, s.y, s.z};577 const float sv[3] = {s.x, s.y, s.z};
578#pragma unroll578#pragma unroll
579 for (int i = 0; i < 3; ++i)579 for (int i = 0; i < 3; ++i)
580#pragma unroll580#pragma unroll
581 for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];581 for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
582 return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),582 return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
583 dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),583 dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
584 dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};584 dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
585}585}
586586
587// Intermediate values of the projection shared by forward and VJP.587// Intermediate values of the projection shared by forward and VJP.
588struct ProjectionState {588struct ProjectionState {
589 float3 t; // View-space mean.589 float3 t; // View-space mean.
590 float tx_c, ty_c; // Clamped x, y used by the Jacobian.590 float tx_c, ty_c; // Clamped x, y used by the Jacobian.
591 bool x_clamped, y_clamped;591 bool x_clamped, y_clamped;
592 float3 m0, m1; // Rows of M = J W (2x3).592 float3 m0, m1; // Rows of M = J W (2x3).
593 Sym3 sigma; // World-space 3D covariance.593 Sym3 sigma; // World-space 3D covariance.
594 float a, b, c; // Dilated 2D covariance.594 float a, b, c; // Dilated 2D covariance.
595};595};
596596
597// Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).597// Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
598GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,598GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
599 float4 quat, ProjectionState& st) {599 float4 quat, ProjectionState& st) {
600 st.t = t;600 st.t = t;
601 const float tz = t.z;601 const float tz = t.z;
602 const float limx = kFovClamp * cam.TanHalfFovX();602 const float limx = kFovClamp * cam.TanHalfFovX();
603 const float limy = kFovClamp * cam.TanHalfFovY();603 const float limy = kFovClamp * cam.TanHalfFovY();
604 const float txtz = t.x / tz, tytz = t.y / tz;604 const float txtz = t.x / tz, tytz = t.y / tz;
605 st.x_clamped = txtz < -limx || txtz > limx;605 st.x_clamped = txtz < -limx || txtz > limx;
606 st.y_clamped = tytz < -limy || tytz > limy;606 st.y_clamped = tytz < -limy || tytz > limy;
607 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;607 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
608 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;608 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
609 const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);609 const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
610 const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);610 const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
611 const Mat3& w = cam.rotation;611 const Mat3& w = cam.rotation;
612 st.m0 = j00 * w.row(0) + j02 * w.row(2);612 st.m0 = j00 * w.row(0) + j02 * w.row(2);
613 st.m1 = j11 * w.row(1) + j12 * w.row(2);613 st.m1 = j11 * w.row(1) + j12 * w.row(2);
614 st.sigma = Covariance3D(QuatToRotation(quat), scale);614 st.sigma = Covariance3D(QuatToRotation(quat), scale);
615 const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);615 const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
616 st.a = dot(st.m0, s0) + kLowPass;616 st.a = dot(st.m0, s0) + kLowPass;
617 st.b = dot(st.m0, s1);617 st.b = dot(st.m0, s1);
618 st.c = dot(st.m1, s1) + kLowPass;618 st.c = dot(st.m1, s1) + kLowPass;
619}619}
620620
621// Computes everything up to the 2D covariance. Returns false if culled by the621// Computes everything up to the 2D covariance. Returns false if culled by the
622// near plane.622// near plane.
623GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,623GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
624 float4 quat, ProjectionState& st) {624 float4 quat, ProjectionState& st) {
625 const float3 t = cam.ToView(mean);625 const float3 t = cam.ToView(mean);
626 if (t.z <= kNearPlane) return false;626 if (t.z <= kNearPlane) return false;
627 ProjectCovarianceFromView(cam, t, scale, quat, st);627 ProjectCovarianceFromView(cam, t, scale, quat, st);
628 return true;628 return true;
629}629}
630630
631// Integer bounding radius (pixels) of the 2D covariance, as in the reference.631// Integer bounding radius (pixels) of the 2D covariance, as in the reference.
632GS_D int BoundingRadius(float a, float b, float c) {632GS_D int BoundingRadius(float a, float b, float c) {
633 const float det = a * c - b * b;633 const float det = a * c - b * b;
634 const float mid = 0.5f * (a + c);634 const float mid = 0.5f * (a + c);
635 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));635 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
636 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));636 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
637}637}
638638
639GS_D float2 ProjectMean(const Camera& cam, float3 t) {639GS_D float2 ProjectMean(const Camera& cam, float3 t) {
640 return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};640 return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
641}641}
642642
643GS_D float3 LoadMean(const GaussianParams& p, int i) {643GS_D float3 LoadMean(const GaussianParams& p, int i) {
644 return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};644 return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
645}645}
646646
647// Unit direction from the camera center to `mean`, and 1 / distance.647// Unit direction from the camera center to `mean`, and 1 / distance.
648GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {648GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
649 const float3 v = mean - cam.position;649 const float3 v = mean - cam.position;
650 inv_len = rsqrtf(dot(v, v));650 inv_len = rsqrtf(dot(v, v));
651 return inv_len * v;651 return inv_len * v;
652}652}
653653
654// Gradient of the loss w.r.t. one splat's screen-space quantities.654// Gradient of the loss w.r.t. one splat's screen-space quantities.
655struct ProjectionGrad {655struct ProjectionGrad {
656 float2 xy;656 float2 xy;
657 float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.657 float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.
658 float opacity;658 float opacity;
659 float3 rgb;659 float3 rgb;
660};660};
661661
662// Image geometry passed by value to every kernel.662// Image geometry passed by value to every kernel.
663struct Frame {663struct Frame {
664 const float* camera; // Device pointer to the packed camera.664 const float* camera; // Device pointer to the packed camera.
665 int width, height;665 int width, height;
666 TileGrid grid;666 TileGrid grid;
667 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }667 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
668};668};
669669
670// ---------------------------------------------------------------------------670// ---------------------------------------------------------------------------
671// Per-Gaussian kernels run one thread per Gaussian.671// Per-Gaussian kernels run one thread per Gaussian.
672672
673inline constexpr int kGaussianThreads = 128;673inline constexpr int kGaussianThreads = 128;
674674
675inline int GaussianBlocks(int64_t n) {675inline int GaussianBlocks(int64_t n) {
676 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);676 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
677}677}
678678
679// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.679// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
680inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2680inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
681inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;681inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
682682
683GS_HD float3 Power2Conic(float a, float b, float c) {683GS_HD float3 Power2Conic(float a, float b, float c) {
684 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};684 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
685}685}
686686
687GS_D float Exp2Approx(float x) {687GS_D float Exp2Approx(float x) {
688 float y;688 float y;
689 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));689 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
690 return y;690 return y;
691}691}
692692
693// ---------------------------------------------------------------------------693// ---------------------------------------------------------------------------
694// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into694// Rasterization. One block per 16x16 tile. Each thread shades `kRows`
695// vertically adjacent pixels, so that per-splat costs (shared-memory loads,
696// barriers, the warp reduction in the backward pass) are paid once per
697// `kRows` pixels; more rows per thread also means fewer resident threads and
698// more waiting on the slowest pixel. A warp covers a compact
699// kWarpWidth x kWarpRows region (as square as possible, so per-warp culling
695// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through700// is tight); warps tile the 16x16 tile row-major. Splats are staged through
696// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.701// shared memory, one per thread per batch.
702
703inline constexpr int kRasterRowChoices[] = {1, 2, 4};
697704
698template <int kRows>705template <int kRows>
699struct RasterLayout {706struct RasterLayout {
700 static_assert(kRows == 1 || kRows == 2 || kRows == 4);707 static_assert(kRows == 1 || kRows == 2 || kRows == 4);
701 static constexpr int kThreads = kTilePixels / kRows;708 static constexpr int kThreads = kTilePixels / kRows;
702 static constexpr int kWarps = kThreads / 32;709 static constexpr int kWarps = kThreads / 32;
703 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;710 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
704 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;711 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
705 static constexpr int kWarpsX = kTileSize / kWarpWidth;712 static constexpr int kWarpsX = kTileSize / kWarpWidth;
706 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);713 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
707714
715 // Tile-relative origin of warp `w`'s region.
708 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }716 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
709 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }717 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
710};718};
711719
720// Pixel coordinates of this thread's first pixel; the others follow in y.
712struct RasterPixels {721struct RasterPixels {
713 int x, y0;722 int x, y0;
714 template <int kRows>723 template <int kRows>
715 GS_D static RasterPixels For(int tid) {724 GS_D static RasterPixels For(int tid) {
↕ Expand 88 unchanged lines (lines 725–812) · in WarpMask
716 using L = RasterLayout<kRows>;725 using L = RasterLayout<kRows>;
717 const int warp = tid / 32, lane = tid % 32;726 const int warp = tid / 32, lane = tid % 32;
718 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +727 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
719 lane % L::kWarpWidth,728 lane % L::kWarpWidth,
720 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +729 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
721 (lane / L::kWarpWidth) * kRows};730 (lane / L::kWarpWidth) * kRows};
722 }731 }
723};732};
724733
725// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the734// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
726// pixel region of warp w of this block's tile.735// pixel region of warp w of this block's tile.
727template <int kRows>736template <int kRows>
728GS_D uint32_t WarpMask(float4 p0, float4 p1) {737GS_D uint32_t WarpMask(float4 p0, float4 p1) {
729 using L = RasterLayout<kRows>;738 using L = RasterLayout<kRows>;
730 constexpr int kWarpsY = L::kWarps / L::kWarpsX;739 constexpr int kWarpsY = L::kWarps / L::kWarpsX;
731 float limit = p1.z;740 float limit = p1.z;
732 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;741 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
733 const float a = p0.z, b = p0.w, c = p1.x;742 const float a = p0.z, b = p0.w, c = p1.x;
734 const float det = a * c - b * b;743 const float det = a * c - b * b;
735 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;744 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
736 const float l_over_det = limit / det;745 const float l_over_det = limit / det;
737 const float rx = sqrtf(c * l_over_det);746 const float rx = sqrtf(c * l_over_det);
738 const float ry = sqrtf(a * l_over_det);747 const float ry = sqrtf(a * l_over_det);
739 const float2 xy = {p0.x, p0.y};748 const float2 xy = {p0.x, p0.y};
740 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));749 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
741 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));750 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
742 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));751 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
743 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));752 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
744 if (x0 > x1 || y0 > y1) return 0;753 if (x0 > x1 || y0 > y1) return 0;
745 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;754 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
746 uint32_t mx = 0u;755 uint32_t mx = 0u;
747#pragma unroll756#pragma unroll
748 for (int wx = 0; wx < L::kWarpsX; ++wx) {757 for (int wx = 0; wx < L::kWarpsX; ++wx) {
749 const int wx0 = tx + wx * L::kWarpWidth;758 const int wx0 = tx + wx * L::kWarpWidth;
750 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);759 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
751 }760 }
752 if (mx == 0u) return 0;761 if (mx == 0u) return 0;
753 uint32_t my = 0u;762 uint32_t my = 0u;
754#pragma unroll763#pragma unroll
755 for (int wy = 0; wy < kWarpsY; ++wy) {764 for (int wy = 0; wy < kWarpsY; ++wy) {
756 const int wy0 = ty + wy * L::kWarpRows;765 const int wy0 = ty + wy * L::kWarpRows;
757 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);766 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
758 }767 }
759 if (my == 0u) return 0;768 if (my == 0u) return 0;
760 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {769 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
761 uint32_t mask = 0u;770 uint32_t mask = 0u;
762#pragma unroll771#pragma unroll
763 for (int wy = 0; wy < kWarpsY; ++wy) {772 for (int wy = 0; wy < kWarpsY; ++wy) {
764 if (!((my >> wy) & 1u)) continue;773 if (!((my >> wy) & 1u)) continue;
765#pragma unroll774#pragma unroll
766 for (int wx = 0; wx < L::kWarpsX; ++wx) {775 for (int wx = 0; wx < L::kWarpsX; ++wx) {
767 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));776 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
768 }777 }
769 }778 }
770 return mask;779 return mask;
771 }780 }
772 EllipseBounds eb;781 EllipseBounds eb;
773 eb.a = a;782 eb.a = a;
774 eb.b = b;783 eb.b = b;
775 eb.c = c;784 eb.c = c;
776 eb.det = det;785 eb.det = det;
777 eb.inv_a = 1.0f / a;786 eb.inv_a = 1.0f / a;
778 eb.a_limit = a * limit;787 eb.a_limit = a * limit;
779 eb.y_ext = ry;788 eb.y_ext = ry;
780 eb.dy_peak = b * sqrtf(l_over_det / c);789 eb.dy_peak = b * sqrtf(l_over_det / c);
781 uint32_t mask = 0u;790 uint32_t mask = 0u;
782#pragma unroll791#pragma unroll
783 for (int wy = 0; wy < kWarpsY; ++wy) {792 for (int wy = 0; wy < kWarpsY; ++wy) {
784 if (!((my >> wy) & 1u)) continue;793 if (!((my >> wy) & 1u)) continue;
785 const int wy0 = ty + wy * L::kWarpRows;794 const int wy0 = ty + wy * L::kWarpRows;
786 int x0_px, x1_px;795 int x0_px, x1_px;
787 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {796 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
788 continue;797 continue;
789 }798 }
790#pragma unroll799#pragma unroll
791 for (int wx = 0; wx < L::kWarpsX; ++wx) {800 for (int wx = 0; wx < L::kWarpsX; ++wx) {
792 const int wx0 = tx + wx * L::kWarpWidth;801 const int wx0 = tx + wx * L::kWarpWidth;
793 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {802 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
794 mask |= 1u << (wy * L::kWarpsX + wx);803 mask |= 1u << (wy * L::kWarpsX + wx);
795 }804 }
796 }805 }
797 }806 }
798 return mask;807 return mask;
799}808}
800809
801} // namespace gs810} // namespace gs
802811
803#endif812#endif
↕ Expand 689 unchanged lines (lines 1–689) · in Exp2Approx
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
1414
1515 #include <cuda_runtime.h>
1616
1717 #include <cmath>
1818 #include <cstdint>
1919
2020 #include "config.h"
2121 #include "stdio.h"
2222
2323 #define BLOCK_SIZE (BLOCK_X * BLOCK_Y)
2424 #define NUM_WARPS (BLOCK_SIZE / 32)
2525
2626 // Spherical harmonics coefficients
2727 __device__ const float SH_C0 = 0.28209479177387814f;
2828 __device__ const float SH_C1 = 0.4886025119029199f;
2929 __device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,
3030 0.31539156525252005f, -1.0925484305920792f,
3131 0.5462742152960396f};
3232 __device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,
3333 -0.4570457994644658f, 0.3731763325901154f,
3434 -0.4570457994644658f, 1.445305721320277f,
3535 -0.5900435899266435f};
3636
3737 namespace gs {
3838
3939 #define GS_HD __host__ __device__ __forceinline__
4040 #define GS_D __device__ __forceinline__
4141
4242 GS_HD float3 operator+(float3 a, float3 b) {
4343 return {a.x + b.x, a.y + b.y, a.z + b.z};
4444 }
4545 GS_HD float3 operator-(float3 a, float3 b) {
4646 return {a.x - b.x, a.y - b.y, a.z - b.z};
4747 }
4848 GS_HD float3 operator*(float s, float3 a) {
4949 return {s * a.x, s * a.y, s * a.z};
5050 }
5151 GS_HD float3& operator+=(float3& a, float3 b) {
5252 a.x += b.x;
5353 a.y += b.y;
5454 a.z += b.z;
5555 return a;
5656 }
5757 GS_HD float dot(float3 a, float3 b) {
5858 return a.x * b.x + a.y * b.y + a.z * b.z;
5959 }
6060
6161 // Row-major 3x3 matrix.
6262 struct Mat3 {
6363 float m[3][3];
6464 GS_HD float3 row(int i) const { return {m[i][0], m[i][1], m[i][2]}; }
6565 GS_HD float3 col(int j) const { return {m[0][j], m[1][j], m[2][j]}; }
6666 };
6767
6868 GS_HD float3 Mul(const Mat3& a, float3 v) {
6969 return {dot(a.row(0), v), dot(a.row(1), v), dot(a.row(2), v)};
7070 }
7171 // a^T v.
7272 GS_HD float3 MulT(const Mat3& a, float3 v) {
7373 return {dot(a.col(0), v), dot(a.col(1), v), dot(a.col(2), v)};
7474 }
7575
7676 // Reinterprets bits; used to pack small integers into float records.
7777 GS_HD float AsFloat(uint32_t u) {
7878 union {
7979 uint32_t u;
8080 float f;
8181 } c{u};
8282 return c.f;
8383 }
8484 GS_HD uint32_t AsUint(float f) {
8585 union {
8686 float f;
8787 uint32_t u;
8888 } c{f};
8989 return c.u;
9090 }
9191
9292 // Number of floats in the packed camera buffer built by `camera.py`:
9393 // [R (9, row-major world-to-camera), t (3), fx, fy, cx, cy, position (3), 0].
9494 inline constexpr int kCameraFloats = 20;
9595
9696 // Pinhole camera with OpenCV axes (+x right, +y down, +z forward). Pixel (i, j)
9797 // is sampled at coordinate (j, i), as in the Inria rasterizer.
9898 struct Camera {
9999 Mat3 rotation; // World-to-camera rotation.
100100 float3 translation;
101101 float fx, fy, cx, cy;
102102 float3 position; // Camera center in world space.
103103 int width, height;
104104
105105 GS_HD static Camera Load(const float* p, int width, int height) {
106106 Camera c;
107107 for (int i = 0; i < 3; ++i)
108108 for (int j = 0; j < 3; ++j) c.rotation.m[i][j] = p[3 * i + j];
109109 c.translation = {p[9], p[10], p[11]};
110110 c.fx = p[12];
111111 c.fy = p[13];
112112 c.cx = p[14];
113113 c.cy = p[15];
114114 c.position = {p[16], p[17], p[18]};
115115 c.width = width;
116116 c.height = height;
117117 return c;
118118 }
119119
120120 GS_HD float3 ToView(float3 p) const { return Mul(rotation, p) + translation; }
121121 // Tangent of the half field of view, as the Inria rasterizer defines it.
122122 GS_HD float TanHalfFovX() const { return width / (2.0f * fx); }
123123 GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }
124124 };
125125
126126 // Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block
127127 // rasterizes one tile.
128128 inline constexpr int kTileSize = 16;
129129 inline constexpr int kTilePixels = kTileSize * kTileSize;
130130
131131 struct TileGrid {
132132 int tiles_x, tiles_y;
133133 GS_HD static TileGrid For(int width, int height) {
134134 return {(width + kTileSize - 1) / kTileSize,
135135 (height + kTileSize - 1) / kTileSize};
136136 }
137137 GS_HD int num_tiles() const { return tiles_x * tiles_y; }
138138 };
139139
140140 // Smallest alpha that contributes to a pixel; smaller ones are skipped.
141141 inline constexpr float kMinAlpha = 1.0f / 255.0f;
142142 inline constexpr float kMaxAlpha = 0.99f;
143143 // Transmittance below which a pixel stops accumulating.
144144 inline constexpr float kMinTransmittance = 1e-4f;
145145
146146 // One splat as stored in memory: three float4 so the rasterizer issues three
147147 // 16-byte loads. Layout:
148148 // p0 = (x, y, conic.a, conic.b)
149149 // p1 = (conic.c, opacity, limit, radius) radius 0 => culled
150150 // p2 = (r, g, b, min_power)
151151 inline constexpr int kSplatFloats = 12;
152152 struct alignas(16) Splat {
153153 float4 p0, p1, p2;
154154 GS_HD float2 xy() const { return {p0.x, p0.y}; }
155155 GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }
156156 GS_HD float opacity() const { return p1.y; }
157157 GS_HD float limit() const { return p1.z; }
158158 GS_HD int radius() const { return static_cast<int>(p1.w); }
159159 GS_HD float3 rgb() const { return {p2.x, p2.y, p2.z}; }
160160 };
161161
162162 // Gradient w.r.t. a splat's rasterized quantities, accumulated by the raster
163163 // backward pass: (dx, dy, da, db, dc, dopacity, dr, dg, db).
164164 inline constexpr int kSplatGradFloats = 9;
165165
166166 // Half-open tile rectangle [x0, x1) x [y0, y1).
167167 struct TileRect {
168168 int x0, y0, x1, y1;
169169 GS_HD int area() const { return (x1 - x0) * (y1 - y0); }
170170 };
171171
172172 // Tile rectangle of a splat of integer bounding `radius`, identical to the
173173 // reference rasterizer.
174174 GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {
175175 const float r = static_cast<float>(radius);
176176 auto lo = [](float v, int n) {
177177 return min(n, max(0, static_cast<int>(v / kTileSize)));
178178 };
179179 auto hi = [](float v, int n) {
180180 return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));
181181 };
182182 return {lo(xy.x - r, grid.tiles_x), lo(xy.y - r, grid.tiles_y),
183183 hi(xy.x + r, grid.tiles_x), hi(xy.y + r, grid.tiles_y)};
184184 }
185185
186186 // How splats are assigned to tiles.
187187 enum class Culling : int {
188188 // Every tile of the bounding rectangle (reference behaviour).
189189 kRect = 0,
190190 // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives
191191 // the same image: skipped (splat, tile) pairs never contribute.
192192 kEllipse = 1,
193193 };
194194
195195 // alpha = o exp(-q/2) >= 1/255 <=> q <= 2 ln(255 o). Writes the bound on q
196196 // with margins that keep tests conservative under float rounding of the
197197 // rasterizer's power. Returns false if no pixel can reach 1/255.
198198 GS_HD bool AlphaLimit(float opacity, float& limit) {
199199 const float o255 = opacity * 255.0f;
200200 if (!(o255 >= 0.999f)) return false;
201201 limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;
202202 return true;
203203 }
204204
205205 // Pad (pixels) on ellipse extents, against rounding.
206206 inline constexpr float kSpanPadPx = 0.01f;
207207
208208 // Precomputed splat-level ellipse constants for fast band intersection tests
209209 // of {q(d) <= limit}, where q(d) = a dx^2 + 2 b dx dy + c dy^2.
210210 struct EllipseBounds {
211211 float a, b, c;
212212 float det;
213213 float inv_a;
214214 float a_limit;
215215 float y_ext;
216216 float dy_peak;
217217
218218 GS_D static bool Init(float3 conic, float opacity, float cached_limit,
219219 EllipseBounds& out) {
220220 float limit = cached_limit;
221221 if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;
222222 out.a = conic.x;
223223 out.b = conic.y;
224224 out.c = conic.z;
225225 out.det = out.a * out.c - out.b * out.b;
226226 if (!(out.det > 0.0f && out.a > 0.0f && out.c > 0.0f)) return false;
227227 out.inv_a = 1.0f / out.a;
228228 out.a_limit = out.a * limit;
229229 const float limit_over_det = limit / out.det;
230230 out.y_ext = sqrtf(out.a * limit_over_det);
231231 out.dy_peak = out.b * sqrtf(limit_over_det / out.c);
232232 return true;
233233 }
234234
235235 // Span [x0_px, x1_px] of integer pixel columns hit across integer pixel rows
236236 // Y in [y0_px, y1_px] (inclusive). Returns false if no integer pixel hits.
237237 GS_D bool SpanForPixelRows(float2 xy, int y0_px, int y1_px, int& x0_px,
238238 int& x1_px) const {
239239 const int ylo_i =
240240 max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));
241241 const int yhi_i =
242242 min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));
243243 if (ylo_i > yhi_i) return false;
244244 const float ylo = static_cast<float>(ylo_i) - xy.y;
245245 const float yhi = static_cast<float>(yhi_i) - xy.y;
246246 // Boundary x(dy) = (-b dy +- sqrt(a limit - det dy^2)) / a; the right
247247 // (left) branch is concave (convex) and peaks at dy = -+dy_peak.
248248 const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));
249249 const float dyl = fminf(yhi, fmaxf(ylo, dy_peak));
250250 const float dyr2 = dyr * dyr, dyl2 = dyl * dyl;
251251 const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));
252252 const float rl =
253253 (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));
254254 const float hi = (-b * dyr + rr) * inv_a;
255255 const float lo = (-b * dyl - rl) * inv_a;
256256 x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));
257257 x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));
258258 return x0_px <= x1_px;
259259 }
260260 };
261261
262262 // Visits every tile row `ty` with the half-open span [tx0, tx1) of tiles that
263263 // splat `s` (not culled) is binned to. Rows may be empty (tx0 >= tx1).
264264 template <Culling kCulling, typename Fn>
265265 GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,
266266 int height, Fn&& fn) {
267267 const float2 xy = s.xy();
268268 const TileRect rect = RectFor(xy, s.radius(), grid);
269269 if constexpr (kCulling == Culling::kRect) {
270270 for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);
271271 } else {
272272 EllipseBounds eb;
273273 if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;
274274 const int min_y =
275275 max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));
276276 const int max_y =
277277 min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));
278278 if (min_y > max_y) return;
279279 const int ty0 = max(rect.y0, min_y / kTileSize);
280280 const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);
281281 for (int ty = ty0; ty < ty1; ++ty) {
282282 int x0_px, x1_px;
283283 if (!eb.SpanForPixelRows(xy, ty * kTileSize,
284284 min((ty + 1) * kTileSize, height) - 1, x0_px,
285285 x1_px)) {
286286 continue;
287287 }
288288 x0_px = max(0, x0_px);
289289 x1_px = min(width - 1, x1_px);
290290 if (x0_px > x1_px) continue;
291291 const int tx0 = max(rect.x0, x0_px / kTileSize);
292292 const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);
293293 if (tx0 < tx1) fn(ty, tx0, tx1);
294294 }
295295 }
296296 }
297297
298298 namespace warp {
299299
300300 inline constexpr unsigned kFullMask = 0xffffffffu;
301301
302302 // Reduce-scatter of an array of `N` <= 32 per-lane values over the warp.
303303 //
304304 // A butterfly where each step halves the payload: at offset `o`, lanes with
305305 // bit `o` clear keep the low half of the array and send the high half, the
306306 // others do the opposite. After log2(32) steps every lane holds a single
307307 // value, v[0], which is the warp total of element `ScatterIndex(N, lane)`
308308 // (or junk if that is -1). For N = 9 this takes 12 shuffles instead of the
309309 // 45 of nine independent butterflies, and the totals end in distinct lanes,
310310 // so writing them out is one instruction.
311311
312312 namespace internal {
313313 template <int N, int kOffset>
314314 GS_D void HalvingStep(float* v, int lane) {
315315 constexpr int kHalf = (N + 1) / 2;
316316 const bool upper = lane & kOffset;
317317 #pragma unroll
318318 for (int i = 0; i < kHalf; ++i) {
319319 const float lo = v[i];
320320 const float hi = i + kHalf < N ? v[i + kHalf] : 0.0f;
321321 const float send = upper ? lo : hi;
322322 const float keep = upper ? hi : lo;
323323 v[i] = keep + __shfl_xor_sync(kFullMask, send, kOffset);
324324 }
325325 }
326326
327327 template <int N, int kOffset>
328328 GS_D void ReduceScatterImpl(float* v, int lane) {
329329 if constexpr (kOffset > 0) {
330330 HalvingStep<N, kOffset>(v, lane);
331331 ReduceScatterImpl<(N + 1) / 2, kOffset / 2>(v, lane);
332332 }
333333 }
334334 } // namespace internal
335335
336336 template <int N>
337337 GS_D void ReduceScatter(float (&v)[N], int lane) {
338338 static_assert(N >= 1 && N <= 32);
339339 internal::ReduceScatterImpl<N, 16>(v, lane);
340340 }
341341
342342 template <int N>
343343 GS_D void ReduceScatter16(float (&v)[N], int lane) {
344344 static_assert(N >= 1 && N <= 16);
345345 internal::ReduceScatterImpl<N, 8>(v, lane);
346346 }
347347
348348 // Element whose warp total lane `lane` holds after `ReduceScatter<N>`, or -1.
349349 GS_HD constexpr int ScatterIndex(int n, int lane) {
350350 // Payload sizes and halves per step, offset 16 first.
351351 int sizes[5] = {};
352352 int size = n;
353353 for (int s = 0; s < 5; ++s) {
354354 sizes[s] = size;
355355 size = (size + 1) / 2;
356356 }
357357 int index = 0; // Slot in the final (size 1) array.
358358 for (int s = 4; s >= 0; --s) {
359359 const int half = (sizes[s] + 1) / 2;
360360 if (lane & (16 >> s)) index += half;
361361 if (index >= sizes[s]) return -1;
362362 }
363363 return index;
364364 }
365365
366366 // Element whose 16-lane half-warp total lane `lane` holds after
367367 // `ReduceScatter16<N>`, or -1.
368368 GS_HD constexpr int ScatterIndex16(int n, int lane) {
369369 const int half_lane = lane & 15;
370370 int sizes[4] = {};
371371 int size = n;
372372 for (int s = 0; s < 4; ++s) {
373373 sizes[s] = size;
374374 size = (size + 1) / 2;
375375 }
376376 int index = 0;
377377 for (int s = 3; s >= 0; --s) {
378378 const int half = (sizes[s] + 1) / 2;
379379 if (half_lane & (8 >> s)) index += half;
380380 if (index >= sizes[s]) return -1;
381381 }
382382 return index;
383383 }
384384
385385 } // namespace warp
386386
387387 inline constexpr int kShCoeffs = 16; // Degree 3.
388388 inline constexpr float kNearPlane = 0.2f;
389389 inline constexpr float kLowPass = 0.3f;
390390 inline constexpr float kFovClamp = 1.3f;
391391
392392 // Pointers to the trainable parameters, all activated, row-major.
393393 struct GaussianParams {
394394 const float* means; // (N, 3)
395395 const float* scales; // (N, 3)
396396 const float* quats; // (N, 4) as (w, x, y, z), not normalized.
397397 const float* opacities; // (N,)
398398 const void* sh; // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).
399399 int n;
400400 };
401401
402402 // Gradients w.r.t. `GaussianParams`, same layout and SH storage type.
403403 struct GaussianGrads {
404404 float* means;
405405 float* scales;
406406 float* quats;
407407 float* opacities;
408408 void* sh;
409409 };
410410
411411 // Screen-space quantities needed by rasterization. `conic` is the inverse 2D
412412 // covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.
413413 struct Projection {
414414 float2 xy;
415415 float depth;
416416 float3 conic;
417417 float opacity;
418418 float limit;
419419 float3 rgb;
420420 int radius; // 0 means culled.
421421 };
422422
423423 // ----------------------------------------------------------------------------
424424 // SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian
425425 // and is accessed in 4 groups of 4 coefficients (12 values). The coefficients
426426 // may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.
427427
428428 using ShBf16 = uint16_t; // bfloat16 bits.
429429
430430 template <typename T>
431431 struct ShIo;
432432
433433 template <>
434434 struct ShIo<float> {
435435 static constexpr int kRowVec16 = 12; // 16-byte vectors per row.
436436 GS_D static void LoadGroup(const float* row, int grp, float v[12]) {
437437 const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;
438438 const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);
439439 v[0] = a.x, v[1] = a.y, v[2] = a.z, v[3] = a.w;
440440 v[4] = b.x, v[5] = b.y, v[6] = b.z, v[7] = b.w;
441441 v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;
442442 }
443443 GS_D static void StoreGroup(float* row, int grp, const float v[12]) {
444444 float4* p = reinterpret_cast<float4*>(row) + 3 * grp;
445445 p[0] = {v[0], v[1], v[2], v[3]};
446446 p[1] = {v[4], v[5], v[6], v[7]};
447447 p[2] = {v[8], v[9], v[10], v[11]};
448448 }
449449 };
450450
451451 template <>
452452 struct ShIo<ShBf16> {
453453 static constexpr int kRowVec16 = 6;
454454 GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
455455 GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
456456 // Round to nearest even (inputs are finite).
457457 GS_D static uint32_t Round(float f) {
458458 const uint32_t u = __float_as_uint(f);
459459 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
460460 }
461461 GS_D static uint32_t Pack(float lo, float hi) {
462462 return Round(lo) | (Round(hi) << 16);
463463 }
464464 GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {
465465 // 24-byte groups: three 8-byte loads.
466466 const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;
467467 #pragma unroll
468468 for (int k = 0; k < 3; ++k) {
469469 const uint2 w = __ldg(p + k);
470470 v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);
471471 v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);
472472 }
473473 }
474474 GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {
475475 uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;
476476 #pragma unroll
477477 for (int k = 0; k < 3; ++k) {
478478 p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};
479479 }
480480 }
481481 };
482482
483483 template <typename ShT>
484484 GS_D const ShT* ShRow(const void* sh, int i) {
485485 return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
486486 }
487487 template <typename ShT>
488488 GS_D ShT* ShRow(void* sh, int i) {
489489 return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
490490 }
491491
492492 // ----------------------------------------------------------------------------
493493 // Spherical harmonics (degree 3, Inria sign convention).
494494
495495 namespace sh {
496496 inline constexpr float kC0 = 0.28209479177387814f;
497497 inline constexpr float kC1 = 0.4886025119029199f;
498498 inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
499499 0.31539156525252005f, -1.0925484305920792f,
500500 0.5462742152960396f};
501501 inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
502502 -0.4570457994644658f, 0.3731763325901154f,
503503 -0.4570457994644658f, 1.445305721320277f,
504504 -0.5900435899266435f};
505505
506506 // Basis values for a unit direction.
507507 GS_D void Basis(float3 d, float b[kShCoeffs]) {
508508 const float x = d.x, y = d.y, z = d.z;
509509 const float xx = x * x, yy = y * y, zz = z * z;
510510 b[0] = kC0;
511511 b[1] = -kC1 * y;
512512 b[2] = kC1 * z;
513513 b[3] = -kC1 * x;
514514 b[4] = kC2[0] * x * y;
515515 b[5] = kC2[1] * y * z;
516516 b[6] = kC2[2] * (2.0f * zz - xx - yy);
517517 b[7] = kC2[3] * x * z;
518518 b[8] = kC2[4] * (xx - yy);
519519 b[9] = kC3[0] * y * (3.0f * xx - yy);
520520 b[10] = kC3[1] * x * y * z;
521521 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
522522 b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
523523 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
524524 b[14] = kC3[5] * z * (xx - yy);
525525 b[15] = kC3[6] * x * (xx - 3.0f * yy);
526526 }
527527
528528 // Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.
529529 template <typename ShT>
530530 GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
531531 float3 c = {0.5f, 0.5f, 0.5f};
532532 #pragma unroll
533533 for (int grp = 0; grp < 4; ++grp) {
534534 const int k0 = 4 * grp;
535535 float v[12];
536536 ShIo<ShT>::LoadGroup(sh_row, grp, v);
537537 c += b[k0] * float3{v[0], v[1], v[2]};
538538 c += b[k0 + 1] * float3{v[3], v[4], v[5]};
539539 c += b[k0 + 2] * float3{v[6], v[7], v[8]};
540540 c += b[k0 + 3] * float3{v[9], v[10], v[11]};
541541 }
542542 return c;
543543 }
544544 } // namespace sh
545545
546546 // ----------------------------------------------------------------------------
547547 // Covariance helpers.
548548
549549 // Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
550550 GS_D Mat3 QuatToRotation(float4 q) {
551551 const float w = q.x, x = q.y, y = q.z, z = q.w;
552552 Mat3 r;
553553 r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
554554 r.m[0][1] = 2.0f * (x * y - w * z);
555555 r.m[0][2] = 2.0f * (x * z + w * y);
556556 r.m[1][0] = 2.0f * (x * y + w * z);
557557 r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
558558 r.m[1][2] = 2.0f * (y * z - w * x);
559559 r.m[2][0] = 2.0f * (x * z - w * y);
560560 r.m[2][1] = 2.0f * (y * z + w * x);
561561 r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
562562 return r;
563563 }
564564
565565 // Symmetric 3x3 matrix.
566566 struct Sym3 {
567567 float xx, xy, xz, yy, yz, zz;
568568 GS_D float3 Mul(float3 v) const {
569569 return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
570570 xz * v.x + yz * v.y + zz * v.z};
571571 }
572572 };
573573
574574 // Sigma = A A^T with A = R diag(s).
575575 GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
576576 Mat3 a;
577577 const float sv[3] = {s.x, s.y, s.z};
578578 #pragma unroll
579579 for (int i = 0; i < 3; ++i)
580580 #pragma unroll
581581 for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
582582 return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
583583 dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
584584 dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
585585 }
586586
587587 // Intermediate values of the projection shared by forward and VJP.
588588 struct ProjectionState {
589589 float3 t; // View-space mean.
590590 float tx_c, ty_c; // Clamped x, y used by the Jacobian.
591591 bool x_clamped, y_clamped;
592592 float3 m0, m1; // Rows of M = J W (2x3).
593593 Sym3 sigma; // World-space 3D covariance.
594594 float a, b, c; // Dilated 2D covariance.
595595 };
596596
597597 // Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
598598 GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
599599 float4 quat, ProjectionState& st) {
600600 st.t = t;
601601 const float tz = t.z;
602602 const float limx = kFovClamp * cam.TanHalfFovX();
603603 const float limy = kFovClamp * cam.TanHalfFovY();
604604 const float txtz = t.x / tz, tytz = t.y / tz;
605605 st.x_clamped = txtz < -limx || txtz > limx;
606606 st.y_clamped = tytz < -limy || tytz > limy;
607607 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
608608 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
609609 const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
610610 const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
611611 const Mat3& w = cam.rotation;
612612 st.m0 = j00 * w.row(0) + j02 * w.row(2);
613613 st.m1 = j11 * w.row(1) + j12 * w.row(2);
614614 st.sigma = Covariance3D(QuatToRotation(quat), scale);
615615 const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
616616 st.a = dot(st.m0, s0) + kLowPass;
617617 st.b = dot(st.m0, s1);
618618 st.c = dot(st.m1, s1) + kLowPass;
619619 }
620620
621621 // Computes everything up to the 2D covariance. Returns false if culled by the
622622 // near plane.
623623 GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
624624 float4 quat, ProjectionState& st) {
625625 const float3 t = cam.ToView(mean);
626626 if (t.z <= kNearPlane) return false;
627627 ProjectCovarianceFromView(cam, t, scale, quat, st);
628628 return true;
629629 }
630630
631631 // Integer bounding radius (pixels) of the 2D covariance, as in the reference.
632632 GS_D int BoundingRadius(float a, float b, float c) {
633633 const float det = a * c - b * b;
634634 const float mid = 0.5f * (a + c);
635635 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
636636 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
637637 }
638638
639639 GS_D float2 ProjectMean(const Camera& cam, float3 t) {
640640 return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
641641 }
642642
643643 GS_D float3 LoadMean(const GaussianParams& p, int i) {
644644 return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
645645 }
646646
647647 // Unit direction from the camera center to `mean`, and 1 / distance.
648648 GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
649649 const float3 v = mean - cam.position;
650650 inv_len = rsqrtf(dot(v, v));
651651 return inv_len * v;
652652 }
653653
654654 // Gradient of the loss w.r.t. one splat's screen-space quantities.
655655 struct ProjectionGrad {
656656 float2 xy;
657657 float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.
658658 float opacity;
659659 float3 rgb;
660660 };
661661
662662 // Image geometry passed by value to every kernel.
663663 struct Frame {
664664 const float* camera; // Device pointer to the packed camera.
665665 int width, height;
666666 TileGrid grid;
667667 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
668668 };
669669
670670 // ---------------------------------------------------------------------------
671671 // Per-Gaussian kernels run one thread per Gaussian.
672672
673673 inline constexpr int kGaussianThreads = 128;
674674
675675 inline int GaussianBlocks(int64_t n) {
676676 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
677677 }
678678
679679 // Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
680680 inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
681681 inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
682682
683683 GS_HD float3 Power2Conic(float a, float b, float c) {
684684 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
685685 }
686686
687687 GS_D float Exp2Approx(float x) {
688688 float y;
689689 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
690690 return y;
691691 }
692692
693693 // ---------------------------------------------------------------------------
694−// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into
694+// Rasterization. One block per 16x16 tile. Each thread shades `kRows`
695+// vertically adjacent pixels, so that per-splat costs (shared-memory loads,
696+// barriers, the warp reduction in the backward pass) are paid once per
697+// `kRows` pixels; more rows per thread also means fewer resident threads and
698+// more waiting on the slowest pixel. A warp covers a compact
699+// kWarpWidth x kWarpRows region (as square as possible, so per-warp culling
695−// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
700+// is tight); warps tile the 16x16 tile row-major. Splats are staged through
696−// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
701+// shared memory, one per thread per batch.
702+
703+inline constexpr int kRasterRowChoices[] = {1, 2, 4};
697704
698705 template <int kRows>
699706 struct RasterLayout {
700707 static_assert(kRows == 1 || kRows == 2 || kRows == 4);
701708 static constexpr int kThreads = kTilePixels / kRows;
702709 static constexpr int kWarps = kThreads / 32;
703710 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
704711 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
705712 static constexpr int kWarpsX = kTileSize / kWarpWidth;
706713 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
707714
715+ // Tile-relative origin of warp `w`'s region.
708716 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
709717 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
710718 };
711719
720+// Pixel coordinates of this thread's first pixel; the others follow in y.
712721 struct RasterPixels {
713722 int x, y0;
714723 template <int kRows>
715724 GS_D static RasterPixels For(int tid) {
↕ Expand 88 unchanged lines (lines 725–812) · in WarpMask
716725 using L = RasterLayout<kRows>;
717726 const int warp = tid / 32, lane = tid % 32;
718727 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
719728 lane % L::kWarpWidth,
720729 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
721730 (lane / L::kWarpWidth) * kRows};
722731 }
723732 };
724733
725734 // Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
726735 // pixel region of warp w of this block's tile.
727736 template <int kRows>
728737 GS_D uint32_t WarpMask(float4 p0, float4 p1) {
729738 using L = RasterLayout<kRows>;
730739 constexpr int kWarpsY = L::kWarps / L::kWarpsX;
731740 float limit = p1.z;
732741 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
733742 const float a = p0.z, b = p0.w, c = p1.x;
734743 const float det = a * c - b * b;
735744 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
736745 const float l_over_det = limit / det;
737746 const float rx = sqrtf(c * l_over_det);
738747 const float ry = sqrtf(a * l_over_det);
739748 const float2 xy = {p0.x, p0.y};
740749 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
741750 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
742751 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
743752 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
744753 if (x0 > x1 || y0 > y1) return 0;
745754 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
746755 uint32_t mx = 0u;
747756 #pragma unroll
748757 for (int wx = 0; wx < L::kWarpsX; ++wx) {
749758 const int wx0 = tx + wx * L::kWarpWidth;
750759 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
751760 }
752761 if (mx == 0u) return 0;
753762 uint32_t my = 0u;
754763 #pragma unroll
755764 for (int wy = 0; wy < kWarpsY; ++wy) {
756765 const int wy0 = ty + wy * L::kWarpRows;
757766 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
758767 }
759768 if (my == 0u) return 0;
760769 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
761770 uint32_t mask = 0u;
762771 #pragma unroll
763772 for (int wy = 0; wy < kWarpsY; ++wy) {
764773 if (!((my >> wy) & 1u)) continue;
765774 #pragma unroll
766775 for (int wx = 0; wx < L::kWarpsX; ++wx) {
767776 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
768777 }
769778 }
770779 return mask;
771780 }
772781 EllipseBounds eb;
773782 eb.a = a;
774783 eb.b = b;
775784 eb.c = c;
776785 eb.det = det;
777786 eb.inv_a = 1.0f / a;
778787 eb.a_limit = a * limit;
779788 eb.y_ext = ry;
780789 eb.dy_peak = b * sqrtf(l_over_det / c);
781790 uint32_t mask = 0u;
782791 #pragma unroll
783792 for (int wy = 0; wy < kWarpsY; ++wy) {
784793 if (!((my >> wy) & 1u)) continue;
785794 const int wy0 = ty + wy * L::kWarpRows;
786795 int x0_px, x1_px;
787796 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
788797 continue;
789798 }
790799 #pragma unroll
791800 for (int wx = 0; wx < L::kWarpsX; ++wx) {
792801 const int wx0 = tx + wx * L::kWarpWidth;
793802 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
794803 mask |= 1u << (wy * L::kWarpsX + wx);
795804 }
796805 }
797806 }
798807 return mask;
799808 }
800809
801810 } // namespace gs
802811
803812 #endif
csrc/forward.h +3 −2
s03_sh_streaming/csrc/forward.h
s04_raster_pruning/csrc/forward.h
↕ Expand 23 unchanged lines (lines 1–23) · in PreprocessKernel
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 "auxiliary.h"15#include "auxiliary.h"
1616
17namespace gs {17namespace gs {
1818
19template <Culling kCulling, typename ShT>19template <Culling kCulling, typename ShT>
20__global__ void __launch_bounds__(kGaussianThreads, 8)20__global__ void __launch_bounds__(kGaussianThreads, 8)
21 PreprocessKernel(GaussianParams params, Frame frame,21 PreprocessKernel(GaussianParams params, Frame frame,
22 Splat* __restrict__ splats,22 Splat* __restrict__ splats,
23 int32_t* __restrict__ tile_counts,23 int32_t* __restrict__ tile_counts,
24 uint32_t* __restrict__ depth_keys,24 uint32_t* __restrict__ depth_keys,
25 int32_t* __restrict__ ids,25 int32_t* __restrict__ ids,
26 int32_t* __restrict__ num_visible);26 int32_t* __restrict__ num_visible);
2727
28template <int kRows>
28__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)29__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
29 RasterForwardKernel(const int2* __restrict__ ranges,30 RasterForwardKernel(const int2* __restrict__ ranges,
30 const int32_t* __restrict__ point_list,31 int32_t* __restrict__ point_list,
31 const Splat* __restrict__ splats, Frame frame,32 const Splat* __restrict__ splats, Frame frame,
32 float* __restrict__ image, float* __restrict__ final_t,33 float* __restrict__ image, float* __restrict__ final_t,
33 int32_t* __restrict__ n_contrib);34 int32_t* __restrict__ n_contrib);
3435
35} // namespace gs36} // namespace gs
3637
37#endif38#endif
↕ Expand 23 unchanged lines (lines 1–23) · in PreprocessKernel
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
1515 #include "auxiliary.h"
1616
1717 namespace gs {
1818
1919 template <Culling kCulling, typename ShT>
2020 __global__ void __launch_bounds__(kGaussianThreads, 8)
2121 PreprocessKernel(GaussianParams params, Frame frame,
2222 Splat* __restrict__ splats,
2323 int32_t* __restrict__ tile_counts,
2424 uint32_t* __restrict__ depth_keys,
2525 int32_t* __restrict__ ids,
2626 int32_t* __restrict__ num_visible);
2727
28+template <int kRows>
28−__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
29+__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
2930 RasterForwardKernel(const int2* __restrict__ ranges,
30− const int32_t* __restrict__ point_list,
31+ int32_t* __restrict__ point_list,
3132 const Splat* __restrict__ splats, Frame frame,
3233 float* __restrict__ image, float* __restrict__ final_t,
3334 int32_t* __restrict__ n_contrib);
3435
3536 } // namespace gs
3637
3738 #endif
csrc/forward.cu.cc +119 −37
s03_sh_streaming/csrc/forward.cu.cc
s04_raster_pruning/csrc/forward.cu.cc
↕ Expand 102 unchanged lines (lines 1–102) · in PreprocessKernel
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 "auxiliary.h"14#include "auxiliary.h"
1515
16namespace gs {16namespace gs {
1717
18// Screen-space geometry of Gaussian `i`: everything in `Projection` except18// Screen-space geometry of Gaussian `i`: everything in `Projection` except
19// `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the19// `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
20// tile-rectangle cull).20// tile-rectangle cull).
21template <Culling kCulling>21template <Culling kCulling>
22GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,22GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
23 float3 mean, int i) {23 float3 mean, int i) {
24 Projection out;24 Projection out;
25 out.radius = 0;25 out.radius = 0;
26 const float3 t = cam.ToView(mean);26 const float3 t = cam.ToView(mean);
27 if (t.z <= kNearPlane) return out;27 if (t.z <= kNearPlane) return out;
28 const float opacity = p.opacities[i];28 const float opacity = p.opacities[i];
29 float limit = 0.0f;29 float limit = 0.0f;
30 if (!AlphaLimit(opacity, limit)) {30 if (!AlphaLimit(opacity, limit)) {
31 if constexpr (kCulling == Culling::kEllipse) return out;31 if constexpr (kCulling == Culling::kEllipse) return out;
32 }32 }
33 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],33 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
34 p.scales[3 * i + 2]};34 p.scales[3 * i + 2]};
35 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];35 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
36 ProjectionState st;36 ProjectionState st;
37 ProjectCovarianceFromView(cam, t, scale, quat, st);37 ProjectCovarianceFromView(cam, t, scale, quat, st);
38 const float det = st.a * st.c - st.b * st.b;38 const float det = st.a * st.c - st.b * st.b;
39 if (det == 0.0f) return out;39 if (det == 0.0f) return out;
40 const float inv = 1.0f / det;40 const float inv = 1.0f / det;
41 out.conic = {st.c * inv, -st.b * inv, st.a * inv};41 out.conic = {st.c * inv, -st.b * inv, st.a * inv};
42 out.xy = ProjectMean(cam, st.t);42 out.xy = ProjectMean(cam, st.t);
43 out.depth = st.t.z;43 out.depth = st.t.z;
44 out.opacity = opacity;44 out.opacity = opacity;
45 out.limit = limit;45 out.limit = limit;
46 out.radius = BoundingRadius(st.a, st.b, st.c);46 out.radius = BoundingRadius(st.a, st.b, st.c);
47 return out;47 return out;
48}48}
4949
50// Color seen along `dir`, clamped at zero, reading from an SH row.50// Color seen along `dir`, clamped at zero, reading from an SH row.
51template <typename ShT>51template <typename ShT>
52GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {52GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
53 float basis[kShCoeffs];53 float basis[kShCoeffs];
54 sh::Basis(dir, basis);54 sh::Basis(dir, basis);
55 const float3 raw = sh::EvalRow(sh_row, 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)};56 return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
57}57}
5858
59// ---------------------------------------------------------------------------59// ---------------------------------------------------------------------------
60// Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,60// Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
61// evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit61// evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
62// depth sort via per-thread `atomicAdd`.62// depth sort via per-thread `atomicAdd`.
6363
64template <Culling kCulling, typename ShT>64template <Culling kCulling, typename ShT>
65__global__ void __launch_bounds__(kGaussianThreads, 8)65__global__ void __launch_bounds__(kGaussianThreads, 8)
66 PreprocessKernel(GaussianParams params, Frame frame,66 PreprocessKernel(GaussianParams params, Frame frame,
67 Splat* __restrict__ splats,67 Splat* __restrict__ splats,
68 int32_t* __restrict__ tile_counts,68 int32_t* __restrict__ tile_counts,
69 uint32_t* __restrict__ depth_keys,69 uint32_t* __restrict__ depth_keys,
70 int32_t* __restrict__ ids,70 int32_t* __restrict__ ids,
71 int32_t* __restrict__ num_visible) {71 int32_t* __restrict__ num_visible) {
72 const int i = blockIdx.x * blockDim.x + threadIdx.x;72 const int i = blockIdx.x * blockDim.x + threadIdx.x;
73 if (i >= params.n) return;73 if (i >= params.n) return;
74 const Camera cam = frame.LoadCamera();74 const Camera cam = frame.LoadCamera();
75 // Geometry first: only Gaussians that reach a tile need their SH row.75 // Geometry first: only Gaussians that reach a tile need their SH row.
76 const float3 mean = LoadMean(params, i);76 const float3 mean = LoadMean(params, i);
77 const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);77 const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
78 Splat s;78 Splat s;
79 int count = 0;79 int count = 0;
80 if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {80 if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
81 const float min_power =81 const float min_power =
82 pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;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};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)};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};85 s.p2 = {0.0f, 0.0f, 0.0f, min_power};
86 ForEachTileRow<kCulling>(86 ForEachTileRow<kCulling>(
87 s, frame.grid, frame.width, frame.height,87 s, frame.grid, frame.width, frame.height,
88 [&](int, int tx0, int tx1) { count += tx1 - tx0; });88 [&](int, int tx0, int tx1) { count += tx1 - tx0; });
89 }89 }
90 if (count == 0) {90 if (count == 0) {
91 s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};91 s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
92 splats[i] = s;92 splats[i] = s;
93 tile_counts[i] = 0;93 tile_counts[i] = 0;
94 return;94 return;
95 }95 }
96 const int vis_idx = atomicAdd(num_visible, 1);96 const int vis_idx = atomicAdd(num_visible, 1);
97 float inv_len;97 float inv_len;
98 const float3 rgb =98 const float3 rgb =
99 ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));99 ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
100 s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};100 s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
101 splats[i] = s;101 splats[i] = s;
102 tile_counts[i] = count;102 tile_counts[i] = count;
103 depth_keys[vis_idx] = AsUint(pr.depth);103 depth_keys[vis_idx] = AsUint(pr.depth);
104 ids[vis_idx] = i;104 ids[vis_idx] = i;
105}105}
106106
107template <int kRows>
107__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)108__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
108 RasterForwardKernel(const int2* __restrict__ ranges,109 RasterForwardKernel(const int2* __restrict__ ranges,
109 const int32_t* __restrict__ point_list,110 int32_t* __restrict__ point_list,
110 const Splat* __restrict__ splats, Frame frame,111 const Splat* __restrict__ splats, Frame frame,
111 float* __restrict__ image, float* __restrict__ final_t,112 float* __restrict__ image, float* __restrict__ final_t,
112 int32_t* __restrict__ n_contrib) {113 int32_t* __restrict__ n_contrib) {
114 constexpr int kRasterRows = kRows;
113 constexpr int kRasterThreads = RasterLayout<1>::kThreads;115 constexpr int kRasterThreads = RasterLayout<kRows>::kThreads;
116 constexpr int kWarps = RasterLayout<kRows>::kWarps;
114 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;117 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
115 const int tid = threadIdx.x;118 const int tid = threadIdx.x;
116 const int warp = tid / 32;119 const int warp = tid / 32;
117 const int lane = tid % 32;120 const int lane = tid % 32;
118 const RasterPixels px = RasterPixels::For<1>(tid);121 const RasterPixels px = RasterPixels::For<kRows>(tid);
119 const float fx = static_cast<float>(px.x);122 const float fx = static_cast<float>(px.x);
120 const float fy = static_cast<float>(px.y0);
121 const int2 range = ranges[tile];123 const int2 range = ranges[tile];
122124
125 float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
126 int last[kRasterRows];
127 bool done[kRasterRows];
128#pragma unroll
129 for (int k = 0; k < kRasterRows; ++k) {
130 t[k] = 1.0f;
123 float t = 1.0f, r = 0.0f, g = 0.0f, b = 0.0f;131 r[k] = g[k] = b[k] = 0.0f;
124 int last = 0;132 last[k] = 0;
125 bool done = px.x >= frame.width || px.y0 >= frame.height;133 done[k] = px.x >= frame.width || px.y0 + k >= frame.height;
134 }
135 auto all_done = [&] {
136 bool d = true;
137#pragma unroll
138 for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
139 return d;
140 };
126141
127 __shared__ float4 batch_geom0[kRasterThreads];142 __shared__ float4 batch_geom0[kRasterThreads];
128 __shared__ float2 batch_geom1[kRasterThreads];143 __shared__ float2 batch_geom1[kRasterThreads];
129 __shared__ float4 batch_rgb_o[kRasterThreads];144 __shared__ float4 batch_rgb_o[kRasterThreads];
130 __shared__ uint8_t batch_mask[kRasterThreads];145 __shared__ uint8_t batch_mask[kRasterThreads];
146 __shared__ uint32_t warp_contrib[kWarps][kWarps];
131 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);147 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
148 bool thread_done = all_done();
149 int tile_compact = 0;
132 for (int start = range.x; start < range.y; start += kRasterThreads) {150 for (int start = range.x; start < range.y; start += kRasterThreads) {
133 if (__syncthreads_count(done) == kRasterThreads) break;
134 const int idx = start + tid;151 const int idx = start + tid;
152 int32_t id = 0;
135 if (idx < range.y) {153 if (idx < range.y) {
136 const int32_t id = point_list[idx];154 id = point_list[idx];
137 const float4 p0 = __ldg(&splats[id].p0);155 const float4 p0 = __ldg(&splats[id].p0);
138 const float4 p1 = __ldg(&splats[id].p1);156 const float4 p1 = __ldg(&splats[id].p1);
139 const float4 p2 = __ldg(&splats[id].p2);157 const float4 p2 = __ldg(&splats[id].p2);
140 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);158 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
141 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};159 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
142 batch_geom1[tid] = {c2.z, p2.w};160 batch_geom1[tid] = {c2.z, p2.w};
143 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};161 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
144 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));162 batch_mask[tid] = static_cast<uint8_t>(WarpMask<kRows>(p0, p1));
145 } else {163 } else {
146 batch_mask[tid] = 0u;164 batch_mask[tid] = 0u;
147 }165 }
148 __syncthreads();166 __syncthreads();
149 const int count = min(kRasterThreads, range.y - start);167 const int count = min(kRasterThreads, range.y - start);
150 const int base = start - range.x;
151 bool warp_done = __all_sync(warp::kFullMask, done);168 bool warp_done = __all_sync(warp::kFullMask, thread_done);
152#pragma unroll169#pragma unroll
153 for (int c = 0; c < kRasterThreads; c += 32) {170 for (int c = 0; c < kRasterThreads; c += 32) {
171 uint32_t contrib = 0u;
154 if (c < count && !warp_done) {172 if (c < count && !warp_done) {
155 uint32_t mask = __ballot_sync(173 uint32_t mask = __ballot_sync(
156 warp::kFullMask,174 warp::kFullMask,
157 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);175 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
158 while (mask != 0u) {176 while (mask != 0u) {
159 const int bit = __ffs(mask) - 1;177 const int bit = __ffs(mask) - 1;
160 mask &= mask - 1u;178 mask &= mask - 1u;
161 const int j = c + bit;179 const int j = c + bit;
180 bool blended = false;
162 if (!done) {181 if (!thread_done) {
163 const float4 g0 = batch_geom0[j];182 const float4 g0 = batch_geom0[j];
164 const float2 g1 = batch_geom1[j];183 const float2 g1 = batch_geom1[j];
165 const float dx = g0.x - fx;184 const float dx = g0.x - fx;
166 const float dy = g0.y - fy;185 const float cx_dx2 = g0.z * dx * dx;
167 const float power =186 const float cy_dx = g0.w * dx;
187#pragma unroll
188 for (int k = 0; k < kRasterRows; ++k) {
189 if (done[k]) continue;
190 const float dy = g0.y - static_cast<float>(px.y0 + k);
168 g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;191 const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
169 if (power <= 0.0f && power >= g1.y) {192 if (power > 0.0f || power < g1.y) continue;
170 const float4 ro = batch_rgb_o[j];193 const float4 ro = batch_rgb_o[j];
171 const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));194 const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
172 if (alpha >= kMinAlpha) {195 if (alpha < kMinAlpha) continue;
173 const float next_t = t * (1.0f - alpha);196 const float next_t = t[k] * (1.0f - alpha);
174 if (next_t < kMinTransmittance) {197 if (next_t < kMinTransmittance) {
175 done = true;198 done[k] = 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;
199 continue;
183 }200 }
201 const float w = alpha * t[k];
202 r[k] += w * ro.x;
203 g[k] += w * ro.y;
204 b[k] += w * ro.z;
205 t[k] = next_t;
206 last[k] = -1 - j;
207 blended = true;
184 }208 }
209 thread_done = all_done();
185 }210 }
211 if (__any_sync(warp::kFullMask, blended)) {
212 contrib |= (1u << bit);
186 }213 }
187 if (__all_sync(warp::kFullMask, done)) {214 if (__all_sync(warp::kFullMask, thread_done)) {
188 warp_done = true;215 warp_done = true;
189 break;216 break;
190 }217 }
191 }218 }
192 warp_done = __all_sync(warp::kFullMask, done);219 warp_done = __all_sync(warp::kFullMask, thread_done);
220 }
221 if (lane == 0) {
222 warp_contrib[warp][c / 32] = contrib;
193 }223 }
194 }224 }
225 const bool block_done =
226 (__syncthreads_count(thread_done) == kRasterThreads);
227 uint32_t chunk_any[kWarps];
228 int chunk_prefix[kWarps];
229 int batch_active = 0;
230#pragma unroll
231 for (int u = 0; u < kWarps; ++u) {
232 uint32_t any_u = 0u;
233#pragma unroll
234 for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
235 any_u |= warp_contrib[w_idx][u];
195 }236 }
196 if (px.x < frame.width && px.y0 < frame.height) {
237 chunk_any[u] = any_u;
238 chunk_prefix[u] = tile_compact + batch_active;
239 batch_active += __popc(any_u);
240 }
241 if ((chunk_any[warp] >> lane) & 1u) {
242 uint32_t cmask = 0u;
243#pragma unroll
244 for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
245 cmask |= ((warp_contrib[w_idx][warp] >> lane) & 1u) << w_idx;
246 }
247 const int dst =
248 chunk_prefix[warp] + __popc(chunk_any[warp] & ((1u << lane) - 1u));
249 point_list[range.x + dst] =
250 static_cast<int32_t>((cmask << 24) | static_cast<uint32_t>(id));
251 }
252#pragma unroll
253 for (int k = 0; k < kRasterRows; ++k) {
254 if (last[k] < 0) {
255 const int j = -1 - last[k];
256 const int u = j >> 5;
257 const int b_idx = j & 31;
258 last[k] =
259 chunk_prefix[u] + __popc(chunk_any[u] & ((1u << b_idx) - 1u)) + 1;
260 }
261 }
262 tile_compact += batch_active;
263 if (block_done) break;
264 }
197 const int pixels = frame.width * frame.height;265 const int pixels = frame.width * frame.height;
266#pragma unroll
267 for (int k = 0; k < kRasterRows; ++k) {
268 const int y = px.y0 + k;
269 if (px.x >= frame.width || y >= frame.height) continue;
198 const int p = px.y0 * frame.width + px.x;270 const int p = y * frame.width + px.x;
199 image[0 * pixels + p] = r;271 image[0 * pixels + p] = r[k];
200 image[1 * pixels + p] = g;272 image[1 * pixels + p] = g[k];
201 image[2 * pixels + p] = b;273 image[2 * pixels + p] = b[k];
202 final_t[p] = t;274 final_t[p] = t[k];
203 n_contrib[p] = last;275 n_contrib[p] = last[k];
204 }276 }
205}277}
206278
207template __global__ void PreprocessKernel<Culling::kRect, float>(279template __global__ void PreprocessKernel<Culling::kRect, float>(
↕ Expand 8 unchanged lines (lines 280–287) · in RasterForwardKernel
208 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,280 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
209 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);281 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
210template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(282template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
211 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,283 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
212 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);284 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
213template __global__ void PreprocessKernel<Culling::kEllipse, float>(285template __global__ void PreprocessKernel<Culling::kEllipse, float>(
214 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,286 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
215 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);287 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
216template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(288template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
217 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,289 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
218 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);290 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
219291
292template __global__ void RasterForwardKernel<1>(
293 const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
294 Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
295template __global__ void RasterForwardKernel<2>(
296 const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
297 Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
298template __global__ void RasterForwardKernel<4>(
299 const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
300 Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
301
220} // namespace gs302} // namespace gs
↕ Expand 102 unchanged lines (lines 1–102) · in PreprocessKernel
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
1414 #include "auxiliary.h"
1515
1616 namespace gs {
1717
1818 // Screen-space geometry of Gaussian `i`: everything in `Projection` except
1919 // `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
2020 // tile-rectangle cull).
2121 template <Culling kCulling>
2222 GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
2323 float3 mean, int i) {
2424 Projection out;
2525 out.radius = 0;
2626 const float3 t = cam.ToView(mean);
2727 if (t.z <= kNearPlane) return out;
2828 const float opacity = p.opacities[i];
2929 float limit = 0.0f;
3030 if (!AlphaLimit(opacity, limit)) {
3131 if constexpr (kCulling == Culling::kEllipse) return out;
3232 }
3333 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
3434 p.scales[3 * i + 2]};
3535 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
3636 ProjectionState st;
3737 ProjectCovarianceFromView(cam, t, scale, quat, st);
3838 const float det = st.a * st.c - st.b * st.b;
3939 if (det == 0.0f) return out;
4040 const float inv = 1.0f / det;
4141 out.conic = {st.c * inv, -st.b * inv, st.a * inv};
4242 out.xy = ProjectMean(cam, st.t);
4343 out.depth = st.t.z;
4444 out.opacity = opacity;
4545 out.limit = limit;
4646 out.radius = BoundingRadius(st.a, st.b, st.c);
4747 return out;
4848 }
4949
5050 // Color seen along `dir`, clamped at zero, reading from an SH row.
5151 template <typename ShT>
5252 GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
5353 float basis[kShCoeffs];
5454 sh::Basis(dir, basis);
5555 const float3 raw = sh::EvalRow(sh_row, basis);
5656 return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
5757 }
5858
5959 // ---------------------------------------------------------------------------
6060 // Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
6161 // evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
6262 // depth sort via per-thread `atomicAdd`.
6363
6464 template <Culling kCulling, typename ShT>
6565 __global__ void __launch_bounds__(kGaussianThreads, 8)
6666 PreprocessKernel(GaussianParams params, Frame frame,
6767 Splat* __restrict__ splats,
6868 int32_t* __restrict__ tile_counts,
6969 uint32_t* __restrict__ depth_keys,
7070 int32_t* __restrict__ ids,
7171 int32_t* __restrict__ num_visible) {
7272 const int i = blockIdx.x * blockDim.x + threadIdx.x;
7373 if (i >= params.n) return;
7474 const Camera cam = frame.LoadCamera();
7575 // Geometry first: only Gaussians that reach a tile need their SH row.
7676 const float3 mean = LoadMean(params, i);
7777 const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
7878 Splat s;
7979 int count = 0;
8080 if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
8181 const float min_power =
8282 pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;
8383 s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
8484 s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
8585 s.p2 = {0.0f, 0.0f, 0.0f, min_power};
8686 ForEachTileRow<kCulling>(
8787 s, frame.grid, frame.width, frame.height,
8888 [&](int, int tx0, int tx1) { count += tx1 - tx0; });
8989 }
9090 if (count == 0) {
9191 s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
9292 splats[i] = s;
9393 tile_counts[i] = 0;
9494 return;
9595 }
9696 const int vis_idx = atomicAdd(num_visible, 1);
9797 float inv_len;
9898 const float3 rgb =
9999 ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
100100 s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
101101 splats[i] = s;
102102 tile_counts[i] = count;
103103 depth_keys[vis_idx] = AsUint(pr.depth);
104104 ids[vis_idx] = i;
105105 }
106106
107+template <int kRows>
107−__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 8)
108+__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
108109 RasterForwardKernel(const int2* __restrict__ ranges,
109− const int32_t* __restrict__ point_list,
110+ int32_t* __restrict__ point_list,
110111 const Splat* __restrict__ splats, Frame frame,
111112 float* __restrict__ image, float* __restrict__ final_t,
112113 int32_t* __restrict__ n_contrib) {
114+ constexpr int kRasterRows = kRows;
113− constexpr int kRasterThreads = RasterLayout<1>::kThreads;
115+ constexpr int kRasterThreads = RasterLayout<kRows>::kThreads;
116+ constexpr int kWarps = RasterLayout<kRows>::kWarps;
114117 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
115118 const int tid = threadIdx.x;
116119 const int warp = tid / 32;
117120 const int lane = tid % 32;
118− const RasterPixels px = RasterPixels::For<1>(tid);
121+ const RasterPixels px = RasterPixels::For<kRows>(tid);
119122 const float fx = static_cast<float>(px.x);
120− const float fy = static_cast<float>(px.y0);
121123 const int2 range = ranges[tile];
122124
125+ float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
126+ int last[kRasterRows];
127+ bool done[kRasterRows];
128+#pragma unroll
129+ for (int k = 0; k < kRasterRows; ++k) {
130+ t[k] = 1.0f;
123− float t = 1.0f, r = 0.0f, g = 0.0f, b = 0.0f;
131+ r[k] = g[k] = b[k] = 0.0f;
124− int last = 0;
132+ last[k] = 0;
125− bool done = px.x >= frame.width || px.y0 >= frame.height;
133+ done[k] = px.x >= frame.width || px.y0 + k >= frame.height;
134+ }
135+ auto all_done = [&] {
136+ bool d = true;
137+#pragma unroll
138+ for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
139+ return d;
140+ };
126141
127142 __shared__ float4 batch_geom0[kRasterThreads];
128143 __shared__ float2 batch_geom1[kRasterThreads];
129144 __shared__ float4 batch_rgb_o[kRasterThreads];
130145 __shared__ uint8_t batch_mask[kRasterThreads];
146+ __shared__ uint32_t warp_contrib[kWarps][kWarps];
131147 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
148+ bool thread_done = all_done();
149+ int tile_compact = 0;
132150 for (int start = range.x; start < range.y; start += kRasterThreads) {
133− if (__syncthreads_count(done) == kRasterThreads) break;
134151 const int idx = start + tid;
152+ int32_t id = 0;
135153 if (idx < range.y) {
136− const int32_t id = point_list[idx];
154+ id = point_list[idx];
137155 const float4 p0 = __ldg(&splats[id].p0);
138156 const float4 p1 = __ldg(&splats[id].p1);
139157 const float4 p2 = __ldg(&splats[id].p2);
140158 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
141159 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
142160 batch_geom1[tid] = {c2.z, p2.w};
143161 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
144− batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
162+ batch_mask[tid] = static_cast<uint8_t>(WarpMask<kRows>(p0, p1));
145163 } else {
146164 batch_mask[tid] = 0u;
147165 }
148166 __syncthreads();
149167 const int count = min(kRasterThreads, range.y - start);
150− const int base = start - range.x;
151− bool warp_done = __all_sync(warp::kFullMask, done);
168+ bool warp_done = __all_sync(warp::kFullMask, thread_done);
152169 #pragma unroll
153170 for (int c = 0; c < kRasterThreads; c += 32) {
171+ uint32_t contrib = 0u;
154172 if (c < count && !warp_done) {
155173 uint32_t mask = __ballot_sync(
156174 warp::kFullMask,
157175 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
158176 while (mask != 0u) {
159177 const int bit = __ffs(mask) - 1;
160178 mask &= mask - 1u;
161179 const int j = c + bit;
180+ bool blended = false;
162− if (!done) {
181+ if (!thread_done) {
163182 const float4 g0 = batch_geom0[j];
164183 const float2 g1 = batch_geom1[j];
165184 const float dx = g0.x - fx;
166− const float dy = g0.y - fy;
185+ const float cx_dx2 = g0.z * dx * dx;
167− const float power =
186+ const float cy_dx = g0.w * dx;
187+#pragma unroll
188+ for (int k = 0; k < kRasterRows; ++k) {
189+ if (done[k]) continue;
190+ const float dy = g0.y - static_cast<float>(px.y0 + k);
168− g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
191+ const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
169− if (power <= 0.0f && power >= g1.y) {
192+ if (power > 0.0f || power < g1.y) continue;
170193 const float4 ro = batch_rgb_o[j];
171194 const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
172− if (alpha >= kMinAlpha) {
195+ if (alpha < kMinAlpha) continue;
173− const float next_t = t * (1.0f - alpha);
196+ const float next_t = t[k] * (1.0f - alpha);
174197⇥ if (next_t < kMinTransmittance) {
175− done = true;
198+ done[k] = 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;
199+ continue;
183200⇥ }
201+ const float w = alpha * t[k];
202+ r[k] += w * ro.x;
203+ g[k] += w * ro.y;
204+ b[k] += w * ro.z;
205+ t[k] = next_t;
206+ last[k] = -1 - j;
207+ blended = true;
184208⇥ }
209+ thread_done = all_done();
185210⇥ }
211+ if (__any_sync(warp::kFullMask, blended)) {
212+ contrib |= (1u << bit);
186213 }
187− if (__all_sync(warp::kFullMask, done)) {
214+ if (__all_sync(warp::kFullMask, thread_done)) {
188215 warp_done = true;
189216 break;
190217 }
191218 }
192− warp_done = __all_sync(warp::kFullMask, done);
219+ warp_done = __all_sync(warp::kFullMask, thread_done);
220+ }
221+ if (lane == 0) {
222+ warp_contrib[warp][c / 32] = contrib;
193223 }
194224 }
225+ const bool block_done =
226+ (__syncthreads_count(thread_done) == kRasterThreads);
227+ uint32_t chunk_any[kWarps];
228+ int chunk_prefix[kWarps];
229+ int batch_active = 0;
230+#pragma unroll
231+ for (int u = 0; u < kWarps; ++u) {
232+ uint32_t any_u = 0u;
233+#pragma unroll
234+ for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
235+ any_u |= warp_contrib[w_idx][u];
195236⇥ }
196− if (px.x < frame.width && px.y0 < frame.height) {
237+ chunk_any[u] = any_u;
238+ chunk_prefix[u] = tile_compact + batch_active;
239+ batch_active += __popc(any_u);
240+ }
241+ if ((chunk_any[warp] >> lane) & 1u) {
242+ uint32_t cmask = 0u;
243+#pragma unroll
244+ for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
245+ cmask |= ((warp_contrib[w_idx][warp] >> lane) & 1u) << w_idx;
246+ }
247+ const int dst =
248+ chunk_prefix[warp] + __popc(chunk_any[warp] & ((1u << lane) - 1u));
249+ point_list[range.x + dst] =
250+ static_cast<int32_t>((cmask << 24) | static_cast<uint32_t>(id));
251+ }
252+#pragma unroll
253+ for (int k = 0; k < kRasterRows; ++k) {
254+ if (last[k] < 0) {
255+ const int j = -1 - last[k];
256+ const int u = j >> 5;
257+ const int b_idx = j & 31;
258+ last[k] =
259+ chunk_prefix[u] + __popc(chunk_any[u] & ((1u << b_idx) - 1u)) + 1;
260+ }
261+ }
262+ tile_compact += batch_active;
263+ if (block_done) break;
264+ }
197265⇥ const int pixels = frame.width * frame.height;
266+#pragma unroll
267+ for (int k = 0; k < kRasterRows; ++k) {
268+ const int y = px.y0 + k;
269+ if (px.x >= frame.width || y >= frame.height) continue;
198− const int p = px.y0 * frame.width + px.x;
270+ const int p = y * frame.width + px.x;
199− image[0 * pixels + p] = r;
271+ image[0 * pixels + p] = r[k];
200− image[1 * pixels + p] = g;
272+ image[1 * pixels + p] = g[k];
201− image[2 * pixels + p] = b;
273+ image[2 * pixels + p] = b[k];
202− final_t[p] = t;
274+ final_t[p] = t[k];
203− n_contrib[p] = last;
275+ n_contrib[p] = last[k];
204276 }
205277 }
206278
207279 template __global__ void PreprocessKernel<Culling::kRect, float>(
↕ Expand 8 unchanged lines (lines 280–287) · in RasterForwardKernel
208280 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
209281 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
210282 template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
211283 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
212284 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
213285 template __global__ void PreprocessKernel<Culling::kEllipse, float>(
214286 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
215287 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
216288 template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
217289 GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
218290 uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
219291
292+template __global__ void RasterForwardKernel<1>(
293+ const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
294+ Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
295+template __global__ void RasterForwardKernel<2>(
296+ const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
297+ Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
298+template __global__ void RasterForwardKernel<4>(
299+ const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
300+ Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
301+
220302 } // namespace gs
csrc/backward.h +7 −4
s03_sh_streaming/csrc/backward.h
s04_raster_pruning/csrc/backward.h
↕ Expand 14 unchanged lines (lines 1–14)
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 "auxiliary.h"15#include "auxiliary.h"
1616
17namespace gs {17namespace gs {
1818
19template <int kRows>
19__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)20__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 6)
20 RasterBackwardKernel(const int2* __restrict__ ranges,21 RasterBackwardKernel(const int2* __restrict__ ranges,
21 const int32_t* __restrict__ point_list,22 const int32_t* __restrict__ point_list,
22 const Splat* __restrict__ splats, Frame frame,23 const Splat* __restrict__ splats, Frame frame,
23 const float* __restrict__ final_t,24 const float* __restrict__ final_t,
24 const int32_t* __restrict__ n_contrib,25 const int32_t* __restrict__ n_contrib,
25 const float* __restrict__ dl_dimage,26 const float* __restrict__ dl_dimage,
26 float* __restrict__ splat_grads);27 float* __restrict__ splat_grads,
28 uint8_t* __restrict__ splat_touched);
2729
28template <typename ShT>30template <typename ShT>
29__global__ void __launch_bounds__(kGaussianThreads)31__global__ void __launch_bounds__(kGaussianThreads)
30 PreprocessBackwardKernel(GaussianParams params, Frame frame,32 PreprocessBackwardKernel(GaussianParams params, Frame frame,
31 const Splat* __restrict__ splats,
32 const float* __restrict__ splat_grads,33 const float* __restrict__ splat_grads,
34 const uint8_t* __restrict__ splat_touched,
33 GaussianGrads grads);35 GaussianGrads grads);
3436
35template <typename ShT>37template <typename ShT>
36cudaError_t LaunchPreprocessBackward(const GaussianParams& params,38cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
37 const Frame& frame, const Splat* splats,39 const Frame& frame,
38 const float* splat_grads,40 const float* splat_grads,
41 const uint8_t* splat_touched,
39 const GaussianGrads& grads,42 const GaussianGrads& grads,
40 cudaStream_t stream);43 cudaStream_t stream);
4144
42} // namespace gs45} // namespace gs
4346
44#endif47#endif
↕ Expand 14 unchanged lines (lines 1–14)
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
1515 #include "auxiliary.h"
1616
1717 namespace gs {
1818
19+template <int kRows>
19−__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
20+__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 6)
2021 RasterBackwardKernel(const int2* __restrict__ ranges,
2122 const int32_t* __restrict__ point_list,
2223 const Splat* __restrict__ splats, Frame frame,
2324 const float* __restrict__ final_t,
2425 const int32_t* __restrict__ n_contrib,
2526 const float* __restrict__ dl_dimage,
26− float* __restrict__ splat_grads);
27+ float* __restrict__ splat_grads,
28+ uint8_t* __restrict__ splat_touched);
2729
2830 template <typename ShT>
2931 __global__ void __launch_bounds__(kGaussianThreads)
3032 PreprocessBackwardKernel(GaussianParams params, Frame frame,
31− const Splat* __restrict__ splats,
3233 const float* __restrict__ splat_grads,
34+ const uint8_t* __restrict__ splat_touched,
3335 GaussianGrads grads);
3436
3537 template <typename ShT>
3638 cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
37− const Frame& frame, const Splat* splats,
39+ const Frame& frame,
3840 const float* splat_grads,
41+ const uint8_t* splat_touched,
3942 const GaussianGrads& grads,
4043 cudaStream_t stream);
4144
4245 } // namespace gs
4346
4447 #endif
csrc/backward.cu.cc +118 −65
s03_sh_streaming/csrc/backward.cu.cc
s04_raster_pruning/csrc/backward.cu.cc
↕ Expand 195 unchanged lines (lines 1–195) · in ProjectVjp
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 "auxiliary.h"14#include "auxiliary.h"
1515
16namespace gs {16namespace gs {
1717
18namespace sh {18namespace sh {
19// Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly19// Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly
20// to the row `sh_out` (same storage type), group by group.20// to the row `sh_out` (same storage type), group by group.
21template <typename ShT>21template <typename ShT>
22GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,22GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,
23 ShT* __restrict__ sh_out) {23 ShT* __restrict__ sh_out) {
24 float b[kShCoeffs];24 float b[kShCoeffs];
25 Basis(d, b);25 Basis(d, b);
26 const float3 raw = EvalRow(sh_in, b);26 const float3 raw = EvalRow(sh_in, b);
27 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,27 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
28 raw.y < 0.0f ? 0.0f : g_rgb.y,28 raw.y < 0.0f ? 0.0f : g_rgb.y,
29 raw.z < 0.0f ? 0.0f : g_rgb.z};29 raw.z < 0.0f ? 0.0f : g_rgb.z};
30 const float x = d.x, y = d.y, z = d.z;30 const float x = d.x, y = d.y, z = d.z;
31 const float xx = x * x, yy = y * y, zz = z * z;31 const float xx = x * x, yy = y * y, zz = z * z;
32 const float xy = x * y, xz = x * z, yz = y * z;32 const float xy = x * y, xz = x * z, yz = y * z;
33 float3 ddir = {0.0f, 0.0f, 0.0f};33 float3 ddir = {0.0f, 0.0f, 0.0f};
34#pragma unroll34#pragma unroll
35 for (int grp = 0; grp < 4; ++grp) {35 for (int grp = 0; grp < 4; ++grp) {
36 const int k0 = 4 * grp;36 const int k0 = 4 * grp;
37 float v[12];37 float v[12];
38 ShIo<ShT>::LoadGroup(sh_in, grp, v);38 ShIo<ShT>::LoadGroup(sh_in, grp, v);
39 const float s0 = v[0] * drgb.x + v[1] * drgb.y + v[2] * drgb.z;39 const float s0 = v[0] * drgb.x + v[1] * drgb.y + v[2] * drgb.z;
40 const float s1 = v[3] * drgb.x + v[4] * drgb.y + v[5] * drgb.z;40 const float s1 = v[3] * drgb.x + v[4] * drgb.y + v[5] * drgb.z;
41 const float s2 = v[6] * drgb.x + v[7] * drgb.y + v[8] * drgb.z;41 const float s2 = v[6] * drgb.x + v[7] * drgb.y + v[8] * drgb.z;
42 const float s3 = v[9] * drgb.x + v[10] * drgb.y + v[11] * drgb.z;42 const float s3 = v[9] * drgb.x + v[10] * drgb.y + v[11] * drgb.z;
43 float out[12];43 float out[12];
44#pragma unroll44#pragma unroll
45 for (int q = 0; q < 4; ++q) {45 for (int q = 0; q < 4; ++q) {
46 out[3 * q] = b[k0 + q] * drgb.x;46 out[3 * q] = b[k0 + q] * drgb.x;
47 out[3 * q + 1] = b[k0 + q] * drgb.y;47 out[3 * q + 1] = b[k0 + q] * drgb.y;
48 out[3 * q + 2] = b[k0 + q] * drgb.z;48 out[3 * q + 2] = b[k0 + q] * drgb.z;
49 }49 }
50 ShIo<ShT>::StoreGroup(sh_out, grp, out);50 ShIo<ShT>::StoreGroup(sh_out, grp, out);
51 if (grp == 0) {51 if (grp == 0) {
52 ddir.y += s1 * (-kC1);52 ddir.y += s1 * (-kC1);
53 ddir.z += s2 * kC1;53 ddir.z += s2 * kC1;
54 ddir.x += s3 * (-kC1);54 ddir.x += s3 * (-kC1);
55 } else if (grp == 1) {55 } else if (grp == 1) {
56 ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};56 ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};
57 ddir += (s1 * kC2[1]) * float3{0.0f, z, y};57 ddir += (s1 * kC2[1]) * float3{0.0f, z, y};
58 ddir += (s2 * kC2[2]) * float3{-2.0f * x, -2.0f * y, 4.0f * z};58 ddir += (s2 * kC2[2]) * float3{-2.0f * x, -2.0f * y, 4.0f * z};
59 ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};59 ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};
60 } else if (grp == 2) {60 } else if (grp == 2) {
61 ddir += (s0 * kC2[4]) * float3{2.0f * x, -2.0f * y, 0.0f};61 ddir += (s0 * kC2[4]) * float3{2.0f * x, -2.0f * y, 0.0f};
62 ddir += (s1 * kC3[0]) * float3{6.0f * xy, 3.0f * (xx - yy), 0.0f};62 ddir += (s1 * kC3[0]) * float3{6.0f * xy, 3.0f * (xx - yy), 0.0f};
63 ddir += (s2 * kC3[1]) * float3{yz, xz, xy};63 ddir += (s2 * kC3[1]) * float3{yz, xz, xy};
64 ddir += (s3 * kC3[2]) *64 ddir += (s3 * kC3[2]) *
65 float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};65 float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};
66 } else {66 } else {
67 ddir += (s0 * kC3[3]) *67 ddir += (s0 * kC3[3]) *
68 float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};68 float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};
69 ddir += (s1 * kC3[4]) *69 ddir += (s1 * kC3[4]) *
70 float3{4.0f * zz - 3.0f * xx - yy, -2.0f * xy, 8.0f * xz};70 float3{4.0f * zz - 3.0f * xx - yy, -2.0f * xy, 8.0f * xz};
71 ddir += (s2 * kC3[5]) * float3{2.0f * xz, -2.0f * yz, xx - yy};71 ddir += (s2 * kC3[5]) * float3{2.0f * xz, -2.0f * yz, xx - yy};
72 ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};72 ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};
73 }73 }
74 }74 }
75 return ddir;75 return ddir;
76}76}
77} // namespace sh77} // namespace sh
7878
79// VJP of the projection for a touched Gaussian `i`; reads its SH row from79// VJP of the projection for a touched Gaussian `i`; reads its SH row from
80// `sh_in` and writes dL/dsh directly to `sh_out`.80// `sh_in` and writes dL/dsh directly to `sh_out`.
81// Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.81// Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
82template <typename ShT>82template <typename ShT>
83GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,83GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
84 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,84 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
85 const ProjectionGrad& g, float3& dmean_out,85 const ProjectionGrad& g, float3& dmean_out,
86 float3& dscale_out, float4& dquat_out) {86 float3& dscale_out, float4& dquat_out) {
87 const float3 mean = LoadMean(p, i);87 const float3 mean = LoadMean(p, i);
88 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],88 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
89 p.scales[3 * i + 2]};89 p.scales[3 * i + 2]};
90 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];90 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
9191
92 // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).92 // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
93 float3 dmean = {0.0f, 0.0f, 0.0f};93 float3 dmean = {0.0f, 0.0f, 0.0f};
94 {94 {
95 float inv_len;95 float inv_len;
96 const float3 dir = ViewDir(cam, mean, inv_len);96 const float3 dir = ViewDir(cam, mean, inv_len);
97 const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);97 const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
98 // d(v/|v|)/dv = (I - dir dir^T) / |v|.98 // d(v/|v|)/dv = (I - dir dir^T) / |v|.
99 dmean += inv_len * (ddir - dot(dir, ddir) * dir);99 dmean += inv_len * (ddir - dot(dir, ddir) * dir);
100 }100 }
101101
102 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).102 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
103 const float3 t = cam.ToView(mean);103 const float3 t = cam.ToView(mean);
104 const float tz = t.z;104 const float tz = t.z;
105 const float limx = kFovClamp * cam.TanHalfFovX();105 const float limx = kFovClamp * cam.TanHalfFovX();
106 const float limy = kFovClamp * cam.TanHalfFovY();106 const float limy = kFovClamp * cam.TanHalfFovY();
107 const float txtz = t.x / tz, tytz = t.y / tz;107 const float txtz = t.x / tz, tytz = t.y / tz;
108 const bool x_clamped = txtz < -limx || txtz > limx;108 const bool x_clamped = txtz < -limx || txtz > limx;
109 const bool y_clamped = tytz < -limy || tytz > limy;109 const bool y_clamped = tytz < -limy || tytz > limy;
110 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;110 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
111 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;111 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
112 const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);112 const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
113 const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);113 const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
114 const Mat3& w = cam.rotation;114 const Mat3& w = cam.rotation;
115 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);115 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
116 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);116 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
117117
118 const Mat3 r = QuatToRotation(quat);118 const Mat3 r = QuatToRotation(quat);
119 const Sym3 sigma = Covariance3D(r, scale);119 const Sym3 sigma = Covariance3D(r, scale);
120 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);120 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
121 const float cov_a = dot(m0, s0) + kLowPass;121 const float cov_a = dot(m0, s0) + kLowPass;
122 const float cov_b = dot(m0, s1);122 const float cov_b = dot(m0, s1);
123 const float cov_c = dot(m1, s1) + kLowPass;123 const float cov_c = dot(m1, s1) + kLowPass;
124124
125 // --- Conic -> 2D covariance: G_cov = -K G_K K.125 // --- Conic -> 2D covariance: G_cov = -K G_K K.
126 const float det = cov_a * cov_c - cov_b * cov_b;126 const float det = cov_a * cov_c - cov_b * cov_b;
127 const float inv = 1.0f / det;127 const float inv = 1.0f / det;
128 const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;128 const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
129 const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;129 const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
130 // P = G_K K.130 // P = G_K K.
131 const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;131 const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
132 const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;132 const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
133 // G_cov = -K P (symmetric).133 // G_cov = -K P (symmetric).
134 const float ca = -(ka * p00 + kb * p10);134 const float ca = -(ka * p00 + kb * p10);
135 const float cb = -(ka * p01 + kb * p11);135 const float cb = -(ka * p01 + kb * p11);
136 const float cc = -(kb * p01 + kc * p11);136 const float cc = -(kb * p01 + kc * p11);
137137
138 // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.138 // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
139 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);139 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
140 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);140 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
141 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.141 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
142 Sym3 dsig;142 Sym3 dsig;
143 dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;143 dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
144 dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;144 dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
145 dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;145 dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
146 dsig.xy =146 dsig.xy =
147 ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;147 ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
148 dsig.xz =148 dsig.xz =
149 ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;149 ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
150 dsig.yz =150 dsig.yz =
151 ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;151 ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
152152
153 // --- M = J W: dJ = dM W^T (only the nonzero entries of J).153 // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
154 const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));154 const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
155 const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));155 const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
156 const float tz2 = tz * tz, tz3 = tz2 * tz;156 const float tz2 = tz * tz, tz3 = tz2 * tz;
157 float3 dt;157 float3 dt;
158 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;158 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
159 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;159 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
160 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +160 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
161 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;161 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
162 // --- Screen-space mean u = fx x / z + cx.162 // --- Screen-space mean u = fx x / z + cx.
163 dt.x += cam.fx / tz * g.xy.x;163 dt.x += cam.fx / tz * g.xy.x;
164 dt.y += cam.fy / tz * g.xy.y;164 dt.y += cam.fy / tz * g.xy.y;
165 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;165 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
166 dmean += MulT(w, dt);166 dmean += MulT(w, dt);
167 dmean_out = dmean;167 dmean_out = dmean;
168168
169 // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.169 // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
170 const float sv[3] = {scale.x, scale.y, scale.z};170 const float sv[3] = {scale.x, scale.y, scale.z};
171 float ds[3];171 float ds[3];
172 Mat3 dr;172 Mat3 dr;
173#pragma unroll173#pragma unroll
174 for (int j = 0; j < 3; ++j) {174 for (int j = 0; j < 3; ++j) {
175 const float3 r_col = r.col(j);175 const float3 r_col = r.col(j);
176 const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);176 const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
177 ds[j] = dot(da_col, r_col);177 ds[j] = dot(da_col, r_col);
178 dr.m[0][j] = da_col.x * sv[j];178 dr.m[0][j] = da_col.x * sv[j];
179 dr.m[1][j] = da_col.y * sv[j];179 dr.m[1][j] = da_col.y * sv[j];
180 dr.m[2][j] = da_col.z * sv[j];180 dr.m[2][j] = da_col.z * sv[j];
181 }181 }
182 dscale_out = {ds[0], ds[1], ds[2]};182 dscale_out = {ds[0], ds[1], ds[2]};
183183
184 // --- R(q) -> q.184 // --- R(q) -> q.
185 const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;185 const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
186 const auto& G = dr.m;186 const auto& G = dr.m;
187 dquat_out = {187 dquat_out = {
188 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -188 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
189 qy * G[2][0] + qx * G[2][1]),189 qy * G[2][0] + qx * G[2][1]),
190 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -190 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
191 qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),191 qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
192 2.0f *192 2.0f *
193 (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +193 (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
194 qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),194 qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
195 2.0f *195 2.0f *
196 (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -196 (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
197 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};197 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
198}198}
199199
200// Reverse traversal per pixel; accumulates splat gradients via Half-Warp200// Reverse traversal per pixel; accumulates splat gradients with atomics.
201// Two-Splat Folding and Halving Register Reduce-Scatter.
201template <int kRows>
202__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 4)202__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 4)
203 RasterBackwardKernel(const int2* __restrict__ ranges,203 RasterBackwardKernel(const int2* __restrict__ ranges,
204 const int32_t* __restrict__ point_list,204 const int32_t* __restrict__ point_list,
205 const Splat* __restrict__ splats, Frame frame,205 const Splat* __restrict__ splats, Frame frame,
206 const float* __restrict__ final_t,206 const float* __restrict__ final_t,
207 const int32_t* __restrict__ n_contrib,207 const int32_t* __restrict__ n_contrib,
208 const float* __restrict__ dl_dimage,208 const float* __restrict__ dl_dimage,
209 float* __restrict__ splat_grads) {209 float* __restrict__ splat_grads,
210 uint8_t* __restrict__ splat_touched) {
210 using L = RasterLayout<1>;211 using L = RasterLayout<kRows>;
212 constexpr int kRasterRows = kRows;
211 constexpr int kRasterThreads = L::kThreads;213 constexpr int kRasterThreads = L::kThreads;
212 constexpr int kWarps = L::kWarps;214 constexpr int kWarps = L::kWarps;
213 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;215 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
214 const int tid = threadIdx.x;216 const int tid = threadIdx.x;
215 const int warp = tid / 32;217 const int warp = tid / 32;
216 const int lane = tid % 32;218 const int lane = tid % 32;
217 const RasterPixels px = RasterPixels::For<1>(tid);219 const RasterPixels px = RasterPixels::For<kRows>(tid);
218 const float fx = static_cast<float>(px.x);220 const float fx = static_cast<float>(px.x);
219 const float fy = static_cast<float>(px.y0);
220 const int2 range = ranges[tile];221 const int2 range = ranges[tile];
221222
222 const bool inside = px.x < frame.width && px.y0 < frame.height;
223 struct PixelState {
224 float t;
225 float3 dpix;
226 float accum; // dot(color accumulated behind, dpix)
227 int last;
228 };
229 PixelState st[kRasterRows];
223 const int pixels = frame.width * frame.height;230 const int pixels = frame.width * frame.height;
231 int thread_last = 0;
232#pragma unroll
233 for (int k = 0; k < kRasterRows; ++k) {
234 const int y = px.y0 + k;
235 const bool inside = px.x < frame.width && y < frame.height;
224 const int p = px.y0 * frame.width + px.x;236 const int p = y * frame.width + px.x;
225 float t = inside ? final_t[p] : 0.0f;237 st[k].t = inside ? final_t[p] : 0.0f;
226 const float3 dpix =238 st[k].dpix =
227 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],239 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
228 dl_dimage[2 * pixels + p]}240 dl_dimage[2 * pixels + p]}
229 : float3{0.0f, 0.0f, 0.0f};241 : float3{0.0f, 0.0f, 0.0f};
230 float accum = 0.0f; // dot(color accumulated behind, dpix)
242 st[k].accum = 0.0f;
231 const int last = inside ? n_contrib[p] : 0;243 st[k].last = inside ? n_contrib[p] : 0;
244 thread_last = max(thread_last, st[k].last);
245 }
232246
233 int warp_max_last = last;247 int warp_max_last = thread_last;
234#pragma unroll248#pragma unroll
235 for (int offset = 16; offset > 0; offset /= 2) {249 for (int offset = 16; offset > 0; offset /= 2) {
236 warp_max_last = max(250 warp_max_last = max(
237 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));251 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
↕ Expand 15 unchanged lines (lines 252–266) · in PixelState
238 }252 }
239 __shared__ int shared_max_last[kWarps];253 __shared__ int shared_max_last[kWarps];
240 if (lane == 0) shared_max_last[warp] = warp_max_last;254 if (lane == 0) shared_max_last[warp] = warp_max_last;
241 __syncthreads();255 __syncthreads();
242 int block_max_last = 0;256 int block_max_last = 0;
243#pragma unroll257#pragma unroll
244 for (int w = 0; w < kWarps; ++w) {258 for (int w = 0; w < kWarps; ++w) {
245 block_max_last = max(block_max_last, shared_max_last[w]);259 block_max_last = max(block_max_last, shared_max_last[w]);
246 }260 }
247 if (block_max_last <= 0) return;261 if (block_max_last <= 0) return;
248262
249 __shared__ float4 batch_geom0[kRasterThreads];263 __shared__ float4 batch_geom0[kRasterThreads];
250 __shared__ float2 batch_geom1[kRasterThreads];264 __shared__ float2 batch_geom1[kRasterThreads];
251 __shared__ float4 batch_rgb_o[kRasterThreads];265 __shared__ float4 batch_rgb_o[kRasterThreads];
252 __shared__ uint32_t batch_id[kRasterThreads];266 __shared__ uint32_t batch_id[kRasterThreads];
253 __shared__ uint8_t batch_mask[kRasterThreads];267 __shared__ uint8_t batch_mask[kRasterThreads];
254 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);268 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
255 int scatter = -1;269 int scatter = -1;
256 int scatter2 = -1;270 int scatter2 = -1;
257 {
271 if constexpr (L::kWarpWidth == 8) {
258 const int sub_lane = lane & 15;272 const int sub_lane = lane & 15;
259 const int row = sub_lane >> 3, col = sub_lane & 7;273 const int row = sub_lane >> 3, col = sub_lane & 7;
260 if (col == 0) {274 if (col == 0) {
261 scatter2 = row;275 scatter2 = row;
↕ Expand 13 unchanged lines (lines 276–288) · in PixelState
262 } else if (col == 2) {276 } else if (col == 2) {
263 scatter2 = 2 + row;277 scatter2 = 2 + row;
264 } else if (col == 4) {278 } else if (col == 4) {
265 scatter2 = row == 0 ? 4 : 7;279 scatter2 = row == 0 ? 4 : 7;
266 } else if (col == 6) {280 } else if (col == 6) {
267 scatter2 = row == 0 ? 6 : 8;281 scatter2 = row == 0 ? 6 : 8;
268 } else if (col == 1 && row == 0) {282 } else if (col == 1 && row == 0) {
269 scatter2 = 5;283 scatter2 = 5;
270 }284 }
271 const int full_row = lane >> 3, full_col = lane & 7;285 const int full_row = lane >> 3, full_col = lane & 7;
272 if (full_col == 0) {286 if (full_col == 0) {
273 scatter = full_row;287 scatter = full_row;
274 } else if (full_col == 2) {288 } else if (full_col == 2) {
275 scatter = full_row == 0 ? 4 : 5 + full_row;289 scatter = full_row == 0 ? 4 : 5 + full_row;
276 } else if (full_col == 4 && full_row == 0) {290 } else if (full_col == 4 && full_row == 0) {
277 scatter = 5;291 scatter = 5;
278 }292 }
293 } else {
294 scatter = warp::ScatterIndex(kSplatGradFloats, lane);
279 }295 }
280296
281 auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,297 auto eval_splat = [&](int j, float4 g0, float2 g1, float4 ro, int entry,
282 float& ga_sum, float& ga_dy, float& ga_dy2,298 float& dx, float& ga_sum, float& ga_dy, float& ga_dy2,
283 float3& drgb) -> bool {299 float3& drgb) {
284 dx = g0.x - fx;300 dx = g0.x - fx;
285 ga_sum = 0.0f;301 ga_sum = 0.0f;
286 ga_dy = 0.0f;302 ga_dy = 0.0f;
287 ga_dy2 = 0.0f;303 ga_dy2 = 0.0f;
288 drgb = {0.0f, 0.0f, 0.0f};304 drgb = {0.0f, 0.0f, 0.0f};
289 if (entry < last) {305 if (entry < thread_last) {
306 const float cx_dx2 = g0.z * dx * dx;
290 const float dy = g0.y - fy;307 const float cy_dx = g0.w * dx;
308#pragma unroll
309 for (int k = 0; k < kRasterRows; ++k) {
310 PixelState& ps = st[k];
311 if (entry >= ps.last) continue;
312 const float dy = g0.y - static_cast<float>(px.y0 + k);
291 const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;313 const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
292 if (power <= 0.0f && power >= g1.y) {314 if (power > 0.0f || power < g1.y) continue;
293 const float g_val = Exp2Approx(power);315 const float g_val = Exp2Approx(power);
294 const float alpha = fminf(kMaxAlpha, ro.w * g_val);316 const float alpha = fminf(kMaxAlpha, ro.w * g_val);
295 if (alpha >= kMinAlpha) {317 if (alpha < kMinAlpha) continue;
296 t = __fdividef(t, 1.0f - alpha);318 ps.t = __fdividef(ps.t, 1.0f - alpha);
297 const float w = alpha * t;319 const float w = alpha * ps.t;
320 const float color_dot =
298 const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;321 ro.x * ps.dpix.x + ro.y * ps.dpix.y + ro.z * ps.dpix.z;
299 const float diff = color_dot - accum;322 const float diff = color_dot - ps.accum;
300 const float dl_dalpha = t * diff;323 const float dl_dalpha = ps.t * diff;
301 accum += alpha * diff;324 ps.accum += alpha * diff;
302 const float ga = g_val * dl_dalpha;325 const float ga = g_val * dl_dalpha;
303 ga_sum = ga;326 ga_sum += ga;
304 const float ga_y = ga * dy;327 const float ga_y = ga * dy;
305 ga_dy = ga_y;328 ga_dy += ga_y;
306 ga_dy2 = ga_y * dy;329 ga_dy2 += ga_y * dy;
307 drgb = w * dpix;330 drgb += w * ps.dpix;
308 return true;
309 }331 }
310 }332 }
311 }
312 return false;
313 };333 };
314 auto raw_conic = [](float4 g0, float2 g1) {334 auto raw_conic = [](float4 g0, float2 g1) {
315 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,335 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
316 kInvPower2Scale * g1.x};336 kInvPower2Scale * g1.x};
317 };337 };
318338
319 for (int end = block_max_last; end > 0; end -= kRasterThreads) {339 for (int end = block_max_last; end > 0; end -= kRasterThreads) {
320 const int e = end - 1 - tid;340 const int e = end - 1 - tid;
321 if (e >= 0) {341 if (e >= 0) {
322 const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);342 const uint32_t packed = static_cast<uint32_t>(point_list[range.x + e]);
343 const uint32_t id = packed & 0x00ffffffu;
344 splat_touched[id] = 1u;
323 const float4 p0 = __ldg(&splats[id].p0);345 const float4 p0 = __ldg(&splats[id].p0);
324 const float4 p1 = __ldg(&splats[id].p1);346 const float4 p1 = __ldg(&splats[id].p1);
325 const float4 p2 = __ldg(&splats[id].p2);347 const float4 p2 = __ldg(&splats[id].p2);
326 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);348 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
327 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};349 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
328 batch_geom1[tid] = {c2.z, p2.w};350 batch_geom1[tid] = {c2.z, p2.w};
329 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};351 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
330 batch_id[tid] = id;352 batch_id[tid] = id;
331 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));353 batch_mask[tid] = static_cast<uint8_t>(packed >> 24);
332 } else {354 } else {
333 batch_mask[tid] = 0u;355 batch_mask[tid] = 0u;
334 }356 }
335 __syncthreads();357 __syncthreads();
336 const int count = min(kRasterThreads, end);358 const int count = min(kRasterThreads, end);
337#pragma unroll359#pragma unroll
338 for (int c = 0; c < kRasterThreads; c += 32) {360 for (int c = 0; c < kRasterThreads; c += 32) {
339 if (c < count && end - 32 - c < warp_max_last) {361 if (c < count && end - 32 - c < warp_max_last) {
340 uint32_t mask = __ballot_sync(362 uint32_t mask = __ballot_sync(
341 warp::kFullMask,363 warp::kFullMask,
342 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);364 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
365 if constexpr (L::kWarpWidth == 8) {
343 while ((mask & (mask - 1u)) != 0u) {366 while ((mask & (mask - 1u)) != 0u) {
344 const int bit0 = __ffs(mask) - 1;367 const int bit0 = __ffs(mask) - 1;
345 mask &= mask - 1u;368 mask &= mask - 1u;
346 const int bit1 = __ffs(mask) - 1;369 const int bit1 = __ffs(mask) - 1;
347 mask &= mask - 1u;370 mask &= mask - 1u;
348 const int j0 = c + bit0, j1 = c + bit1;371 const int j0 = c + bit0, j1 = c + bit1;
349 const float4 g0_0 = batch_geom0[j0];372 const float4 g0_0 = batch_geom0[j0];
350 const float2 g1_0 = batch_geom1[j0];373 const float2 g1_0 = batch_geom1[j0];
351 const float4 ro_0 = batch_rgb_o[j0];374 const float4 ro_0 = batch_rgb_o[j0];
352 const float4 g0_1 = batch_geom0[j1];375 const float4 g0_1 = batch_geom0[j1];
353 const float2 g1_1 = batch_geom1[j1];376 const float2 g1_1 = batch_geom1[j1];
354 const float4 ro_1 = batch_rgb_o[j1];377 const float4 ro_1 = batch_rgb_o[j1];
355378
356 float dx0, ga_sum0, ga_dy0, ga_dy2_0;379 float dx0, ga_sum0, ga_dy0, ga_dy2_0;
357 float3 drgb0;380 float3 drgb0;
358 const bool hit0 = eval_splat(g0_0, g1_0, ro_0, end - 1 - j0, dx0,381 eval_splat(j0, g0_0, g1_0, ro_0, end - 1 - j0, dx0, ga_sum0, ga_dy0,
359 ga_sum0, ga_dy0, ga_dy2_0, drgb0);382 ga_dy2_0, drgb0);
360383
361 float dx1, ga_sum1, ga_dy1, ga_dy2_1;384 float dx1, ga_sum1, ga_dy1, ga_dy2_1;
362 float3 drgb1;385 float3 drgb1;
363 const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,386 eval_splat(j1, g0_1, g1_1, ro_1, end - 1 - j1, dx1, ga_sum1, ga_dy1,
364 ga_sum1, ga_dy1, ga_dy2_1, drgb1);387 ga_dy2_1, drgb1);
365 // As in s01: a pair that no pixel of the warp blends has zero
366 // gradients, so skip its reduction and atomics entirely.
367 if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;
368388
369 const bool hi16 = (lane & 16) != 0;389 const bool hi16 = (lane & 16) != 0;
370 const float opacity = hi16 ? ro_1.w : ro_0.w;390 const float opacity = hi16 ? ro_1.w : ro_0.w;
371 const float dx = hi16 ? dx1 : dx0;391 const float dx = hi16 ? dx1 : dx0;
372 const float3 cn =392 const float3 cn =
373 hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);393 hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);
374394
375 auto fold16 = [&](float v0, float v1) {395 auto fold16 = [&](float v0, float v1) {
376 const float send = hi16 ? v0 : v1;396 const float send = hi16 ? v0 : v1;
377 const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);397 const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);
378 return (hi16 ? v1 : v0) + recv;398 return (hi16 ? v1 : v0) + recv;
379 };399 };
380 float ga_sum = fold16(ga_sum0, ga_sum1);400 float ga_sum = fold16(ga_sum0, ga_sum1);
381 float ga_dy = fold16(ga_dy0, ga_dy1);401 float ga_dy = fold16(ga_dy0, ga_dy1);
382 float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),402 float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),
383 fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),403 fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),
384 fold16(drgb0.z, drgb1.z)};404 fold16(drgb0.z, drgb1.z)};
385 warp::internal::HalvingStep<4, 8>(b4, lane);405 warp::internal::HalvingStep<4, 8>(b4, lane);
386 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);406 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
387 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);407 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
388408
389 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;409 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
390 const float p_dx = dx * p_sum;410 const float p_dx = dx * p_sum;
391 const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)411 const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
392 : -(cn.y * p_dx + cn.z * p_dy);412 : -(cn.y * p_dx + cn.z * p_dy);
413 const float v_dx1 =
393 const float v_dx1 = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;414 (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
394 float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};415 float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};
395 warp::internal::HalvingStep<4, 4>(v4, lane);416 warp::internal::HalvingStep<4, 4>(v4, lane);
396 warp::internal::HalvingStep<2, 2>(v4, lane);417 warp::internal::HalvingStep<2, 2>(v4, lane);
397 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);418 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);
398 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);419 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);
399 const float v_sum =420 const float v_sum =
400 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);421 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);
401 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);422 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);
402 const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;423 const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
403 if (scatter2 >= 0 && out_grad != 0.0f) {424 if (scatter2 >= 0 && out_grad != 0.0f) {
404 const uint32_t target_id = batch_id[hi16 ? j1 : j0];425 const uint32_t target_id = batch_id[hi16 ? j1 : j0];
405 atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,426 atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,
406 out_grad);427 out_grad);
407 }428 }
408 }429 }
430 }
409 while (mask != 0u) {431 while (mask != 0u) {
410 const int bit = __ffs(mask) - 1;432 const int bit = __ffs(mask) - 1;
411 mask &= mask - 1u;433 mask &= mask - 1u;
412 const int j = c + bit;434 const int j = c + bit;
413 const float4 g0 = batch_geom0[j];435 const float4 g0 = batch_geom0[j];
414 const float2 g1 = batch_geom1[j];436 const float2 g1 = batch_geom1[j];
415 const float4 ro = batch_rgb_o[j];437 const float4 ro = batch_rgb_o[j];
416 float dx, ga_sum, ga_dy, ga_dy2;438 float dx, ga_sum, ga_dy, ga_dy2;
417 float3 drgb;439 float3 drgb;
418 const bool hit =
419 eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);440 eval_splat(j, g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2,
420 if (!__any_sync(warp::kFullMask, hit)) continue;
441 drgb);
421 const float opacity = ro.w;442 const float opacity = ro.w;
422 const float3 cn = raw_conic(g0, g1);443 const float3 cn = raw_conic(g0, g1);
444 float out_grad = 0.0f;
445 if constexpr (L::kWarpWidth == 8) {
423 float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};446 float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};
424 warp::internal::HalvingStep<4, 16>(b4, lane);447 warp::internal::HalvingStep<4, 16>(b4, lane);
425 warp::internal::HalvingStep<2, 8>(b4, lane);448 warp::internal::HalvingStep<2, 8>(b4, lane);
426 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);449 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);
427 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);450 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);
428 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);451 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
429 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);452 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
430 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;453 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
431 const float p_dx = dx * p_sum;454 const float p_dx = dx * p_sum;
432 float v_dx;455 float v_dx;
433 if ((lane & 16) == 0) {456 if ((lane & 16) == 0) {
434 v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)457 v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
435 : -(cn.y * p_dx + cn.z * p_dy);458 : -(cn.y * p_dx + cn.z * p_dy);
436 } else {459 } else {
437 v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;460 v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
438 }461 }
439 float v[3] = {v_dx, b4[0], ga_sum};462 float v[3] = {v_dx, b4[0], ga_sum};
440 warp::internal::HalvingStep<3, 4>(v, lane);463 warp::internal::HalvingStep<3, 4>(v, lane);
441 warp::internal::HalvingStep<2, 2>(v, lane);464 warp::internal::HalvingStep<2, 2>(v, lane);
442 const float out_grad =
443 v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);465 out_grad = v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
466 } else {
467 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
468 const float p_dx = dx * p_sum;
469 float gr[kSplatGradFloats] = {-(cn.x * p_dx + cn.y * p_dy), // xy
470 -(cn.y * p_dx + cn.z * p_dy),
471 -0.5f * dx * p_dx, // conic
472 -dx * p_dy,
473 -0.5f * opacity * ga_dy2,
474 ga_sum, // opacity
475 drgb.x,
476 drgb.y,
477 drgb.z}; // rgb
478 warp::ReduceScatter(gr, lane);
479 out_grad = gr[0];
480 }
444 if (scatter >= 0 && out_grad != 0.0f) {481 if (scatter >= 0 && out_grad != 0.0f) {
445 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,482 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
446 out_grad);483 out_grad);
447 }484 }
448 }485 }
449 }486 }
450 }487 }
451 __syncthreads();488 __syncthreads();
452 }489 }
453}490}
454491
455// ---------------------------------------------------------------------------492// ---------------------------------------------------------------------------
456// Preprocess backward: splat gradients -> parameter gradients (dense).493// Preprocess backward: splat gradients -> parameter gradients (dense).
457// Gaussians stored as culled (radius 0, i.e. no tile pairs) are untouched.
494// Uses the 1-byte `splat_touched` flag written by `RasterBackwardKernel` to
495// skip reading `splat_grads` and `params.sh` for untouched Gaussians, and
458// Uses warp-coalesced 16-byte stores (0 shared memory) to zero untouched SH496// zeros untouched SH rows with warp-coalesced float4 stores (0 shared memory).
459// rows and streams active SH rows in 4 groups of 12 floats via `ShIo<ShT>`.
460497
461template <typename ShT>498template <typename ShT>
462__global__ void __launch_bounds__(kGaussianThreads)499__global__ void __launch_bounds__(kGaussianThreads)
463 PreprocessBackwardKernel(GaussianParams params, Frame frame,500 PreprocessBackwardKernel(GaussianParams params, Frame frame,
464 const Splat* __restrict__ splats,
465 const float* __restrict__ splat_grads,501 const float* __restrict__ splat_grads,
502 const uint8_t* __restrict__ splat_touched,
466 GaussianGrads grads) {503 GaussianGrads grads) {
467 constexpr int kShVec16 = ShIo<ShT>::kRowVec16; // 16-byte stores per row.504 constexpr int kShVec16 = ShIo<ShT>::kRowVec16; // 16-byte stores per row.
468 const int i = blockIdx.x * blockDim.x + threadIdx.x;505 const int i = blockIdx.x * blockDim.x + threadIdx.x;
469 const int lane = threadIdx.x % 32;506 const int lane = threadIdx.x % 32;
470 bool touched = (i < params.n && splats[i].radius() != 0);507 bool touched = (i < params.n && splat_touched[i] != 0u);
471 const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);508 const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);
472 if (init_mask == 0u) {509 if (init_mask == 0u) {
473 if (i < params.n) {510 if (i < params.n) {
474 grads.means[3 * i] = 0.0f;511 grads.means[3 * i] = 0.0f;
↕ Expand 87 unchanged lines (lines 512–598) · in PreprocessBackwardKernel
475 grads.means[3 * i + 1] = 0.0f;512 grads.means[3 * i + 1] = 0.0f;
476 grads.means[3 * i + 2] = 0.0f;513 grads.means[3 * i + 2] = 0.0f;
477 grads.scales[3 * i] = 0.0f;514 grads.scales[3 * i] = 0.0f;
478 grads.scales[3 * i + 1] = 0.0f;515 grads.scales[3 * i + 1] = 0.0f;
479 grads.scales[3 * i + 2] = 0.0f;516 grads.scales[3 * i + 2] = 0.0f;
480 reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};517 reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};
481 grads.opacities[i] = 0.0f;518 grads.opacities[i] = 0.0f;
482 }519 }
483 const int warp_i0 =520 const int warp_i0 =
484 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;521 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
485 const int warp_n = max(0, min(32, params.n - warp_i0));522 const int warp_n = max(0, min(32, params.n - warp_i0));
486 uint4* warp_sh_dst =523 uint4* warp_sh_dst =
487 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));524 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
488 const uint4 zero = {0u, 0u, 0u, 0u};525 const uint4 zero = {0u, 0u, 0u, 0u};
489 if (warp_n == 32) {526 if (warp_n == 32) {
490#pragma unroll527#pragma unroll
491 for (int s = 0; s < kShVec16; ++s) {528 for (int s = 0; s < kShVec16; ++s) {
492 warp_sh_dst[s * 32 + lane] = zero;529 warp_sh_dst[s * 32 + lane] = zero;
493 }530 }
494 } else {531 } else {
495 for (int k = lane; k < warp_n * kShVec16; k += 32) {532 for (int k = lane; k < warp_n * kShVec16; k += 32) {
496 warp_sh_dst[k] = zero;533 warp_sh_dst[k] = zero;
497 }534 }
498 }535 }
499 return;536 return;
500 }537 }
501538
502 float3 dmean = {0.0f, 0.0f, 0.0f};539 float3 dmean = {0.0f, 0.0f, 0.0f};
503 float3 dscale = {0.0f, 0.0f, 0.0f};540 float3 dscale = {0.0f, 0.0f, 0.0f};
504 float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};541 float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};
505 float dopacity = 0.0f;542 float dopacity = 0.0f;
506 if (touched) {543 if (touched) {
507 float g[kSplatGradFloats];544 float g[kSplatGradFloats];
508 bool nonzero = false;545 bool nonzero = false;
509#pragma unroll546#pragma unroll
510 for (int q = 0; q < kSplatGradFloats; ++q) {547 for (int q = 0; q < kSplatGradFloats; ++q) {
511 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);548 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
512 nonzero = nonzero || g[q] != 0.0f;549 nonzero = nonzero || g[q] != 0.0f;
513 }550 }
514 touched = nonzero;551 touched = nonzero;
515 if (touched) {552 if (touched) {
516 ProjectionGrad pg;553 ProjectionGrad pg;
517 pg.xy = {g[0], g[1]};554 pg.xy = {g[0], g[1]};
518 pg.conic = {g[2], g[3], g[4]};555 pg.conic = {g[2], g[3], g[4]};
519 pg.opacity = g[5];556 pg.opacity = g[5];
520 pg.rgb = {g[6], g[7], g[8]};557 pg.rgb = {g[6], g[7], g[8]};
521 dopacity = g[5];558 dopacity = g[5];
522 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),559 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
523 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);560 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
524 }561 }
525 }562 }
526 if (i < params.n) {563 if (i < params.n) {
527 grads.means[3 * i] = dmean.x;564 grads.means[3 * i] = dmean.x;
528 grads.means[3 * i + 1] = dmean.y;565 grads.means[3 * i + 1] = dmean.y;
529 grads.means[3 * i + 2] = dmean.z;566 grads.means[3 * i + 2] = dmean.z;
530 grads.scales[3 * i] = dscale.x;567 grads.scales[3 * i] = dscale.x;
531 grads.scales[3 * i + 1] = dscale.y;568 grads.scales[3 * i + 1] = dscale.y;
532 grads.scales[3 * i + 2] = dscale.z;569 grads.scales[3 * i + 2] = dscale.z;
533 reinterpret_cast<float4*>(grads.quats)[i] = dquat;570 reinterpret_cast<float4*>(grads.quats)[i] = dquat;
534 grads.opacities[i] = dopacity;571 grads.opacities[i] = dopacity;
535 }572 }
536 const uint32_t mask = __ballot_sync(warp::kFullMask, touched);573 const uint32_t mask = __ballot_sync(warp::kFullMask, touched);
537 if (mask != warp::kFullMask) {574 if (mask != warp::kFullMask) {
538 const int warp_i0 =575 const int warp_i0 =
539 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;576 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
540 const int warp_n = max(0, min(32, params.n - warp_i0));577 const int warp_n = max(0, min(32, params.n - warp_i0));
541 uint4* warp_sh_dst =578 uint4* warp_sh_dst =
542 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));579 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
543 const uint4 zero = {0u, 0u, 0u, 0u};580 const uint4 zero = {0u, 0u, 0u, 0u};
544 if (warp_n == 32) {581 if (warp_n == 32) {
545#pragma unroll582#pragma unroll
546 for (int s = 0; s < kShVec16; ++s) {583 for (int s = 0; s < kShVec16; ++s) {
547 const int k = s * 32 + lane;584 const int k = s * 32 + lane;
548 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;585 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
549 if (((mask >> row_in_warp) & 1u) == 0u) {586 if (((mask >> row_in_warp) & 1u) == 0u) {
550 warp_sh_dst[k] = zero;587 warp_sh_dst[k] = zero;
551 }588 }
552 }589 }
553 } else {590 } else {
554 for (int k = lane; k < warp_n * kShVec16; k += 32) {591 for (int k = lane; k < warp_n * kShVec16; k += 32) {
555 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;592 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
556 if (((mask >> row_in_warp) & 1u) == 0u) {593 if (((mask >> row_in_warp) & 1u) == 0u) {
557 warp_sh_dst[k] = zero;594 warp_sh_dst[k] = zero;
558 }595 }
559 }596 }
560 }597 }
561 }598 }
562}599}
563600
564template <typename ShT>601template <typename ShT>
565cudaError_t LaunchPreprocessBackward(const GaussianParams& params,602cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
566 const Frame& frame, const Splat* splats,603 const Frame& frame,
567 const float* splat_grads,604 const float* splat_grads,
605 const uint8_t* splat_touched,
568 const GaussianGrads& grads,606 const GaussianGrads& grads,
569 cudaStream_t stream) {607 cudaStream_t stream) {
570 PreprocessBackwardKernel<ShT>608 PreprocessBackwardKernel<ShT>
571 <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(609 <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(
572 params, frame, splats, splat_grads, grads);610 params, frame, splat_grads, splat_touched, grads);
573 return cudaSuccess;611 return cudaSuccess;
574}612}
575613
614template __global__ void RasterBackwardKernel<1>(
615 const int2* __restrict__, const int32_t* __restrict__,
616 const Splat* __restrict__, Frame, const float* __restrict__,
617 const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
618 uint8_t* __restrict__);
619template __global__ void RasterBackwardKernel<2>(
620 const int2* __restrict__, const int32_t* __restrict__,
621 const Splat* __restrict__, Frame, const float* __restrict__,
622 const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
623 uint8_t* __restrict__);
624template __global__ void RasterBackwardKernel<4>(
625 const int2* __restrict__, const int32_t* __restrict__,
626 const Splat* __restrict__, Frame, const float* __restrict__,
627 const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
628 uint8_t* __restrict__);
576template __global__ void PreprocessBackwardKernel<float>(629template __global__ void PreprocessBackwardKernel<float>(
577 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,630 GaussianParams, Frame, const float* __restrict__,
578 GaussianGrads);631 const uint8_t* __restrict__, GaussianGrads);
579template __global__ void PreprocessBackwardKernel<ShBf16>(632template __global__ void PreprocessBackwardKernel<ShBf16>(
580 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,633 GaussianParams, Frame, const float* __restrict__,
581 GaussianGrads);634 const uint8_t* __restrict__, GaussianGrads);
582template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,635template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
583 const Frame&, const Splat*,636 const Frame&, const float*,
584 const float*,637 const uint8_t*,
585 const GaussianGrads&,638 const GaussianGrads&,
586 cudaStream_t);639 cudaStream_t);
587template cudaError_t LaunchPreprocessBackward<ShBf16>(640template cudaError_t LaunchPreprocessBackward<ShBf16>(
588 const GaussianParams&, const Frame&, const Splat*, const float*,641 const GaussianParams&, const Frame&, const float*, const uint8_t*,
589 const GaussianGrads&, cudaStream_t);642 const GaussianGrads&, cudaStream_t);
590643
591} // namespace gs644} // namespace gs
↕ Expand 195 unchanged lines (lines 1–195) · in ProjectVjp
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
1414 #include "auxiliary.h"
1515
1616 namespace gs {
1717
1818 namespace sh {
1919 // Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly
2020 // to the row `sh_out` (same storage type), group by group.
2121 template <typename ShT>
2222 GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,
2323 ShT* __restrict__ sh_out) {
2424 float b[kShCoeffs];
2525 Basis(d, b);
2626 const float3 raw = EvalRow(sh_in, b);
2727 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
2828 raw.y < 0.0f ? 0.0f : g_rgb.y,
2929 raw.z < 0.0f ? 0.0f : g_rgb.z};
3030 const float x = d.x, y = d.y, z = d.z;
3131 const float xx = x * x, yy = y * y, zz = z * z;
3232 const float xy = x * y, xz = x * z, yz = y * z;
3333 float3 ddir = {0.0f, 0.0f, 0.0f};
3434 #pragma unroll
3535 for (int grp = 0; grp < 4; ++grp) {
3636 const int k0 = 4 * grp;
3737 float v[12];
3838 ShIo<ShT>::LoadGroup(sh_in, grp, v);
3939 const float s0 = v[0] * drgb.x + v[1] * drgb.y + v[2] * drgb.z;
4040 const float s1 = v[3] * drgb.x + v[4] * drgb.y + v[5] * drgb.z;
4141 const float s2 = v[6] * drgb.x + v[7] * drgb.y + v[8] * drgb.z;
4242 const float s3 = v[9] * drgb.x + v[10] * drgb.y + v[11] * drgb.z;
4343 float out[12];
4444 #pragma unroll
4545 for (int q = 0; q < 4; ++q) {
4646 out[3 * q] = b[k0 + q] * drgb.x;
4747 out[3 * q + 1] = b[k0 + q] * drgb.y;
4848 out[3 * q + 2] = b[k0 + q] * drgb.z;
4949 }
5050 ShIo<ShT>::StoreGroup(sh_out, grp, out);
5151 if (grp == 0) {
5252 ddir.y += s1 * (-kC1);
5353 ddir.z += s2 * kC1;
5454 ddir.x += s3 * (-kC1);
5555 } else if (grp == 1) {
5656 ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};
5757 ddir += (s1 * kC2[1]) * float3{0.0f, z, y};
5858 ddir += (s2 * kC2[2]) * float3{-2.0f * x, -2.0f * y, 4.0f * z};
5959 ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};
6060 } else if (grp == 2) {
6161 ddir += (s0 * kC2[4]) * float3{2.0f * x, -2.0f * y, 0.0f};
6262 ddir += (s1 * kC3[0]) * float3{6.0f * xy, 3.0f * (xx - yy), 0.0f};
6363 ddir += (s2 * kC3[1]) * float3{yz, xz, xy};
6464 ddir += (s3 * kC3[2]) *
6565 float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};
6666 } else {
6767 ddir += (s0 * kC3[3]) *
6868 float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};
6969 ddir += (s1 * kC3[4]) *
7070 float3{4.0f * zz - 3.0f * xx - yy, -2.0f * xy, 8.0f * xz};
7171 ddir += (s2 * kC3[5]) * float3{2.0f * xz, -2.0f * yz, xx - yy};
7272 ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};
7373 }
7474 }
7575 return ddir;
7676 }
7777 } // namespace sh
7878
7979 // VJP of the projection for a touched Gaussian `i`; reads its SH row from
8080 // `sh_in` and writes dL/dsh directly to `sh_out`.
8181 // Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
8282 template <typename ShT>
8383 GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
8484 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
8585 const ProjectionGrad& g, float3& dmean_out,
8686 float3& dscale_out, float4& dquat_out) {
8787 const float3 mean = LoadMean(p, i);
8888 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
8989 p.scales[3 * i + 2]};
9090 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
9191
9292 // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
9393 float3 dmean = {0.0f, 0.0f, 0.0f};
9494 {
9595 float inv_len;
9696 const float3 dir = ViewDir(cam, mean, inv_len);
9797 const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
9898 // d(v/|v|)/dv = (I - dir dir^T) / |v|.
9999 dmean += inv_len * (ddir - dot(dir, ddir) * dir);
100100 }
101101
102102 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
103103 const float3 t = cam.ToView(mean);
104104 const float tz = t.z;
105105 const float limx = kFovClamp * cam.TanHalfFovX();
106106 const float limy = kFovClamp * cam.TanHalfFovY();
107107 const float txtz = t.x / tz, tytz = t.y / tz;
108108 const bool x_clamped = txtz < -limx || txtz > limx;
109109 const bool y_clamped = tytz < -limy || tytz > limy;
110110 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
111111 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
112112 const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
113113 const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
114114 const Mat3& w = cam.rotation;
115115 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
116116 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
117117
118118 const Mat3 r = QuatToRotation(quat);
119119 const Sym3 sigma = Covariance3D(r, scale);
120120 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
121121 const float cov_a = dot(m0, s0) + kLowPass;
122122 const float cov_b = dot(m0, s1);
123123 const float cov_c = dot(m1, s1) + kLowPass;
124124
125125 // --- Conic -> 2D covariance: G_cov = -K G_K K.
126126 const float det = cov_a * cov_c - cov_b * cov_b;
127127 const float inv = 1.0f / det;
128128 const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
129129 const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
130130 // P = G_K K.
131131 const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
132132 const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
133133 // G_cov = -K P (symmetric).
134134 const float ca = -(ka * p00 + kb * p10);
135135 const float cb = -(ka * p01 + kb * p11);
136136 const float cc = -(kb * p01 + kc * p11);
137137
138138 // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
139139 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
140140 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
141141 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
142142 Sym3 dsig;
143143 dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
144144 dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
145145 dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
146146 dsig.xy =
147147 ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
148148 dsig.xz =
149149 ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
150150 dsig.yz =
151151 ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
152152
153153 // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
154154 const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
155155 const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
156156 const float tz2 = tz * tz, tz3 = tz2 * tz;
157157 float3 dt;
158158 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
159159 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
160160 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
161161 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
162162 // --- Screen-space mean u = fx x / z + cx.
163163 dt.x += cam.fx / tz * g.xy.x;
164164 dt.y += cam.fy / tz * g.xy.y;
165165 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
166166 dmean += MulT(w, dt);
167167 dmean_out = dmean;
168168
169169 // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
170170 const float sv[3] = {scale.x, scale.y, scale.z};
171171 float ds[3];
172172 Mat3 dr;
173173 #pragma unroll
174174 for (int j = 0; j < 3; ++j) {
175175 const float3 r_col = r.col(j);
176176 const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
177177 ds[j] = dot(da_col, r_col);
178178 dr.m[0][j] = da_col.x * sv[j];
179179 dr.m[1][j] = da_col.y * sv[j];
180180 dr.m[2][j] = da_col.z * sv[j];
181181 }
182182 dscale_out = {ds[0], ds[1], ds[2]};
183183
184184 // --- R(q) -> q.
185185 const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
186186 const auto& G = dr.m;
187187 dquat_out = {
188188 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
189189 qy * G[2][0] + qx * G[2][1]),
190190 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
191191 qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
192192 2.0f *
193193 (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
194194 qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
195195 2.0f *
196196 (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
197197 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
198198 }
199199
200−// Reverse traversal per pixel; accumulates splat gradients via Half-Warp
200+// Reverse traversal per pixel; accumulates splat gradients with atomics.
201−// Two-Splat Folding and Halving Register Reduce-Scatter.
201+template <int kRows>
202−__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 4)
202+__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 4)
203203 RasterBackwardKernel(const int2* __restrict__ ranges,
204204 const int32_t* __restrict__ point_list,
205205 const Splat* __restrict__ splats, Frame frame,
206206 const float* __restrict__ final_t,
207207 const int32_t* __restrict__ n_contrib,
208208 const float* __restrict__ dl_dimage,
209− float* __restrict__ splat_grads) {
209+ float* __restrict__ splat_grads,
210+ uint8_t* __restrict__ splat_touched) {
210− using L = RasterLayout<1>;
211+ using L = RasterLayout<kRows>;
212+ constexpr int kRasterRows = kRows;
211213 constexpr int kRasterThreads = L::kThreads;
212214 constexpr int kWarps = L::kWarps;
213215 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
214216 const int tid = threadIdx.x;
215217 const int warp = tid / 32;
216218 const int lane = tid % 32;
217− const RasterPixels px = RasterPixels::For<1>(tid);
219+ const RasterPixels px = RasterPixels::For<kRows>(tid);
218220 const float fx = static_cast<float>(px.x);
219− const float fy = static_cast<float>(px.y0);
220221 const int2 range = ranges[tile];
221222
222− const bool inside = px.x < frame.width && px.y0 < frame.height;
223+ struct PixelState {
224+ float t;
225+ float3 dpix;
226+ float accum; // dot(color accumulated behind, dpix)
227+ int last;
228+ };
229+ PixelState st[kRasterRows];
223230 const int pixels = frame.width * frame.height;
231+ int thread_last = 0;
232+#pragma unroll
233+ for (int k = 0; k < kRasterRows; ++k) {
234+ const int y = px.y0 + k;
235+ const bool inside = px.x < frame.width && y < frame.height;
224− const int p = px.y0 * frame.width + px.x;
236+ const int p = y * frame.width + px.x;
225− float t = inside ? final_t[p] : 0.0f;
237+ st[k].t = inside ? final_t[p] : 0.0f;
226− const float3 dpix =
238+ st[k].dpix =
227239⇥ inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
228240⇥ dl_dimage[2 * pixels + p]}
229241⇥ : float3{0.0f, 0.0f, 0.0f};
230− float accum = 0.0f; // dot(color accumulated behind, dpix)
242+ st[k].accum = 0.0f;
231− const int last = inside ? n_contrib[p] : 0;
243+ st[k].last = inside ? n_contrib[p] : 0;
244+ thread_last = max(thread_last, st[k].last);
245+ }
232246
233− int warp_max_last = last;
247+ int warp_max_last = thread_last;
234248 #pragma unroll
235249 for (int offset = 16; offset > 0; offset /= 2) {
236250 warp_max_last = max(
237251 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
↕ Expand 15 unchanged lines (lines 252–266) · in PixelState
238252 }
239253 __shared__ int shared_max_last[kWarps];
240254 if (lane == 0) shared_max_last[warp] = warp_max_last;
241255 __syncthreads();
242256 int block_max_last = 0;
243257 #pragma unroll
244258 for (int w = 0; w < kWarps; ++w) {
245259 block_max_last = max(block_max_last, shared_max_last[w]);
246260 }
247261 if (block_max_last <= 0) return;
248262
249263 __shared__ float4 batch_geom0[kRasterThreads];
250264 __shared__ float2 batch_geom1[kRasterThreads];
251265 __shared__ float4 batch_rgb_o[kRasterThreads];
252266 __shared__ uint32_t batch_id[kRasterThreads];
253267 __shared__ uint8_t batch_mask[kRasterThreads];
254268 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
255269 int scatter = -1;
256270 int scatter2 = -1;
257− {
271+ if constexpr (L::kWarpWidth == 8) {
258272 const int sub_lane = lane & 15;
259273 const int row = sub_lane >> 3, col = sub_lane & 7;
260274 if (col == 0) {
261275 scatter2 = row;
↕ Expand 13 unchanged lines (lines 276–288) · in PixelState
262276 } else if (col == 2) {
263277 scatter2 = 2 + row;
264278 } else if (col == 4) {
265279 scatter2 = row == 0 ? 4 : 7;
266280 } else if (col == 6) {
267281 scatter2 = row == 0 ? 6 : 8;
268282 } else if (col == 1 && row == 0) {
269283 scatter2 = 5;
270284 }
271285 const int full_row = lane >> 3, full_col = lane & 7;
272286 if (full_col == 0) {
273287 scatter = full_row;
274288 } else if (full_col == 2) {
275289 scatter = full_row == 0 ? 4 : 5 + full_row;
276290 } else if (full_col == 4 && full_row == 0) {
277291 scatter = 5;
278292 }
293+ } else {
294+ scatter = warp::ScatterIndex(kSplatGradFloats, lane);
279295 }
280296
281− auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,
297+ auto eval_splat = [&](int j, float4 g0, float2 g1, float4 ro, int entry,
282− float& ga_sum, float& ga_dy, float& ga_dy2,
298+ float& dx, float& ga_sum, float& ga_dy, float& ga_dy2,
283− float3& drgb) -> bool {
299+ float3& drgb) {
284300 dx = g0.x - fx;
285301 ga_sum = 0.0f;
286302 ga_dy = 0.0f;
287303 ga_dy2 = 0.0f;
288304 drgb = {0.0f, 0.0f, 0.0f};
289− if (entry < last) {
305+ if (entry < thread_last) {
306+ const float cx_dx2 = g0.z * dx * dx;
290− const float dy = g0.y - fy;
307+ const float cy_dx = g0.w * dx;
308+#pragma unroll
309+ for (int k = 0; k < kRasterRows; ++k) {
310+ PixelState& ps = st[k];
311+ if (entry >= ps.last) continue;
312+ const float dy = g0.y - static_cast<float>(px.y0 + k);
291− const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
313+ const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
292− if (power <= 0.0f && power >= g1.y) {
314+ if (power > 0.0f || power < g1.y) continue;
293315 const float g_val = Exp2Approx(power);
294316 const float alpha = fminf(kMaxAlpha, ro.w * g_val);
295− if (alpha >= kMinAlpha) {
317+ if (alpha < kMinAlpha) continue;
296− t = __fdividef(t, 1.0f - alpha);
318+ ps.t = __fdividef(ps.t, 1.0f - alpha);
297− const float w = alpha * t;
319+ const float w = alpha * ps.t;
320+ const float color_dot =
298− const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
321+ ro.x * ps.dpix.x + ro.y * ps.dpix.y + ro.z * ps.dpix.z;
299− const float diff = color_dot - accum;
322+ const float diff = color_dot - ps.accum;
300− const float dl_dalpha = t * diff;
323+ const float dl_dalpha = ps.t * diff;
301− accum += alpha * diff;
324+ ps.accum += alpha * diff;
302325⇥ const float ga = g_val * dl_dalpha;
303− ga_sum = ga;
326+ ga_sum += ga;
304327⇥ const float ga_y = ga * dy;
305− ga_dy = ga_y;
328+ ga_dy += ga_y;
306− ga_dy2 = ga_y * dy;
329+ ga_dy2 += ga_y * dy;
307− drgb = w * dpix;
330+ drgb += w * ps.dpix;
308− return true;
309331⇥ }
310332⇥ }
311− }
312− return false;
313333 };
314334 auto raw_conic = [](float4 g0, float2 g1) {
315335 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
316336 kInvPower2Scale * g1.x};
317337 };
318338
319339 for (int end = block_max_last; end > 0; end -= kRasterThreads) {
320340 const int e = end - 1 - tid;
321341 if (e >= 0) {
322− const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
342+ const uint32_t packed = static_cast<uint32_t>(point_list[range.x + e]);
343+ const uint32_t id = packed & 0x00ffffffu;
344+ splat_touched[id] = 1u;
323345 const float4 p0 = __ldg(&splats[id].p0);
324346 const float4 p1 = __ldg(&splats[id].p1);
325347 const float4 p2 = __ldg(&splats[id].p2);
326348 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
327349 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
328350 batch_geom1[tid] = {c2.z, p2.w};
329351 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
330352 batch_id[tid] = id;
331− batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
353+ batch_mask[tid] = static_cast<uint8_t>(packed >> 24);
332354 } else {
333355 batch_mask[tid] = 0u;
334356 }
335357 __syncthreads();
336358 const int count = min(kRasterThreads, end);
337359 #pragma unroll
338360 for (int c = 0; c < kRasterThreads; c += 32) {
339361 if (c < count && end - 32 - c < warp_max_last) {
340362 uint32_t mask = __ballot_sync(
341363 warp::kFullMask,
342364 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
365+ if constexpr (L::kWarpWidth == 8) {
343366⇥ while ((mask & (mask - 1u)) != 0u) {
344367⇥ const int bit0 = __ffs(mask) - 1;
345368⇥ mask &= mask - 1u;
346369⇥ const int bit1 = __ffs(mask) - 1;
347370⇥ mask &= mask - 1u;
348371⇥ const int j0 = c + bit0, j1 = c + bit1;
349372⇥ const float4 g0_0 = batch_geom0[j0];
350373⇥ const float2 g1_0 = batch_geom1[j0];
351374⇥ const float4 ro_0 = batch_rgb_o[j0];
352375⇥ const float4 g0_1 = batch_geom0[j1];
353376⇥ const float2 g1_1 = batch_geom1[j1];
354377⇥ const float4 ro_1 = batch_rgb_o[j1];
355378
356379⇥ float dx0, ga_sum0, ga_dy0, ga_dy2_0;
357380⇥ float3 drgb0;
358− const bool hit0 = eval_splat(g0_0, g1_0, ro_0, end - 1 - j0, dx0,
381+ eval_splat(j0, g0_0, g1_0, ro_0, end - 1 - j0, dx0, ga_sum0, ga_dy0,
359− ga_sum0, ga_dy0, ga_dy2_0, drgb0);
382+ ga_dy2_0, drgb0);
360383
361384⇥ float dx1, ga_sum1, ga_dy1, ga_dy2_1;
362385⇥ float3 drgb1;
363− const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,
386+ eval_splat(j1, g0_1, g1_1, ro_1, end - 1 - j1, dx1, ga_sum1, ga_dy1,
364− ga_sum1, ga_dy1, ga_dy2_1, drgb1);
387+ ga_dy2_1, drgb1);
365− // As in s01: a pair that no pixel of the warp blends has zero
366− // gradients, so skip its reduction and atomics entirely.
367− if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;
368388
369389⇥ const bool hi16 = (lane & 16) != 0;
370390⇥ const float opacity = hi16 ? ro_1.w : ro_0.w;
371391⇥ const float dx = hi16 ? dx1 : dx0;
372392⇥ const float3 cn =
373393⇥ hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);
374394
375395⇥ auto fold16 = [&](float v0, float v1) {
376396⇥ const float send = hi16 ? v0 : v1;
377397⇥ const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);
378398⇥ return (hi16 ? v1 : v0) + recv;
379399⇥ };
380400⇥ float ga_sum = fold16(ga_sum0, ga_sum1);
381401⇥ float ga_dy = fold16(ga_dy0, ga_dy1);
382402⇥ float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),
383403⇥ fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),
384404⇥ fold16(drgb0.z, drgb1.z)};
385405⇥ warp::internal::HalvingStep<4, 8>(b4, lane);
386406⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
387407⇥ ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
388408
389409⇥ const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
390410⇥ const float p_dx = dx * p_sum;
391411⇥ const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
392412⇥ : -(cn.y * p_dx + cn.z * p_dy);
413+ const float v_dx1 =
393− const float v_dx1 = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
414+ (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
394415⇥ float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};
395416⇥ warp::internal::HalvingStep<4, 4>(v4, lane);
396417⇥ warp::internal::HalvingStep<2, 2>(v4, lane);
397418⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);
398419⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);
399420⇥ const float v_sum =
400421⇥ v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);
401422⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);
402423⇥ const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
403424⇥ if (scatter2 >= 0 && out_grad != 0.0f) {
404425⇥ const uint32_t target_id = batch_id[hi16 ? j1 : j0];
405426⇥ atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,
406427⇥ out_grad);
407428⇥ }
408429⇥ }
430+ }
409431 while (mask != 0u) {
410432 const int bit = __ffs(mask) - 1;
411433 mask &= mask - 1u;
412434 const int j = c + bit;
413435 const float4 g0 = batch_geom0[j];
414436 const float2 g1 = batch_geom1[j];
415437 const float4 ro = batch_rgb_o[j];
416438 float dx, ga_sum, ga_dy, ga_dy2;
417439 float3 drgb;
418− const bool hit =
419− eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
440+ eval_splat(j, g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2,
420− if (!__any_sync(warp::kFullMask, hit)) continue;
441+ drgb);
421442 const float opacity = ro.w;
422443 const float3 cn = raw_conic(g0, g1);
444+ float out_grad = 0.0f;
445+ if constexpr (L::kWarpWidth == 8) {
423446⇥ float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};
424447⇥ warp::internal::HalvingStep<4, 16>(b4, lane);
425448⇥ warp::internal::HalvingStep<2, 8>(b4, lane);
426449⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);
427450⇥ ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);
428451⇥ ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
429452⇥ ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
430453⇥ const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
431454⇥ const float p_dx = dx * p_sum;
432455⇥ float v_dx;
433456⇥ if ((lane & 16) == 0) {
434457⇥ v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
435458⇥ : -(cn.y * p_dx + cn.z * p_dy);
436459⇥ } else {
437460⇥ v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
438461⇥ }
439462⇥ float v[3] = {v_dx, b4[0], ga_sum};
440463⇥ warp::internal::HalvingStep<3, 4>(v, lane);
441464⇥ warp::internal::HalvingStep<2, 2>(v, lane);
442− const float out_grad =
443− v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
465+ out_grad = v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
466+ } else {
467+ const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
468+ const float p_dx = dx * p_sum;
469+ float gr[kSplatGradFloats] = {-(cn.x * p_dx + cn.y * p_dy), // xy
470+ -(cn.y * p_dx + cn.z * p_dy),
471+ -0.5f * dx * p_dx, // conic
472+ -dx * p_dy,
473+ -0.5f * opacity * ga_dy2,
474+ ga_sum, // opacity
475+ drgb.x,
476+ drgb.y,
477+ drgb.z}; // rgb
478+ warp::ReduceScatter(gr, lane);
479+ out_grad = gr[0];
480+ }
444481 if (scatter >= 0 && out_grad != 0.0f) {
445482 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
446483 out_grad);
447484 }
448485 }
449486 }
450487 }
451488 __syncthreads();
452489 }
453490 }
454491
455492 // ---------------------------------------------------------------------------
456493 // Preprocess backward: splat gradients -> parameter gradients (dense).
457−// Gaussians stored as culled (radius 0, i.e. no tile pairs) are untouched.
494+// Uses the 1-byte `splat_touched` flag written by `RasterBackwardKernel` to
495+// skip reading `splat_grads` and `params.sh` for untouched Gaussians, and
458−// Uses warp-coalesced 16-byte stores (0 shared memory) to zero untouched SH
496+// zeros untouched SH rows with warp-coalesced float4 stores (0 shared memory).
459−// rows and streams active SH rows in 4 groups of 12 floats via `ShIo<ShT>`.
460497
461498 template <typename ShT>
462499 __global__ void __launch_bounds__(kGaussianThreads)
463500 PreprocessBackwardKernel(GaussianParams params, Frame frame,
464− const Splat* __restrict__ splats,
465501 const float* __restrict__ splat_grads,
502+ const uint8_t* __restrict__ splat_touched,
466503 GaussianGrads grads) {
467504 constexpr int kShVec16 = ShIo<ShT>::kRowVec16; // 16-byte stores per row.
468505 const int i = blockIdx.x * blockDim.x + threadIdx.x;
469506 const int lane = threadIdx.x % 32;
470− bool touched = (i < params.n && splats[i].radius() != 0);
507+ bool touched = (i < params.n && splat_touched[i] != 0u);
471508 const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);
472509 if (init_mask == 0u) {
473510 if (i < params.n) {
474511 grads.means[3 * i] = 0.0f;
↕ Expand 87 unchanged lines (lines 512–598) · in PreprocessBackwardKernel
475512 grads.means[3 * i + 1] = 0.0f;
476513 grads.means[3 * i + 2] = 0.0f;
477514 grads.scales[3 * i] = 0.0f;
478515 grads.scales[3 * i + 1] = 0.0f;
479516 grads.scales[3 * i + 2] = 0.0f;
480517 reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};
481518 grads.opacities[i] = 0.0f;
482519 }
483520 const int warp_i0 =
484521 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
485522 const int warp_n = max(0, min(32, params.n - warp_i0));
486523 uint4* warp_sh_dst =
487524 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
488525 const uint4 zero = {0u, 0u, 0u, 0u};
489526 if (warp_n == 32) {
490527 #pragma unroll
491528 for (int s = 0; s < kShVec16; ++s) {
492529 warp_sh_dst[s * 32 + lane] = zero;
493530 }
494531 } else {
495532 for (int k = lane; k < warp_n * kShVec16; k += 32) {
496533 warp_sh_dst[k] = zero;
497534 }
498535 }
499536 return;
500537 }
501538
502539 float3 dmean = {0.0f, 0.0f, 0.0f};
503540 float3 dscale = {0.0f, 0.0f, 0.0f};
504541 float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};
505542 float dopacity = 0.0f;
506543 if (touched) {
507544 float g[kSplatGradFloats];
508545 bool nonzero = false;
509546 #pragma unroll
510547 for (int q = 0; q < kSplatGradFloats; ++q) {
511548 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
512549 nonzero = nonzero || g[q] != 0.0f;
513550 }
514551 touched = nonzero;
515552 if (touched) {
516553 ProjectionGrad pg;
517554 pg.xy = {g[0], g[1]};
518555 pg.conic = {g[2], g[3], g[4]};
519556 pg.opacity = g[5];
520557 pg.rgb = {g[6], g[7], g[8]};
521558 dopacity = g[5];
522559 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
523560 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
524561 }
525562 }
526563 if (i < params.n) {
527564 grads.means[3 * i] = dmean.x;
528565 grads.means[3 * i + 1] = dmean.y;
529566 grads.means[3 * i + 2] = dmean.z;
530567 grads.scales[3 * i] = dscale.x;
531568 grads.scales[3 * i + 1] = dscale.y;
532569 grads.scales[3 * i + 2] = dscale.z;
533570 reinterpret_cast<float4*>(grads.quats)[i] = dquat;
534571 grads.opacities[i] = dopacity;
535572 }
536573 const uint32_t mask = __ballot_sync(warp::kFullMask, touched);
537574 if (mask != warp::kFullMask) {
538575 const int warp_i0 =
539576 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
540577 const int warp_n = max(0, min(32, params.n - warp_i0));
541578 uint4* warp_sh_dst =
542579 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
543580 const uint4 zero = {0u, 0u, 0u, 0u};
544581 if (warp_n == 32) {
545582 #pragma unroll
546583 for (int s = 0; s < kShVec16; ++s) {
547584 const int k = s * 32 + lane;
548585 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
549586 if (((mask >> row_in_warp) & 1u) == 0u) {
550587 warp_sh_dst[k] = zero;
551588 }
552589 }
553590 } else {
554591 for (int k = lane; k < warp_n * kShVec16; k += 32) {
555592 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
556593 if (((mask >> row_in_warp) & 1u) == 0u) {
557594 warp_sh_dst[k] = zero;
558595 }
559596 }
560597 }
561598 }
562599 }
563600
564601 template <typename ShT>
565602 cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
566− const Frame& frame, const Splat* splats,
603+ const Frame& frame,
567604 const float* splat_grads,
605+ const uint8_t* splat_touched,
568606 const GaussianGrads& grads,
569607 cudaStream_t stream) {
570608 PreprocessBackwardKernel<ShT>
571609 <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(
572− params, frame, splats, splat_grads, grads);
610+ params, frame, splat_grads, splat_touched, grads);
573611 return cudaSuccess;
574612 }
575613
614+template __global__ void RasterBackwardKernel<1>(
615+ const int2* __restrict__, const int32_t* __restrict__,
616+ const Splat* __restrict__, Frame, const float* __restrict__,
617+ const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
618+ uint8_t* __restrict__);
619+template __global__ void RasterBackwardKernel<2>(
620+ const int2* __restrict__, const int32_t* __restrict__,
621+ const Splat* __restrict__, Frame, const float* __restrict__,
622+ const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
623+ uint8_t* __restrict__);
624+template __global__ void RasterBackwardKernel<4>(
625+ const int2* __restrict__, const int32_t* __restrict__,
626+ const Splat* __restrict__, Frame, const float* __restrict__,
627+ const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
628+ uint8_t* __restrict__);
576629 template __global__ void PreprocessBackwardKernel<float>(
577− GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
630+ GaussianParams, Frame, const float* __restrict__,
578− GaussianGrads);
631+ const uint8_t* __restrict__, GaussianGrads);
579632 template __global__ void PreprocessBackwardKernel<ShBf16>(
580− GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
633+ GaussianParams, Frame, const float* __restrict__,
581− GaussianGrads);
634+ const uint8_t* __restrict__, GaussianGrads);
582635 template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
583− const Frame&, const Splat*,
636+ const Frame&, const float*,
584− const float*,
637+ const uint8_t*,
585638 const GaussianGrads&,
586639 cudaStream_t);
587640 template cudaError_t LaunchPreprocessBackward<ShBf16>(
588− const GaussianParams&, const Frame&, const Splat*, const float*,
641+ const GaussianParams&, const Frame&, const float*, const uint8_t*,
589642 const GaussianGrads&, cudaStream_t);
590643
591644 } // namespace gs
csrc/rasterizer_impl.h +3 −0
s03_sh_streaming/csrc/rasterizer_impl.h
s04_raster_pruning/csrc/rasterizer_impl.h
↕ Expand 16 unchanged lines (lines 1–16)
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 <cstddef>
15#include <cstdint>15#include <cstdint>
16#include <iterator>16#include <iterator>
1717
18#include "auxiliary.h"18#include "auxiliary.h"
1919
20namespace gs {20namespace gs {
21
22inline constexpr int kDefaultForwardRows = 1;
23inline constexpr int kDefaultBackwardRows = 1;
2124
22// Number of bits needed to represent tile ids in [0, num_tiles).25// Number of bits needed to represent tile ids in [0, num_tiles).
23inline int TileBits(int num_tiles) {26inline int TileBits(int num_tiles) {
24 int bits = 1;27 int bits = 1;
↕ Expand 41 unchanged lines (lines 28–68) · in TileRangesKernel
25 while ((1 << bits) < num_tiles) ++bits;28 while ((1 << bits) < num_tiles) ++bits;
26 return bits;29 return bits;
27}30}
2831
29struct GatherCountsIter {32struct GatherCountsIter {
30 const int32_t* counts;33 const int32_t* counts;
31 const int32_t* order;34 const int32_t* order;
32 __host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {35 __host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
33 return counts[order[idx]];36 return counts[order[idx]];
34 }37 }
35 __host__ __device__ __forceinline__ GatherCountsIter38 __host__ __device__ __forceinline__ GatherCountsIter
36 operator+(ptrdiff_t offset) const {39 operator+(ptrdiff_t offset) const {
37 return {counts, order + offset};40 return {counts, order + offset};
38 }41 }
39 __host__ __device__ __forceinline__ int32_t operator*() const {42 __host__ __device__ __forceinline__ int32_t operator*() const {
40 return counts[*order];43 return counts[*order];
41 }44 }
42};45};
4346
44template <Culling kCulling>47template <Culling kCulling>
45__global__ void EmitPairsKernel(const Splat* __restrict__ splats,48__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
46 const int32_t* __restrict__ sorted_ids,49 const int32_t* __restrict__ sorted_ids,
47 const int32_t* __restrict__ offsets, int n_vis,50 const int32_t* __restrict__ offsets, int n_vis,
48 Frame frame, uint16_t* __restrict__ keys,51 Frame frame, uint16_t* __restrict__ keys,
49 int32_t* __restrict__ values);52 int32_t* __restrict__ values);
5053
51__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,54__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
52 int num_pairs, int2* __restrict__ ranges);55 int num_pairs, int2* __restrict__ ranges);
5356
54} // namespace gs57} // namespace gs
5558
56namespace std {59namespace std {
57template <>60template <>
58struct iterator_traits<gs::GatherCountsIter> {61struct iterator_traits<gs::GatherCountsIter> {
59 using difference_type = ptrdiff_t;62 using difference_type = ptrdiff_t;
60 using value_type = int32_t;63 using value_type = int32_t;
61 using pointer = const int32_t*;64 using pointer = const int32_t*;
62 using reference = int32_t;65 using reference = int32_t;
63 using iterator_category = random_access_iterator_tag;66 using iterator_category = random_access_iterator_tag;
64};67};
65} // namespace std68} // namespace std
↕ Expand 16 unchanged lines (lines 1–16)
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
1414 #include <cstddef>
1515 #include <cstdint>
1616 #include <iterator>
1717
1818 #include "auxiliary.h"
1919
2020 namespace gs {
21+
22+inline constexpr int kDefaultForwardRows = 1;
23+inline constexpr int kDefaultBackwardRows = 1;
2124
2225 // Number of bits needed to represent tile ids in [0, num_tiles).
2326 inline int TileBits(int num_tiles) {
2427 int bits = 1;
↕ Expand 41 unchanged lines (lines 28–68) · in TileRangesKernel
2528 while ((1 << bits) < num_tiles) ++bits;
2629 return bits;
2730 }
2831
2932 struct GatherCountsIter {
3033 const int32_t* counts;
3134 const int32_t* order;
3235 __host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
3336 return counts[order[idx]];
3437 }
3538 __host__ __device__ __forceinline__ GatherCountsIter
3639 operator+(ptrdiff_t offset) const {
3740 return {counts, order + offset};
3841 }
3942 __host__ __device__ __forceinline__ int32_t operator*() const {
4043 return counts[*order];
4144 }
4245 };
4346
4447 template <Culling kCulling>
4548 __global__ void EmitPairsKernel(const Splat* __restrict__ splats,
4649 const int32_t* __restrict__ sorted_ids,
4750 const int32_t* __restrict__ offsets, int n_vis,
4851 Frame frame, uint16_t* __restrict__ keys,
4952 int32_t* __restrict__ values);
5053
5154 __global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
5255 int num_pairs, int2* __restrict__ ranges);
5356
5457 } // namespace gs
5558
5659 namespace std {
5760 template <>
5861 struct iterator_traits<gs::GatherCountsIter> {
5962 using difference_type = ptrdiff_t;
6063 using value_type = int32_t;
6164 using pointer = const int32_t*;
6265 using reference = int32_t;
6366 using iterator_category = random_access_iterator_tag;
6467 };
6568 } // namespace std
Unchanged files in s04_raster_pruning (identical to s03_sh_streaming): csrc/config.h, csrc/rasterizer_impl.cu.cc