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

s02_reduce_scatter → s03_sh_streaming

Step 3: 4-Group Register-Window SH Streaming & Cooperative Warp Zeroing  ·  +157 / −129 lines across 2 changed files

Core files: csrc/auxiliary.h +34/−84csrc/backward.cu.cc +123/−45
csrc/auxiliary.h +34 −84
s02_reduce_scatter/csrc/auxiliary.h
s03_sh_streaming/csrc/auxiliary.h
↕ Expand 420 unchanged lines (lines 1–420) · in Projection
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 coefficients
425// and 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.
426427
427using ShBf16 = uint16_t; // bfloat16 bits.428using ShBf16 = uint16_t; // bfloat16 bits.
428429
429template <typename T>430template <typename T>
430struct ShIo;431struct ShIo;
431432
432template <>433template <>
433struct ShIo<float> {434struct ShIo<float> {
435 static constexpr int kRowVec16 = 12; // 16-byte vectors per row.
434 GS_D static void LoadRow(const float* row, float3 coeffs[kShCoeffs]) {436 GS_D static void LoadGroup(const float* row, int grp, float v[12]) {
435 const float4* p = reinterpret_cast<const float4*>(row);437 const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;
436 float raw[3 * kShCoeffs];
437#pragma unroll
438 for (int k = 0; k < 12; ++k) {
439 const float4 q = __ldg(p + k);438 const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);
440 raw[4 * k] = q.x;
441 raw[4 * k + 1] = q.y;
442 raw[4 * k + 2] = q.z;
443 raw[4 * k + 3] = q.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;
441 v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;
444 }442 }
445#pragma unroll
446 for (int c = 0; c < kShCoeffs; ++c) {
447 coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
448 }
449 }
450 GS_D static void StoreRow(float* row, const float3 dsh[kShCoeffs]) {443 GS_D static void StoreGroup(float* row, int grp, const float v[12]) {
451 float raw[3 * kShCoeffs];
452#pragma unroll
453 for (int c = 0; c < kShCoeffs; ++c) {
454 raw[3 * c] = dsh[c].x;
455 raw[3 * c + 1] = dsh[c].y;
456 raw[3 * c + 2] = dsh[c].z;
457 }
458 float4* p = reinterpret_cast<float4*>(row);444 float4* p = reinterpret_cast<float4*>(row) + 3 * grp;
459#pragma unroll
460 for (int k = 0; k < 12; ++k) {
461 p[k] = {raw[4 * k], raw[4 * k + 1], raw[4 * k + 2], raw[4 * k + 3]};445 p[0] = {v[0], v[1], v[2], v[3]};
462 }
446 p[1] = {v[4], v[5], v[6], v[7]};
447 p[2] = {v[8], v[9], v[10], v[11]};
463 }448 }
464};449};
465450
466template <>451template <>
467struct ShIo<ShBf16> {452struct ShIo<ShBf16> {
453 static constexpr int kRowVec16 = 6;
468 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); }
469 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); }
470 // Round to nearest even (inputs are finite).456 // Round to nearest even (inputs are finite).
471 GS_D static uint32_t Round(float f) {457 GS_D static uint32_t Round(float f) {
472 const uint32_t u = __float_as_uint(f);458 const uint32_t u = __float_as_uint(f);
473 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;459 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
474 }460 }
475 GS_D static uint32_t Pack(float lo, float hi) {461 GS_D static uint32_t Pack(float lo, float hi) {
476 return Round(lo) | (Round(hi) << 16);462 return Round(lo) | (Round(hi) << 16);
477 }463 }
478 GS_D static void LoadRow(const ShBf16* row, float3 coeffs[kShCoeffs]) {464 GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {
465 // 24-byte groups: three 8-byte loads.
479 const uint2* p = reinterpret_cast<const uint2*>(row);466 const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;
480 float raw[3 * kShCoeffs];
481#pragma unroll467#pragma unroll
482 for (int k = 0; k < 12; ++k) {468 for (int k = 0; k < 3; ++k) {
483 const uint2 w = __ldg(p + k);469 const uint2 w = __ldg(p + k);
484 raw[4 * k] = Lo(w.x);
485 raw[4 * k + 1] = Hi(w.x);470 v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);
486 raw[4 * k + 2] = Lo(w.y);
487 raw[4 * k + 3] = Hi(w.y);471 v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);
488 }
489#pragma unroll
490 for (int c = 0; c < kShCoeffs; ++c) {
491 coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
492 }472 }
493 }473 }
494 GS_D static void StoreRow(ShBf16* row, const float3 dsh[kShCoeffs]) {474 GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {
495 float raw[3 * kShCoeffs];
475 uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;
496#pragma unroll476#pragma unroll
497 for (int c = 0; c < kShCoeffs; ++c) {
498 raw[3 * c] = dsh[c].x;
499 raw[3 * c + 1] = dsh[c].y;
500 raw[3 * c + 2] = dsh[c].z;
501 }
502 uint2* p = reinterpret_cast<uint2*>(row);
503#pragma unroll
504 for (int k = 0; k < 12; ++k) {477 for (int k = 0; k < 3; ++k) {
505 p[k] = {Pack(raw[4 * k], raw[4 * k + 1]),478 p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};
506 Pack(raw[4 * k + 2], raw[4 * k + 3])};
507 }479 }
508 }480 }
509};481};
510482
↕ Expand 41 unchanged lines (lines 483–523) · in Basis
511template <typename ShT>483template <typename ShT>
512GS_D const ShT* ShRow(const void* sh, int i) {484GS_D const ShT* ShRow(const void* sh, int i) {
513 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;
514}486}
515template <typename ShT>487template <typename ShT>
516GS_D ShT* ShRow(void* sh, int i) {488GS_D ShT* ShRow(void* sh, int i) {
517 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;
518}490}
519491
520// ----------------------------------------------------------------------------492// ----------------------------------------------------------------------------
521// Spherical harmonics (degree 3, Inria sign convention).493// Spherical harmonics (degree 3, Inria sign convention).
522494
523namespace sh {495namespace sh {
524inline constexpr float kC0 = 0.28209479177387814f;496inline constexpr float kC0 = 0.28209479177387814f;
525inline constexpr float kC1 = 0.4886025119029199f;497inline constexpr float kC1 = 0.4886025119029199f;
526inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,498inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
527 0.31539156525252005f, -1.0925484305920792f,499 0.31539156525252005f, -1.0925484305920792f,
528 0.5462742152960396f};500 0.5462742152960396f};
529inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,501inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
530 -0.4570457994644658f, 0.3731763325901154f,502 -0.4570457994644658f, 0.3731763325901154f,
531 -0.4570457994644658f, 1.445305721320277f,503 -0.4570457994644658f, 1.445305721320277f,
532 -0.5900435899266435f};504 -0.5900435899266435f};
533505
534// Basis values for a unit direction.506// Basis values for a unit direction.
535GS_D void Basis(float3 d, float b[kShCoeffs]) {507GS_D void Basis(float3 d, float b[kShCoeffs]) {
536 const float x = d.x, y = d.y, z = d.z;508 const float x = d.x, y = d.y, z = d.z;
537 const float xx = x * x, yy = y * y, zz = z * z;509 const float xx = x * x, yy = y * y, zz = z * z;
538 b[0] = kC0;510 b[0] = kC0;
539 b[1] = -kC1 * y;511 b[1] = -kC1 * y;
540 b[2] = kC1 * z;512 b[2] = kC1 * z;
541 b[3] = -kC1 * x;513 b[3] = -kC1 * x;
542 b[4] = kC2[0] * x * y;514 b[4] = kC2[0] * x * y;
543 b[5] = kC2[1] * y * z;515 b[5] = kC2[1] * y * z;
544 b[6] = kC2[2] * (2.0f * zz - xx - yy);516 b[6] = kC2[2] * (2.0f * zz - xx - yy);
545 b[7] = kC2[3] * x * z;517 b[7] = kC2[3] * x * z;
546 b[8] = kC2[4] * (xx - yy);518 b[8] = kC2[4] * (xx - yy);
547 b[9] = kC3[0] * y * (3.0f * xx - yy);519 b[9] = kC3[0] * y * (3.0f * xx - yy);
548 b[10] = kC3[1] * x * y * z;520 b[10] = kC3[1] * x * y * z;
549 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);521 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
550 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);
551 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);523 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
552 b[14] = kC3[5] * z * (xx - yy);524 b[14] = kC3[5] * z * (xx - yy);
553 b[15] = kC3[6] * x * (xx - 3.0f * yy);525 b[15] = kC3[6] * x * (xx - 3.0f * yy);
554}526}
555527
556// Partial derivatives of every basis function w.r.t. (x, y, z) of the unit
557// direction (treated as free coordinates).
558GS_D void BasisGrad(float3 d, float3 g[kShCoeffs]) {
559 const float x = d.x, y = d.y, z = d.z;
560 const float xx = x * x, yy = y * y, zz = z * z;
561 g[0] = {0.0f, 0.0f, 0.0f};
562 g[1] = {0.0f, -kC1, 0.0f};
563 g[2] = {0.0f, 0.0f, kC1};
564 g[3] = {-kC1, 0.0f, 0.0f};
565 g[4] = kC2[0] * float3{y, x, 0.0f};
566 g[5] = kC2[1] * float3{0.0f, z, y};
567 g[6] = kC2[2] * float3{-2.0f * x, -2.0f * y, 4.0f * z};
568 g[7] = kC2[3] * float3{z, 0.0f, x};
569 g[8] = kC2[4] * float3{2.0f * x, -2.0f * y, 0.0f};
570 g[9] = kC3[0] * float3{6.0f * x * y, 3.0f * (xx - yy), 0.0f};
571 g[10] = kC3[1] * float3{y * z, x * z, x * y};
572 g[11] =
573 kC3[2] * float3{-2.0f * x * y, 4.0f * zz - xx - 3.0f * yy, 8.0f * y * z};
574 g[12] = kC3[3] * float3{-6.0f * x * z, -6.0f * y * z,
575 6.0f * zz - 3.0f * xx - 3.0f * yy};
576 g[13] =
577 kC3[4] * float3{4.0f * zz - 3.0f * xx - yy, -2.0f * x * y, 8.0f * x * z};
578 g[14] = kC3[5] * float3{2.0f * x * z, -2.0f * y * z, xx - yy};
579 g[15] = kC3[6] * float3{3.0f * (xx - yy), -6.0f * x * y, 0.0f};
580}
581
582// Pre-clamp color 0.5 + sum_k b_k sh_k, loading the full 16-basis SH row.528// Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.
583template <typename ShT>529template <typename ShT>
584GS_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]) {
585 float3 coeffs[kShCoeffs];
586 ShIo<ShT>::LoadRow(sh_row, coeffs);
587 float3 c = {0.5f, 0.5f, 0.5f};531 float3 c = {0.5f, 0.5f, 0.5f};
588#pragma unroll532#pragma unroll
589 for (int k = 0; k < kShCoeffs; ++k) {533 for (int grp = 0; grp < 4; ++grp) {
534 const int k0 = 4 * grp;
535 float v[12];
536 ShIo<ShT>::LoadGroup(sh_row, grp, v);
590 c += b[k] * coeffs[k];537 c += b[k0] * float3{v[0], v[1], v[2]};
538 c += b[k0 + 1] * float3{v[3], v[4], v[5]};
539 c += b[k0 + 2] * float3{v[6], v[7], v[8]};
540 c += b[k0 + 3] * float3{v[9], v[10], v[11]};
591 }541 }
592 return c;542 return c;
593}543}
594} // namespace sh544} // namespace sh
↕ Expand 259 unchanged lines (lines 545–803) · in WarpMask
595545
596// ----------------------------------------------------------------------------546// ----------------------------------------------------------------------------
597// Covariance helpers.547// Covariance helpers.
598548
599// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).549// Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
600GS_D Mat3 QuatToRotation(float4 q) {550GS_D Mat3 QuatToRotation(float4 q) {
601 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;
602 Mat3 r;552 Mat3 r;
603 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);
604 r.m[0][1] = 2.0f * (x * y - w * z);554 r.m[0][1] = 2.0f * (x * y - w * z);
605 r.m[0][2] = 2.0f * (x * z + w * y);555 r.m[0][2] = 2.0f * (x * z + w * y);
606 r.m[1][0] = 2.0f * (x * y + w * z);556 r.m[1][0] = 2.0f * (x * y + w * z);
607 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);
608 r.m[1][2] = 2.0f * (y * z - w * x);558 r.m[1][2] = 2.0f * (y * z - w * x);
609 r.m[2][0] = 2.0f * (x * z - w * y);559 r.m[2][0] = 2.0f * (x * z - w * y);
610 r.m[2][1] = 2.0f * (y * z + w * x);560 r.m[2][1] = 2.0f * (y * z + w * x);
611 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);
612 return r;562 return r;
613}563}
614564
615// Symmetric 3x3 matrix.565// Symmetric 3x3 matrix.
616struct Sym3 {566struct Sym3 {
617 float xx, xy, xz, yy, yz, zz;567 float xx, xy, xz, yy, yz, zz;
618 GS_D float3 Mul(float3 v) const {568 GS_D float3 Mul(float3 v) const {
619 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,
620 xz * v.x + yz * v.y + zz * v.z};570 xz * v.x + yz * v.y + zz * v.z};
621 }571 }
622};572};
623573
624// Sigma = A A^T with A = R diag(s).574// Sigma = A A^T with A = R diag(s).
625GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {575GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
626 Mat3 a;576 Mat3 a;
627 const float sv[3] = {s.x, s.y, s.z};577 const float sv[3] = {s.x, s.y, s.z};
628#pragma unroll578#pragma unroll
629 for (int i = 0; i < 3; ++i)579 for (int i = 0; i < 3; ++i)
630#pragma unroll580#pragma unroll
631 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];
632 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)),
633 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)),
634 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))};
635}585}
636586
637// Intermediate values of the projection shared by forward and VJP.587// Intermediate values of the projection shared by forward and VJP.
638struct ProjectionState {588struct ProjectionState {
639 float3 t; // View-space mean.589 float3 t; // View-space mean.
640 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.
641 bool x_clamped, y_clamped;591 bool x_clamped, y_clamped;
642 float3 m0, m1; // Rows of M = J W (2x3).592 float3 m0, m1; // Rows of M = J W (2x3).
643 Sym3 sigma; // World-space 3D covariance.593 Sym3 sigma; // World-space 3D covariance.
644 float a, b, c; // Dilated 2D covariance.594 float a, b, c; // Dilated 2D covariance.
645};595};
646596
647// 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).
648GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,598GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
649 float4 quat, ProjectionState& st) {599 float4 quat, ProjectionState& st) {
650 st.t = t;600 st.t = t;
651 const float tz = t.z;601 const float tz = t.z;
652 const float limx = kFovClamp * cam.TanHalfFovX();602 const float limx = kFovClamp * cam.TanHalfFovX();
653 const float limy = kFovClamp * cam.TanHalfFovY();603 const float limy = kFovClamp * cam.TanHalfFovY();
654 const float txtz = t.x / tz, tytz = t.y / tz;604 const float txtz = t.x / tz, tytz = t.y / tz;
655 st.x_clamped = txtz < -limx || txtz > limx;605 st.x_clamped = txtz < -limx || txtz > limx;
656 st.y_clamped = tytz < -limy || tytz > limy;606 st.y_clamped = tytz < -limy || tytz > limy;
657 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;607 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
658 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;608 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
659 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);
660 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);
661 const Mat3& w = cam.rotation;611 const Mat3& w = cam.rotation;
662 st.m0 = j00 * w.row(0) + j02 * w.row(2);612 st.m0 = j00 * w.row(0) + j02 * w.row(2);
663 st.m1 = j11 * w.row(1) + j12 * w.row(2);613 st.m1 = j11 * w.row(1) + j12 * w.row(2);
664 st.sigma = Covariance3D(QuatToRotation(quat), scale);614 st.sigma = Covariance3D(QuatToRotation(quat), scale);
665 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);
666 st.a = dot(st.m0, s0) + kLowPass;616 st.a = dot(st.m0, s0) + kLowPass;
667 st.b = dot(st.m0, s1);617 st.b = dot(st.m0, s1);
668 st.c = dot(st.m1, s1) + kLowPass;618 st.c = dot(st.m1, s1) + kLowPass;
669}619}
670620
671// 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
672// near plane.622// near plane.
673GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,623GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
674 float4 quat, ProjectionState& st) {624 float4 quat, ProjectionState& st) {
675 const float3 t = cam.ToView(mean);625 const float3 t = cam.ToView(mean);
676 if (t.z <= kNearPlane) return false;626 if (t.z <= kNearPlane) return false;
677 ProjectCovarianceFromView(cam, t, scale, quat, st);627 ProjectCovarianceFromView(cam, t, scale, quat, st);
678 return true;628 return true;
679}629}
680630
681// 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.
682GS_D int BoundingRadius(float a, float b, float c) {632GS_D int BoundingRadius(float a, float b, float c) {
683 const float det = a * c - b * b;633 const float det = a * c - b * b;
684 const float mid = 0.5f * (a + c);634 const float mid = 0.5f * (a + c);
685 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));635 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
686 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));636 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
687}637}
688638
689GS_D float2 ProjectMean(const Camera& cam, float3 t) {639GS_D float2 ProjectMean(const Camera& cam, float3 t) {
690 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};
691}641}
692642
693GS_D float3 LoadMean(const GaussianParams& p, int i) {643GS_D float3 LoadMean(const GaussianParams& p, int i) {
694 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]};
695}645}
696646
697// Unit direction from the camera center to `mean`, and 1 / distance.647// Unit direction from the camera center to `mean`, and 1 / distance.
698GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {648GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
699 const float3 v = mean - cam.position;649 const float3 v = mean - cam.position;
700 inv_len = rsqrtf(dot(v, v));650 inv_len = rsqrtf(dot(v, v));
701 return inv_len * v;651 return inv_len * v;
702}652}
703653
704// 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.
705struct ProjectionGrad {655struct ProjectionGrad {
706 float2 xy;656 float2 xy;
707 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`.
708 float opacity;658 float opacity;
709 float3 rgb;659 float3 rgb;
710};660};
711661
712// Image geometry passed by value to every kernel.662// Image geometry passed by value to every kernel.
713struct Frame {663struct Frame {
714 const float* camera; // Device pointer to the packed camera.664 const float* camera; // Device pointer to the packed camera.
715 int width, height;665 int width, height;
716 TileGrid grid;666 TileGrid grid;
717 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }667 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
718};668};
719669
720// ---------------------------------------------------------------------------670// ---------------------------------------------------------------------------
721// Per-Gaussian kernels run one thread per Gaussian.671// Per-Gaussian kernels run one thread per Gaussian.
722672
723inline constexpr int kGaussianThreads = 128;673inline constexpr int kGaussianThreads = 128;
724674
725inline int GaussianBlocks(int64_t n) {675inline int GaussianBlocks(int64_t n) {
726 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);676 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
727}677}
728678
729// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.679// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
730inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2680inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
731inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;681inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
732682
733GS_HD float3 Power2Conic(float a, float b, float c) {683GS_HD float3 Power2Conic(float a, float b, float c) {
734 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};684 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
735}685}
736686
737GS_D float Exp2Approx(float x) {687GS_D float Exp2Approx(float x) {
738 float y;688 float y;
739 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));689 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
740 return y;690 return y;
741}691}
742692
743// ---------------------------------------------------------------------------693// ---------------------------------------------------------------------------
744// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into694// Rasterization. One block per 16x16 tile, one pixel per thread, arranged into
745// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through695// eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
746// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.696// shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
747697
748template <int kRows>698template <int kRows>
749struct RasterLayout {699struct RasterLayout {
750 static_assert(kRows == 1 || kRows == 2 || kRows == 4);700 static_assert(kRows == 1 || kRows == 2 || kRows == 4);
751 static constexpr int kThreads = kTilePixels / kRows;701 static constexpr int kThreads = kTilePixels / kRows;
752 static constexpr int kWarps = kThreads / 32;702 static constexpr int kWarps = kThreads / 32;
753 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;703 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
754 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;704 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
755 static constexpr int kWarpsX = kTileSize / kWarpWidth;705 static constexpr int kWarpsX = kTileSize / kWarpWidth;
756 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);706 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
757707
758 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }708 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
759 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }709 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
760};710};
761711
762struct RasterPixels {712struct RasterPixels {
763 int x, y0;713 int x, y0;
764 template <int kRows>714 template <int kRows>
765 GS_D static RasterPixels For(int tid) {715 GS_D static RasterPixels For(int tid) {
766 using L = RasterLayout<kRows>;716 using L = RasterLayout<kRows>;
767 const int warp = tid / 32, lane = tid % 32;717 const int warp = tid / 32, lane = tid % 32;
768 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +718 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
769 lane % L::kWarpWidth,719 lane % L::kWarpWidth,
770 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +720 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
771 (lane / L::kWarpWidth) * kRows};721 (lane / L::kWarpWidth) * kRows};
772 }722 }
773};723};
774724
775// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the725// Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
776// pixel region of warp w of this block's tile.726// pixel region of warp w of this block's tile.
777template <int kRows>727template <int kRows>
778GS_D uint32_t WarpMask(float4 p0, float4 p1) {728GS_D uint32_t WarpMask(float4 p0, float4 p1) {
779 using L = RasterLayout<kRows>;729 using L = RasterLayout<kRows>;
780 constexpr int kWarpsY = L::kWarps / L::kWarpsX;730 constexpr int kWarpsY = L::kWarps / L::kWarpsX;
781 float limit = p1.z;731 float limit = p1.z;
782 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;732 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
783 const float a = p0.z, b = p0.w, c = p1.x;733 const float a = p0.z, b = p0.w, c = p1.x;
784 const float det = a * c - b * b;734 const float det = a * c - b * b;
785 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;735 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
786 const float l_over_det = limit / det;736 const float l_over_det = limit / det;
787 const float rx = sqrtf(c * l_over_det);737 const float rx = sqrtf(c * l_over_det);
788 const float ry = sqrtf(a * l_over_det);738 const float ry = sqrtf(a * l_over_det);
789 const float2 xy = {p0.x, p0.y};739 const float2 xy = {p0.x, p0.y};
790 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));740 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
791 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));741 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
792 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));742 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
793 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));743 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
794 if (x0 > x1 || y0 > y1) return 0;744 if (x0 > x1 || y0 > y1) return 0;
795 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;745 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
796 uint32_t mx = 0u;746 uint32_t mx = 0u;
797#pragma unroll747#pragma unroll
798 for (int wx = 0; wx < L::kWarpsX; ++wx) {748 for (int wx = 0; wx < L::kWarpsX; ++wx) {
799 const int wx0 = tx + wx * L::kWarpWidth;749 const int wx0 = tx + wx * L::kWarpWidth;
800 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);750 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
801 }751 }
802 if (mx == 0u) return 0;752 if (mx == 0u) return 0;
803 uint32_t my = 0u;753 uint32_t my = 0u;
804#pragma unroll754#pragma unroll
805 for (int wy = 0; wy < kWarpsY; ++wy) {755 for (int wy = 0; wy < kWarpsY; ++wy) {
806 const int wy0 = ty + wy * L::kWarpRows;756 const int wy0 = ty + wy * L::kWarpRows;
807 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);757 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
808 }758 }
809 if (my == 0u) return 0;759 if (my == 0u) return 0;
810 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {760 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
811 uint32_t mask = 0u;761 uint32_t mask = 0u;
812#pragma unroll762#pragma unroll
813 for (int wy = 0; wy < kWarpsY; ++wy) {763 for (int wy = 0; wy < kWarpsY; ++wy) {
814 if (!((my >> wy) & 1u)) continue;764 if (!((my >> wy) & 1u)) continue;
815#pragma unroll765#pragma unroll
816 for (int wx = 0; wx < L::kWarpsX; ++wx) {766 for (int wx = 0; wx < L::kWarpsX; ++wx) {
817 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));767 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
818 }768 }
819 }769 }
820 return mask;770 return mask;
821 }771 }
822 EllipseBounds eb;772 EllipseBounds eb;
823 eb.a = a;773 eb.a = a;
824 eb.b = b;774 eb.b = b;
825 eb.c = c;775 eb.c = c;
826 eb.det = det;776 eb.det = det;
827 eb.inv_a = 1.0f / a;777 eb.inv_a = 1.0f / a;
828 eb.a_limit = a * limit;778 eb.a_limit = a * limit;
829 eb.y_ext = ry;779 eb.y_ext = ry;
830 eb.dy_peak = b * sqrtf(l_over_det / c);780 eb.dy_peak = b * sqrtf(l_over_det / c);
831 uint32_t mask = 0u;781 uint32_t mask = 0u;
832#pragma unroll782#pragma unroll
833 for (int wy = 0; wy < kWarpsY; ++wy) {783 for (int wy = 0; wy < kWarpsY; ++wy) {
834 if (!((my >> wy) & 1u)) continue;784 if (!((my >> wy) & 1u)) continue;
835 const int wy0 = ty + wy * L::kWarpRows;785 const int wy0 = ty + wy * L::kWarpRows;
836 int x0_px, x1_px;786 int x0_px, x1_px;
837 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {787 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
838 continue;788 continue;
839 }789 }
840#pragma unroll790#pragma unroll
841 for (int wx = 0; wx < L::kWarpsX; ++wx) {791 for (int wx = 0; wx < L::kWarpsX; ++wx) {
842 const int wx0 = tx + wx * L::kWarpWidth;792 const int wx0 = tx + wx * L::kWarpWidth;
843 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {793 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
844 mask |= 1u << (wy * L::kWarpsX + wx);794 mask |= 1u << (wy * L::kWarpsX + wx);
845 }795 }
846 }796 }
847 }797 }
848 return mask;798 return mask;
849}799}
850800
851} // namespace gs801} // namespace gs
852802
853#endif803#endif
↕ Expand 420 unchanged lines (lines 1–420) · in Projection
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
425+// and is accessed in 4 groups of 4 coefficients (12 values). The coefficients
425−// and 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.
426427
427428 using ShBf16 = uint16_t; // bfloat16 bits.
428429
429430 template <typename T>
430431 struct ShIo;
431432
432433 template <>
433434 struct ShIo<float> {
435+ static constexpr int kRowVec16 = 12; // 16-byte vectors per row.
434− GS_D static void LoadRow(const float* row, float3 coeffs[kShCoeffs]) {
436+ GS_D static void LoadGroup(const float* row, int grp, float v[12]) {
435− const float4* p = reinterpret_cast<const float4*>(row);
437+ const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;
436− float raw[3 * kShCoeffs];
437−#pragma unroll
438− for (int k = 0; k < 12; ++k) {
439− const float4 q = __ldg(p + k);
438+ const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);
440− raw[4 * k] = q.x;
441− raw[4 * k + 1] = q.y;
442− raw[4 * k + 2] = q.z;
443− raw[4 * k + 3] = q.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;
441+ v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;
444442⇥ }
445−#pragma unroll
446− for (int c = 0; c < kShCoeffs; ++c) {
447− coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
448− }
449− }
450− GS_D static void StoreRow(float* row, const float3 dsh[kShCoeffs]) {
443+ GS_D static void StoreGroup(float* row, int grp, const float v[12]) {
451− float raw[3 * kShCoeffs];
452−#pragma unroll
453− for (int c = 0; c < kShCoeffs; ++c) {
454− raw[3 * c] = dsh[c].x;
455− raw[3 * c + 1] = dsh[c].y;
456− raw[3 * c + 2] = dsh[c].z;
457− }
458− float4* p = reinterpret_cast<float4*>(row);
444+ float4* p = reinterpret_cast<float4*>(row) + 3 * grp;
459−#pragma unroll
460− for (int k = 0; k < 12; ++k) {
461− p[k] = {raw[4 * k], raw[4 * k + 1], raw[4 * k + 2], raw[4 * k + 3]};
445+ p[0] = {v[0], v[1], v[2], v[3]};
462− }
446+ p[1] = {v[4], v[5], v[6], v[7]};
447+ p[2] = {v[8], v[9], v[10], v[11]};
463448 }
464449 };
465450
466451 template <>
467452 struct ShIo<ShBf16> {
453+ static constexpr int kRowVec16 = 6;
468454 GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
469455 GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
470456 // Round to nearest even (inputs are finite).
471457 GS_D static uint32_t Round(float f) {
472458 const uint32_t u = __float_as_uint(f);
473459 return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
474460 }
475461 GS_D static uint32_t Pack(float lo, float hi) {
476462 return Round(lo) | (Round(hi) << 16);
477463 }
478− GS_D static void LoadRow(const ShBf16* row, float3 coeffs[kShCoeffs]) {
464+ GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {
465+ // 24-byte groups: three 8-byte loads.
479− const uint2* p = reinterpret_cast<const uint2*>(row);
466+ const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;
480− float raw[3 * kShCoeffs];
481467 #pragma unroll
482− for (int k = 0; k < 12; ++k) {
468+ for (int k = 0; k < 3; ++k) {
483469 const uint2 w = __ldg(p + k);
484− raw[4 * k] = Lo(w.x);
485− raw[4 * k + 1] = Hi(w.x);
470+ v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);
486− raw[4 * k + 2] = Lo(w.y);
487− raw[4 * k + 3] = Hi(w.y);
471+ v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);
488− }
489−#pragma unroll
490− for (int c = 0; c < kShCoeffs; ++c) {
491− coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
492472 }
493473 }
494− GS_D static void StoreRow(ShBf16* row, const float3 dsh[kShCoeffs]) {
474+ GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {
495− float raw[3 * kShCoeffs];
475+ uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;
496476 #pragma unroll
497− for (int c = 0; c < kShCoeffs; ++c) {
498− raw[3 * c] = dsh[c].x;
499− raw[3 * c + 1] = dsh[c].y;
500− raw[3 * c + 2] = dsh[c].z;
501− }
502− uint2* p = reinterpret_cast<uint2*>(row);
503−#pragma unroll
504− for (int k = 0; k < 12; ++k) {
477+ for (int k = 0; k < 3; ++k) {
505− p[k] = {Pack(raw[4 * k], raw[4 * k + 1]),
478+ p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};
506− Pack(raw[4 * k + 2], raw[4 * k + 3])};
507479 }
508480 }
509481 };
510482
↕ Expand 41 unchanged lines (lines 483–523) · in Basis
511483 template <typename ShT>
512484 GS_D const ShT* ShRow(const void* sh, int i) {
513485 return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
514486 }
515487 template <typename ShT>
516488 GS_D ShT* ShRow(void* sh, int i) {
517489 return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
518490 }
519491
520492 // ----------------------------------------------------------------------------
521493 // Spherical harmonics (degree 3, Inria sign convention).
522494
523495 namespace sh {
524496 inline constexpr float kC0 = 0.28209479177387814f;
525497 inline constexpr float kC1 = 0.4886025119029199f;
526498 inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
527499 0.31539156525252005f, -1.0925484305920792f,
528500 0.5462742152960396f};
529501 inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
530502 -0.4570457994644658f, 0.3731763325901154f,
531503 -0.4570457994644658f, 1.445305721320277f,
532504 -0.5900435899266435f};
533505
534506 // Basis values for a unit direction.
535507 GS_D void Basis(float3 d, float b[kShCoeffs]) {
536508 const float x = d.x, y = d.y, z = d.z;
537509 const float xx = x * x, yy = y * y, zz = z * z;
538510 b[0] = kC0;
539511 b[1] = -kC1 * y;
540512 b[2] = kC1 * z;
541513 b[3] = -kC1 * x;
542514 b[4] = kC2[0] * x * y;
543515 b[5] = kC2[1] * y * z;
544516 b[6] = kC2[2] * (2.0f * zz - xx - yy);
545517 b[7] = kC2[3] * x * z;
546518 b[8] = kC2[4] * (xx - yy);
547519 b[9] = kC3[0] * y * (3.0f * xx - yy);
548520 b[10] = kC3[1] * x * y * z;
549521 b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
550522 b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
551523 b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
552524 b[14] = kC3[5] * z * (xx - yy);
553525 b[15] = kC3[6] * x * (xx - 3.0f * yy);
554526 }
555527
556−// Partial derivatives of every basis function w.r.t. (x, y, z) of the unit
557−// direction (treated as free coordinates).
558−GS_D void BasisGrad(float3 d, float3 g[kShCoeffs]) {
559− const float x = d.x, y = d.y, z = d.z;
560− const float xx = x * x, yy = y * y, zz = z * z;
561− g[0] = {0.0f, 0.0f, 0.0f};
562− g[1] = {0.0f, -kC1, 0.0f};
563− g[2] = {0.0f, 0.0f, kC1};
564− g[3] = {-kC1, 0.0f, 0.0f};
565− g[4] = kC2[0] * float3{y, x, 0.0f};
566− g[5] = kC2[1] * float3{0.0f, z, y};
567− g[6] = kC2[2] * float3{-2.0f * x, -2.0f * y, 4.0f * z};
568− g[7] = kC2[3] * float3{z, 0.0f, x};
569− g[8] = kC2[4] * float3{2.0f * x, -2.0f * y, 0.0f};
570− g[9] = kC3[0] * float3{6.0f * x * y, 3.0f * (xx - yy), 0.0f};
571− g[10] = kC3[1] * float3{y * z, x * z, x * y};
572− g[11] =
573− kC3[2] * float3{-2.0f * x * y, 4.0f * zz - xx - 3.0f * yy, 8.0f * y * z};
574− g[12] = kC3[3] * float3{-6.0f * x * z, -6.0f * y * z,
575− 6.0f * zz - 3.0f * xx - 3.0f * yy};
576− g[13] =
577− kC3[4] * float3{4.0f * zz - 3.0f * xx - yy, -2.0f * x * y, 8.0f * x * z};
578− g[14] = kC3[5] * float3{2.0f * x * z, -2.0f * y * z, xx - yy};
579− g[15] = kC3[6] * float3{3.0f * (xx - yy), -6.0f * x * y, 0.0f};
580−}
581−
582−// Pre-clamp color 0.5 + sum_k b_k sh_k, loading the full 16-basis SH row.
528+// Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.
583529 template <typename ShT>
584530 GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
585− float3 coeffs[kShCoeffs];
586− ShIo<ShT>::LoadRow(sh_row, coeffs);
587531 float3 c = {0.5f, 0.5f, 0.5f};
588532 #pragma unroll
589− for (int k = 0; k < kShCoeffs; ++k) {
533+ for (int grp = 0; grp < 4; ++grp) {
534+ const int k0 = 4 * grp;
535+ float v[12];
536+ ShIo<ShT>::LoadGroup(sh_row, grp, v);
590− c += b[k] * coeffs[k];
537+ c += b[k0] * float3{v[0], v[1], v[2]};
538+ c += b[k0 + 1] * float3{v[3], v[4], v[5]};
539+ c += b[k0 + 2] * float3{v[6], v[7], v[8]};
540+ c += b[k0 + 3] * float3{v[9], v[10], v[11]};
591541 }
592542 return c;
593543 }
594544 } // namespace sh
↕ Expand 259 unchanged lines (lines 545–803) · in WarpMask
595545
596546 // ----------------------------------------------------------------------------
597547 // Covariance helpers.
598548
599549 // Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
600550 GS_D Mat3 QuatToRotation(float4 q) {
601551 const float w = q.x, x = q.y, y = q.z, z = q.w;
602552 Mat3 r;
603553 r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
604554 r.m[0][1] = 2.0f * (x * y - w * z);
605555 r.m[0][2] = 2.0f * (x * z + w * y);
606556 r.m[1][0] = 2.0f * (x * y + w * z);
607557 r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
608558 r.m[1][2] = 2.0f * (y * z - w * x);
609559 r.m[2][0] = 2.0f * (x * z - w * y);
610560 r.m[2][1] = 2.0f * (y * z + w * x);
611561 r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
612562 return r;
613563 }
614564
615565 // Symmetric 3x3 matrix.
616566 struct Sym3 {
617567 float xx, xy, xz, yy, yz, zz;
618568 GS_D float3 Mul(float3 v) const {
619569 return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
620570 xz * v.x + yz * v.y + zz * v.z};
621571 }
622572 };
623573
624574 // Sigma = A A^T with A = R diag(s).
625575 GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
626576 Mat3 a;
627577 const float sv[3] = {s.x, s.y, s.z};
628578 #pragma unroll
629579 for (int i = 0; i < 3; ++i)
630580 #pragma unroll
631581 for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
632582 return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
633583 dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
634584 dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
635585 }
636586
637587 // Intermediate values of the projection shared by forward and VJP.
638588 struct ProjectionState {
639589 float3 t; // View-space mean.
640590 float tx_c, ty_c; // Clamped x, y used by the Jacobian.
641591 bool x_clamped, y_clamped;
642592 float3 m0, m1; // Rows of M = J W (2x3).
643593 Sym3 sigma; // World-space 3D covariance.
644594 float a, b, c; // Dilated 2D covariance.
645595 };
646596
647597 // Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
648598 GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
649599 float4 quat, ProjectionState& st) {
650600 st.t = t;
651601 const float tz = t.z;
652602 const float limx = kFovClamp * cam.TanHalfFovX();
653603 const float limy = kFovClamp * cam.TanHalfFovY();
654604 const float txtz = t.x / tz, tytz = t.y / tz;
655605 st.x_clamped = txtz < -limx || txtz > limx;
656606 st.y_clamped = tytz < -limy || tytz > limy;
657607 st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
658608 st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
659609 const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
660610 const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
661611 const Mat3& w = cam.rotation;
662612 st.m0 = j00 * w.row(0) + j02 * w.row(2);
663613 st.m1 = j11 * w.row(1) + j12 * w.row(2);
664614 st.sigma = Covariance3D(QuatToRotation(quat), scale);
665615 const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
666616 st.a = dot(st.m0, s0) + kLowPass;
667617 st.b = dot(st.m0, s1);
668618 st.c = dot(st.m1, s1) + kLowPass;
669619 }
670620
671621 // Computes everything up to the 2D covariance. Returns false if culled by the
672622 // near plane.
673623 GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
674624 float4 quat, ProjectionState& st) {
675625 const float3 t = cam.ToView(mean);
676626 if (t.z <= kNearPlane) return false;
677627 ProjectCovarianceFromView(cam, t, scale, quat, st);
678628 return true;
679629 }
680630
681631 // Integer bounding radius (pixels) of the 2D covariance, as in the reference.
682632 GS_D int BoundingRadius(float a, float b, float c) {
683633 const float det = a * c - b * b;
684634 const float mid = 0.5f * (a + c);
685635 const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
686636 return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
687637 }
688638
689639 GS_D float2 ProjectMean(const Camera& cam, float3 t) {
690640 return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
691641 }
692642
693643 GS_D float3 LoadMean(const GaussianParams& p, int i) {
694644 return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
695645 }
696646
697647 // Unit direction from the camera center to `mean`, and 1 / distance.
698648 GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
699649 const float3 v = mean - cam.position;
700650 inv_len = rsqrtf(dot(v, v));
701651 return inv_len * v;
702652 }
703653
704654 // Gradient of the loss w.r.t. one splat's screen-space quantities.
705655 struct ProjectionGrad {
706656 float2 xy;
707657 float3 conic; // w.r.t. the scalars (a, b, c) of `Projection::conic`.
708658 float opacity;
709659 float3 rgb;
710660 };
711661
712662 // Image geometry passed by value to every kernel.
713663 struct Frame {
714664 const float* camera; // Device pointer to the packed camera.
715665 int width, height;
716666 TileGrid grid;
717667 GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
718668 };
719669
720670 // ---------------------------------------------------------------------------
721671 // Per-Gaussian kernels run one thread per Gaussian.
722672
723673 inline constexpr int kGaussianThreads = 128;
724674
725675 inline int GaussianBlocks(int64_t n) {
726676 return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
727677 }
728678
729679 // Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
730680 inline constexpr float kPower2Scale = -0.72134752044448170f; // -log2(e)/2
731681 inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
732682
733683 GS_HD float3 Power2Conic(float a, float b, float c) {
734684 return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
735685 }
736686
737687 GS_D float Exp2Approx(float x) {
738688 float y;
739689 asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
740690 return y;
741691 }
742692
743693 // ---------------------------------------------------------------------------
744694 // Rasterization. One block per 16x16 tile, one pixel per thread, arranged into
745695 // eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
746696 // shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
747697
748698 template <int kRows>
749699 struct RasterLayout {
750700 static_assert(kRows == 1 || kRows == 2 || kRows == 4);
751701 static constexpr int kThreads = kTilePixels / kRows;
752702 static constexpr int kWarps = kThreads / 32;
753703 static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
754704 static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
755705 static constexpr int kWarpsX = kTileSize / kWarpWidth;
756706 static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
757707
758708 GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
759709 GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
760710 };
761711
762712 struct RasterPixels {
763713 int x, y0;
764714 template <int kRows>
765715 GS_D static RasterPixels For(int tid) {
766716 using L = RasterLayout<kRows>;
767717 const int warp = tid / 32, lane = tid % 32;
768718 return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
769719 lane % L::kWarpWidth,
770720 static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
771721 (lane / L::kWarpWidth) * kRows};
772722 }
773723 };
774724
775725 // Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
776726 // pixel region of warp w of this block's tile.
777727 template <int kRows>
778728 GS_D uint32_t WarpMask(float4 p0, float4 p1) {
779729 using L = RasterLayout<kRows>;
780730 constexpr int kWarpsY = L::kWarps / L::kWarpsX;
781731 float limit = p1.z;
782732 if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
783733 const float a = p0.z, b = p0.w, c = p1.x;
784734 const float det = a * c - b * b;
785735 if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
786736 const float l_over_det = limit / det;
787737 const float rx = sqrtf(c * l_over_det);
788738 const float ry = sqrtf(a * l_over_det);
789739 const float2 xy = {p0.x, p0.y};
790740 const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
791741 const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
792742 const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
793743 const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
794744 if (x0 > x1 || y0 > y1) return 0;
795745 const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
796746 uint32_t mx = 0u;
797747 #pragma unroll
798748 for (int wx = 0; wx < L::kWarpsX; ++wx) {
799749 const int wx0 = tx + wx * L::kWarpWidth;
800750 if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
801751 }
802752 if (mx == 0u) return 0;
803753 uint32_t my = 0u;
804754 #pragma unroll
805755 for (int wy = 0; wy < kWarpsY; ++wy) {
806756 const int wy0 = ty + wy * L::kWarpRows;
807757 if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
808758 }
809759 if (my == 0u) return 0;
810760 if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
811761 uint32_t mask = 0u;
812762 #pragma unroll
813763 for (int wy = 0; wy < kWarpsY; ++wy) {
814764 if (!((my >> wy) & 1u)) continue;
815765 #pragma unroll
816766 for (int wx = 0; wx < L::kWarpsX; ++wx) {
817767 if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
818768 }
819769 }
820770 return mask;
821771 }
822772 EllipseBounds eb;
823773 eb.a = a;
824774 eb.b = b;
825775 eb.c = c;
826776 eb.det = det;
827777 eb.inv_a = 1.0f / a;
828778 eb.a_limit = a * limit;
829779 eb.y_ext = ry;
830780 eb.dy_peak = b * sqrtf(l_over_det / c);
831781 uint32_t mask = 0u;
832782 #pragma unroll
833783 for (int wy = 0; wy < kWarpsY; ++wy) {
834784 if (!((my >> wy) & 1u)) continue;
835785 const int wy0 = ty + wy * L::kWarpRows;
836786 int x0_px, x1_px;
837787 if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
838788 continue;
839789 }
840790 #pragma unroll
841791 for (int wx = 0; wx < L::kWarpsX; ++wx) {
842792 const int wx0 = tx + wx * L::kWarpWidth;
843793 if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
844794 mask |= 1u << (wy * L::kWarpsX + wx);
845795 }
846796 }
847797 }
848798 return mask;
849799 }
850800
851801 } // namespace gs
852802
853803 #endif
csrc/backward.cu.cc +123 −45
s02_reduce_scatter/csrc/backward.cu.cc
s03_sh_streaming/csrc/backward.cu.cc
↕ 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#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 loading the full 16-basis row `sh_in` and storing19// Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly
20// `dsh` to `sh_out`.
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 float3 coeffs[kShCoeffs];
27 ShIo<ShT>::LoadRow(sh_in, coeffs);
28 float3 raw = {0.5f, 0.5f, 0.5f};26 const float3 raw = EvalRow(sh_in, b);
29#pragma unroll
30 for (int k = 0; k < kShCoeffs; ++k) {
31 raw += b[k] * coeffs[k];
32 }
33 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,27 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
34 raw.y < 0.0f ? 0.0f : g_rgb.y,28 raw.y < 0.0f ? 0.0f : g_rgb.y,
35 raw.z < 0.0f ? 0.0f : g_rgb.z};29 raw.z < 0.0f ? 0.0f : g_rgb.z};
36 float3 bgrad[kShCoeffs];
37 BasisGrad(d, bgrad);
38 float3 dsh[kShCoeffs];
30 const float x = d.x, y = d.y, z = d.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;
39 float3 ddir = {0.0f, 0.0f, 0.0f};33 float3 ddir = {0.0f, 0.0f, 0.0f};
40#pragma unroll34#pragma unroll
35 for (int grp = 0; grp < 4; ++grp) {
36 const int k0 = 4 * grp;
37 float v[12];
38 ShIo<ShT>::LoadGroup(sh_in, grp, v);
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;
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;
43 float out[12];
44#pragma unroll
41 for (int k = 0; k < kShCoeffs; ++k) {45 for (int q = 0; q < 4; ++q) {
42 dsh[k] = b[k] * drgb;46 out[3 * q] = b[k0 + q] * drgb.x;
43 const float s = dot(coeffs[k], drgb);47 out[3 * q + 1] = b[k0 + q] * drgb.y;
44 ddir += s * bgrad[k];
48 out[3 * q + 2] = b[k0 + q] * drgb.z;
45 }49 }
46 ShIo<ShT>::StoreRow(sh_out, dsh);50 ShIo<ShT>::StoreGroup(sh_out, grp, out);
51 if (grp == 0) {
52 ddir.y += s1 * (-kC1);
53 ddir.z += s2 * kC1;
54 ddir.x += s3 * (-kC1);
55 } else if (grp == 1) {
56 ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};
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};
59 ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};
60 } else if (grp == 2) {
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};
63 ddir += (s2 * kC3[1]) * float3{yz, xz, xy};
64 ddir += (s3 * kC3[2]) *
65 float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};
66 } else {
67 ddir += (s0 * kC3[3]) *
68 float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};
69 ddir += (s1 * kC3[4]) *
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};
72 ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};
73 }
74 }
47 return ddir;75 return ddir;
48}76}
49} // namespace sh77} // namespace sh
5078
↕ Expand 373 unchanged lines (lines 79–451) · in RasterBackwardKernel
51// 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
52// `sh_in` and writes dL/dsh directly to `sh_out`.80// `sh_in` and writes dL/dsh directly to `sh_out`.
53// 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`.
54template <typename ShT>82template <typename ShT>
55GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,83GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
56 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,84 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
57 const ProjectionGrad& g, float3& dmean_out,85 const ProjectionGrad& g, float3& dmean_out,
58 float3& dscale_out, float4& dquat_out) {86 float3& dscale_out, float4& dquat_out) {
59 const float3 mean = LoadMean(p, i);87 const float3 mean = LoadMean(p, i);
60 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],
61 p.scales[3 * i + 2]};89 p.scales[3 * i + 2]};
62 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];90 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
6391
64 // --- 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).
65 float3 dmean = {0.0f, 0.0f, 0.0f};93 float3 dmean = {0.0f, 0.0f, 0.0f};
66 {94 {
67 float inv_len;95 float inv_len;
68 const float3 dir = ViewDir(cam, mean, inv_len);96 const float3 dir = ViewDir(cam, mean, inv_len);
69 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);
70 // d(v/|v|)/dv = (I - dir dir^T) / |v|.98 // d(v/|v|)/dv = (I - dir dir^T) / |v|.
71 dmean += inv_len * (ddir - dot(dir, ddir) * dir);99 dmean += inv_len * (ddir - dot(dir, ddir) * dir);
72 }100 }
73101
74 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).102 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
75 const float3 t = cam.ToView(mean);103 const float3 t = cam.ToView(mean);
76 const float tz = t.z;104 const float tz = t.z;
77 const float limx = kFovClamp * cam.TanHalfFovX();105 const float limx = kFovClamp * cam.TanHalfFovX();
78 const float limy = kFovClamp * cam.TanHalfFovY();106 const float limy = kFovClamp * cam.TanHalfFovY();
79 const float txtz = t.x / tz, tytz = t.y / tz;107 const float txtz = t.x / tz, tytz = t.y / tz;
80 const bool x_clamped = txtz < -limx || txtz > limx;108 const bool x_clamped = txtz < -limx || txtz > limx;
81 const bool y_clamped = tytz < -limy || tytz > limy;109 const bool y_clamped = tytz < -limy || tytz > limy;
82 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;110 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
83 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;111 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
84 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);
85 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);
86 const Mat3& w = cam.rotation;114 const Mat3& w = cam.rotation;
87 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);115 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
88 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);116 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
89117
90 const Mat3 r = QuatToRotation(quat);118 const Mat3 r = QuatToRotation(quat);
91 const Sym3 sigma = Covariance3D(r, scale);119 const Sym3 sigma = Covariance3D(r, scale);
92 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);120 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
93 const float cov_a = dot(m0, s0) + kLowPass;121 const float cov_a = dot(m0, s0) + kLowPass;
94 const float cov_b = dot(m0, s1);122 const float cov_b = dot(m0, s1);
95 const float cov_c = dot(m1, s1) + kLowPass;123 const float cov_c = dot(m1, s1) + kLowPass;
96124
97 // --- Conic -> 2D covariance: G_cov = -K G_K K.125 // --- Conic -> 2D covariance: G_cov = -K G_K K.
98 const float det = cov_a * cov_c - cov_b * cov_b;126 const float det = cov_a * cov_c - cov_b * cov_b;
99 const float inv = 1.0f / det;127 const float inv = 1.0f / det;
100 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;
101 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;
102 // P = G_K K.130 // P = G_K K.
103 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;
104 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;
105 // G_cov = -K P (symmetric).133 // G_cov = -K P (symmetric).
106 const float ca = -(ka * p00 + kb * p10);134 const float ca = -(ka * p00 + kb * p10);
107 const float cb = -(ka * p01 + kb * p11);135 const float cb = -(ka * p01 + kb * p11);
108 const float cc = -(kb * p01 + kc * p11);136 const float cc = -(kb * p01 + kc * p11);
109137
110 // --- 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.
111 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);139 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
112 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);140 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
113 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.141 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
114 Sym3 dsig;142 Sym3 dsig;
115 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;
116 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;
117 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;
118 dsig.xy =146 dsig.xy =
119 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;
120 dsig.xz =148 dsig.xz =
121 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;
122 dsig.yz =150 dsig.yz =
123 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;
124152
125 // --- 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).
126 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));
127 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));
128 const float tz2 = tz * tz, tz3 = tz2 * tz;156 const float tz2 = tz * tz, tz3 = tz2 * tz;
129 float3 dt;157 float3 dt;
130 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;158 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
131 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;159 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
132 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +160 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
133 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;161 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
134 // --- Screen-space mean u = fx x / z + cx.162 // --- Screen-space mean u = fx x / z + cx.
135 dt.x += cam.fx / tz * g.xy.x;163 dt.x += cam.fx / tz * g.xy.x;
136 dt.y += cam.fy / tz * g.xy.y;164 dt.y += cam.fy / tz * g.xy.y;
137 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;165 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
138 dmean += MulT(w, dt);166 dmean += MulT(w, dt);
139 dmean_out = dmean;167 dmean_out = dmean;
140168
141 // --- 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.
142 const float sv[3] = {scale.x, scale.y, scale.z};170 const float sv[3] = {scale.x, scale.y, scale.z};
143 float ds[3];171 float ds[3];
144 Mat3 dr;172 Mat3 dr;
145#pragma unroll173#pragma unroll
146 for (int j = 0; j < 3; ++j) {174 for (int j = 0; j < 3; ++j) {
147 const float3 r_col = r.col(j);175 const float3 r_col = r.col(j);
148 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);
149 ds[j] = dot(da_col, r_col);177 ds[j] = dot(da_col, r_col);
150 dr.m[0][j] = da_col.x * sv[j];178 dr.m[0][j] = da_col.x * sv[j];
151 dr.m[1][j] = da_col.y * sv[j];179 dr.m[1][j] = da_col.y * sv[j];
152 dr.m[2][j] = da_col.z * sv[j];180 dr.m[2][j] = da_col.z * sv[j];
153 }181 }
154 dscale_out = {ds[0], ds[1], ds[2]};182 dscale_out = {ds[0], ds[1], ds[2]};
155183
156 // --- R(q) -> q.184 // --- R(q) -> q.
157 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;
158 const auto& G = dr.m;186 const auto& G = dr.m;
159 dquat_out = {187 dquat_out = {
160 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] -
161 qy * G[2][0] + qx * G[2][1]),189 qy * G[2][0] + qx * G[2][1]),
162 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -190 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
163 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]),
164 2.0f *192 2.0f *
165 (-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] +
166 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]),
167 2.0f *195 2.0f *
168 (-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] -
169 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])};
170}198}
171199
172// Reverse traversal per pixel; accumulates splat gradients via Half-Warp200// Reverse traversal per pixel; accumulates splat gradients via Half-Warp
173// Two-Splat Folding and Halving Register Reduce-Scatter.201// Two-Splat Folding and Halving Register Reduce-Scatter.
174__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 4)202__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 4)
175 RasterBackwardKernel(const int2* __restrict__ ranges,203 RasterBackwardKernel(const int2* __restrict__ ranges,
176 const int32_t* __restrict__ point_list,204 const int32_t* __restrict__ point_list,
177 const Splat* __restrict__ splats, Frame frame,205 const Splat* __restrict__ splats, Frame frame,
178 const float* __restrict__ final_t,206 const float* __restrict__ final_t,
179 const int32_t* __restrict__ n_contrib,207 const int32_t* __restrict__ n_contrib,
180 const float* __restrict__ dl_dimage,208 const float* __restrict__ dl_dimage,
181 float* __restrict__ splat_grads) {209 float* __restrict__ splat_grads) {
182 using L = RasterLayout<1>;210 using L = RasterLayout<1>;
183 constexpr int kRasterThreads = L::kThreads;211 constexpr int kRasterThreads = L::kThreads;
184 constexpr int kWarps = L::kWarps;212 constexpr int kWarps = L::kWarps;
185 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;213 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
186 const int tid = threadIdx.x;214 const int tid = threadIdx.x;
187 const int warp = tid / 32;215 const int warp = tid / 32;
188 const int lane = tid % 32;216 const int lane = tid % 32;
189 const RasterPixels px = RasterPixels::For<1>(tid);217 const RasterPixels px = RasterPixels::For<1>(tid);
190 const float fx = static_cast<float>(px.x);218 const float fx = static_cast<float>(px.x);
191 const float fy = static_cast<float>(px.y0);219 const float fy = static_cast<float>(px.y0);
192 const int2 range = ranges[tile];220 const int2 range = ranges[tile];
193221
194 const bool inside = px.x < frame.width && px.y0 < frame.height;222 const bool inside = px.x < frame.width && px.y0 < frame.height;
195 const int pixels = frame.width * frame.height;223 const int pixels = frame.width * frame.height;
196 const int p = px.y0 * frame.width + px.x;224 const int p = px.y0 * frame.width + px.x;
197 float t = inside ? final_t[p] : 0.0f;225 float t = inside ? final_t[p] : 0.0f;
198 const float3 dpix =226 const float3 dpix =
199 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],227 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
200 dl_dimage[2 * pixels + p]}228 dl_dimage[2 * pixels + p]}
201 : float3{0.0f, 0.0f, 0.0f};229 : float3{0.0f, 0.0f, 0.0f};
202 float accum = 0.0f; // dot(color accumulated behind, dpix)230 float accum = 0.0f; // dot(color accumulated behind, dpix)
203 const int last = inside ? n_contrib[p] : 0;231 const int last = inside ? n_contrib[p] : 0;
204232
205 int warp_max_last = last;233 int warp_max_last = last;
206#pragma unroll234#pragma unroll
207 for (int offset = 16; offset > 0; offset /= 2) {235 for (int offset = 16; offset > 0; offset /= 2) {
208 warp_max_last = max(236 warp_max_last = max(
209 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));237 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
210 }238 }
211 __shared__ int shared_max_last[kWarps];239 __shared__ int shared_max_last[kWarps];
212 if (lane == 0) shared_max_last[warp] = warp_max_last;240 if (lane == 0) shared_max_last[warp] = warp_max_last;
213 __syncthreads();241 __syncthreads();
214 int block_max_last = 0;242 int block_max_last = 0;
215#pragma unroll243#pragma unroll
216 for (int w = 0; w < kWarps; ++w) {244 for (int w = 0; w < kWarps; ++w) {
217 block_max_last = max(block_max_last, shared_max_last[w]);245 block_max_last = max(block_max_last, shared_max_last[w]);
218 }246 }
219 if (block_max_last <= 0) return;247 if (block_max_last <= 0) return;
220248
221 __shared__ float4 batch_geom0[kRasterThreads];249 __shared__ float4 batch_geom0[kRasterThreads];
222 __shared__ float2 batch_geom1[kRasterThreads];250 __shared__ float2 batch_geom1[kRasterThreads];
223 __shared__ float4 batch_rgb_o[kRasterThreads];251 __shared__ float4 batch_rgb_o[kRasterThreads];
224 __shared__ uint32_t batch_id[kRasterThreads];252 __shared__ uint32_t batch_id[kRasterThreads];
225 __shared__ uint8_t batch_mask[kRasterThreads];253 __shared__ uint8_t batch_mask[kRasterThreads];
226 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);254 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
227 int scatter = -1;255 int scatter = -1;
228 int scatter2 = -1;256 int scatter2 = -1;
229 {257 {
230 const int sub_lane = lane & 15;258 const int sub_lane = lane & 15;
231 const int row = sub_lane >> 3, col = sub_lane & 7;259 const int row = sub_lane >> 3, col = sub_lane & 7;
232 if (col == 0) {260 if (col == 0) {
233 scatter2 = row;261 scatter2 = row;
234 } else if (col == 2) {262 } else if (col == 2) {
235 scatter2 = 2 + row;263 scatter2 = 2 + row;
236 } else if (col == 4) {264 } else if (col == 4) {
237 scatter2 = row == 0 ? 4 : 7;265 scatter2 = row == 0 ? 4 : 7;
238 } else if (col == 6) {266 } else if (col == 6) {
239 scatter2 = row == 0 ? 6 : 8;267 scatter2 = row == 0 ? 6 : 8;
240 } else if (col == 1 && row == 0) {268 } else if (col == 1 && row == 0) {
241 scatter2 = 5;269 scatter2 = 5;
242 }270 }
243 const int full_row = lane >> 3, full_col = lane & 7;271 const int full_row = lane >> 3, full_col = lane & 7;
244 if (full_col == 0) {272 if (full_col == 0) {
245 scatter = full_row;273 scatter = full_row;
246 } else if (full_col == 2) {274 } else if (full_col == 2) {
247 scatter = full_row == 0 ? 4 : 5 + full_row;275 scatter = full_row == 0 ? 4 : 5 + full_row;
248 } else if (full_col == 4 && full_row == 0) {276 } else if (full_col == 4 && full_row == 0) {
249 scatter = 5;277 scatter = 5;
250 }278 }
251 }279 }
252280
253 auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,281 auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,
254 float& ga_sum, float& ga_dy, float& ga_dy2,282 float& ga_sum, float& ga_dy, float& ga_dy2,
255 float3& drgb) -> bool {283 float3& drgb) -> bool {
256 dx = g0.x - fx;284 dx = g0.x - fx;
257 ga_sum = 0.0f;285 ga_sum = 0.0f;
258 ga_dy = 0.0f;286 ga_dy = 0.0f;
259 ga_dy2 = 0.0f;287 ga_dy2 = 0.0f;
260 drgb = {0.0f, 0.0f, 0.0f};288 drgb = {0.0f, 0.0f, 0.0f};
261 if (entry < last) {289 if (entry < last) {
262 const float dy = g0.y - fy;290 const float dy = g0.y - fy;
263 const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;291 const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
264 if (power <= 0.0f && power >= g1.y) {292 if (power <= 0.0f && power >= g1.y) {
265 const float g_val = Exp2Approx(power);293 const float g_val = Exp2Approx(power);
266 const float alpha = fminf(kMaxAlpha, ro.w * g_val);294 const float alpha = fminf(kMaxAlpha, ro.w * g_val);
267 if (alpha >= kMinAlpha) {295 if (alpha >= kMinAlpha) {
268 t = __fdividef(t, 1.0f - alpha);296 t = __fdividef(t, 1.0f - alpha);
269 const float w = alpha * t;297 const float w = alpha * t;
270 const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;298 const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
271 const float diff = color_dot - accum;299 const float diff = color_dot - accum;
272 const float dl_dalpha = t * diff;300 const float dl_dalpha = t * diff;
273 accum += alpha * diff;301 accum += alpha * diff;
274 const float ga = g_val * dl_dalpha;302 const float ga = g_val * dl_dalpha;
275 ga_sum = ga;303 ga_sum = ga;
276 const float ga_y = ga * dy;304 const float ga_y = ga * dy;
277 ga_dy = ga_y;305 ga_dy = ga_y;
278 ga_dy2 = ga_y * dy;306 ga_dy2 = ga_y * dy;
279 drgb = w * dpix;307 drgb = w * dpix;
280 return true;308 return true;
281 }309 }
282 }310 }
283 }311 }
284 return false;312 return false;
285 };313 };
286 auto raw_conic = [](float4 g0, float2 g1) {314 auto raw_conic = [](float4 g0, float2 g1) {
287 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,315 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
288 kInvPower2Scale * g1.x};316 kInvPower2Scale * g1.x};
289 };317 };
290318
291 for (int end = block_max_last; end > 0; end -= kRasterThreads) {319 for (int end = block_max_last; end > 0; end -= kRasterThreads) {
292 const int e = end - 1 - tid;320 const int e = end - 1 - tid;
293 if (e >= 0) {321 if (e >= 0) {
294 const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);322 const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
295 const float4 p0 = __ldg(&splats[id].p0);323 const float4 p0 = __ldg(&splats[id].p0);
296 const float4 p1 = __ldg(&splats[id].p1);324 const float4 p1 = __ldg(&splats[id].p1);
297 const float4 p2 = __ldg(&splats[id].p2);325 const float4 p2 = __ldg(&splats[id].p2);
298 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);326 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
299 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};327 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
300 batch_geom1[tid] = {c2.z, p2.w};328 batch_geom1[tid] = {c2.z, p2.w};
301 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};329 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
302 batch_id[tid] = id;330 batch_id[tid] = id;
303 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));331 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
304 } else {332 } else {
305 batch_mask[tid] = 0u;333 batch_mask[tid] = 0u;
306 }334 }
307 __syncthreads();335 __syncthreads();
308 const int count = min(kRasterThreads, end);336 const int count = min(kRasterThreads, end);
309#pragma unroll337#pragma unroll
310 for (int c = 0; c < kRasterThreads; c += 32) {338 for (int c = 0; c < kRasterThreads; c += 32) {
311 if (c < count && end - 32 - c < warp_max_last) {339 if (c < count && end - 32 - c < warp_max_last) {
312 uint32_t mask = __ballot_sync(340 uint32_t mask = __ballot_sync(
313 warp::kFullMask,341 warp::kFullMask,
314 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);342 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
315 while ((mask & (mask - 1u)) != 0u) {343 while ((mask & (mask - 1u)) != 0u) {
316 const int bit0 = __ffs(mask) - 1;344 const int bit0 = __ffs(mask) - 1;
317 mask &= mask - 1u;345 mask &= mask - 1u;
318 const int bit1 = __ffs(mask) - 1;346 const int bit1 = __ffs(mask) - 1;
319 mask &= mask - 1u;347 mask &= mask - 1u;
320 const int j0 = c + bit0, j1 = c + bit1;348 const int j0 = c + bit0, j1 = c + bit1;
321 const float4 g0_0 = batch_geom0[j0];349 const float4 g0_0 = batch_geom0[j0];
322 const float2 g1_0 = batch_geom1[j0];350 const float2 g1_0 = batch_geom1[j0];
323 const float4 ro_0 = batch_rgb_o[j0];351 const float4 ro_0 = batch_rgb_o[j0];
324 const float4 g0_1 = batch_geom0[j1];352 const float4 g0_1 = batch_geom0[j1];
325 const float2 g1_1 = batch_geom1[j1];353 const float2 g1_1 = batch_geom1[j1];
326 const float4 ro_1 = batch_rgb_o[j1];354 const float4 ro_1 = batch_rgb_o[j1];
327355
328 float dx0, ga_sum0, ga_dy0, ga_dy2_0;356 float dx0, ga_sum0, ga_dy0, ga_dy2_0;
329 float3 drgb0;357 float3 drgb0;
330 const bool hit0 = eval_splat(g0_0, g1_0, ro_0, end - 1 - j0, dx0,358 const bool hit0 = eval_splat(g0_0, g1_0, ro_0, end - 1 - j0, dx0,
331 ga_sum0, ga_dy0, ga_dy2_0, drgb0);359 ga_sum0, ga_dy0, ga_dy2_0, drgb0);
332360
333 float dx1, ga_sum1, ga_dy1, ga_dy2_1;361 float dx1, ga_sum1, ga_dy1, ga_dy2_1;
334 float3 drgb1;362 float3 drgb1;
335 const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,363 const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,
336 ga_sum1, ga_dy1, ga_dy2_1, drgb1);364 ga_sum1, ga_dy1, ga_dy2_1, drgb1);
337 // As in s01: a pair that no pixel of the warp blends has zero365 // As in s01: a pair that no pixel of the warp blends has zero
338 // gradients, so skip its reduction and atomics entirely.366 // gradients, so skip its reduction and atomics entirely.
339 if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;367 if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;
340368
341 const bool hi16 = (lane & 16) != 0;369 const bool hi16 = (lane & 16) != 0;
342 const float opacity = hi16 ? ro_1.w : ro_0.w;370 const float opacity = hi16 ? ro_1.w : ro_0.w;
343 const float dx = hi16 ? dx1 : dx0;371 const float dx = hi16 ? dx1 : dx0;
344 const float3 cn =372 const float3 cn =
345 hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);373 hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);
346374
347 auto fold16 = [&](float v0, float v1) {375 auto fold16 = [&](float v0, float v1) {
348 const float send = hi16 ? v0 : v1;376 const float send = hi16 ? v0 : v1;
349 const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);377 const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);
350 return (hi16 ? v1 : v0) + recv;378 return (hi16 ? v1 : v0) + recv;
351 };379 };
352 float ga_sum = fold16(ga_sum0, ga_sum1);380 float ga_sum = fold16(ga_sum0, ga_sum1);
353 float ga_dy = fold16(ga_dy0, ga_dy1);381 float ga_dy = fold16(ga_dy0, ga_dy1);
354 float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),382 float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),
355 fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),383 fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),
356 fold16(drgb0.z, drgb1.z)};384 fold16(drgb0.z, drgb1.z)};
357 warp::internal::HalvingStep<4, 8>(b4, lane);385 warp::internal::HalvingStep<4, 8>(b4, lane);
358 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);386 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
359 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);387 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
360388
361 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;389 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
362 const float p_dx = dx * p_sum;390 const float p_dx = dx * p_sum;
363 const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)391 const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
364 : -(cn.y * p_dx + cn.z * p_dy);392 : -(cn.y * p_dx + cn.z * p_dy);
365 const float v_dx1 = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;393 const float v_dx1 = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
366 float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};394 float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};
367 warp::internal::HalvingStep<4, 4>(v4, lane);395 warp::internal::HalvingStep<4, 4>(v4, lane);
368 warp::internal::HalvingStep<2, 2>(v4, lane);396 warp::internal::HalvingStep<2, 2>(v4, lane);
369 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);397 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);
370 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);398 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);
371 const float v_sum =399 const float v_sum =
372 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);400 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);
373 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);401 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);
374 const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;402 const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
375 if (scatter2 >= 0 && out_grad != 0.0f) {403 if (scatter2 >= 0 && out_grad != 0.0f) {
376 const uint32_t target_id = batch_id[hi16 ? j1 : j0];404 const uint32_t target_id = batch_id[hi16 ? j1 : j0];
377 atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,405 atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,
378 out_grad);406 out_grad);
379 }407 }
380 }408 }
381 while (mask != 0u) {409 while (mask != 0u) {
382 const int bit = __ffs(mask) - 1;410 const int bit = __ffs(mask) - 1;
383 mask &= mask - 1u;411 mask &= mask - 1u;
384 const int j = c + bit;412 const int j = c + bit;
385 const float4 g0 = batch_geom0[j];413 const float4 g0 = batch_geom0[j];
386 const float2 g1 = batch_geom1[j];414 const float2 g1 = batch_geom1[j];
387 const float4 ro = batch_rgb_o[j];415 const float4 ro = batch_rgb_o[j];
388 float dx, ga_sum, ga_dy, ga_dy2;416 float dx, ga_sum, ga_dy, ga_dy2;
389 float3 drgb;417 float3 drgb;
390 const bool hit =418 const bool hit =
391 eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);419 eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
392 if (!__any_sync(warp::kFullMask, hit)) continue;420 if (!__any_sync(warp::kFullMask, hit)) continue;
393 const float opacity = ro.w;421 const float opacity = ro.w;
394 const float3 cn = raw_conic(g0, g1);422 const float3 cn = raw_conic(g0, g1);
395 float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};423 float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};
396 warp::internal::HalvingStep<4, 16>(b4, lane);424 warp::internal::HalvingStep<4, 16>(b4, lane);
397 warp::internal::HalvingStep<2, 8>(b4, lane);425 warp::internal::HalvingStep<2, 8>(b4, lane);
398 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);426 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);
399 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);427 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);
400 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);428 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
401 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);429 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
402 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;430 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
403 const float p_dx = dx * p_sum;431 const float p_dx = dx * p_sum;
404 float v_dx;432 float v_dx;
405 if ((lane & 16) == 0) {433 if ((lane & 16) == 0) {
406 v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)434 v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
407 : -(cn.y * p_dx + cn.z * p_dy);435 : -(cn.y * p_dx + cn.z * p_dy);
408 } else {436 } else {
409 v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;437 v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
410 }438 }
411 float v[3] = {v_dx, b4[0], ga_sum};439 float v[3] = {v_dx, b4[0], ga_sum};
412 warp::internal::HalvingStep<3, 4>(v, lane);440 warp::internal::HalvingStep<3, 4>(v, lane);
413 warp::internal::HalvingStep<2, 2>(v, lane);441 warp::internal::HalvingStep<2, 2>(v, lane);
414 const float out_grad =442 const float out_grad =
415 v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);443 v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
416 if (scatter >= 0 && out_grad != 0.0f) {444 if (scatter >= 0 && out_grad != 0.0f) {
417 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,445 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
418 out_grad);446 out_grad);
419 }447 }
420 }448 }
421 }449 }
422 }450 }
423 __syncthreads();451 __syncthreads();
424 }452 }
425}453}
426454
427// ---------------------------------------------------------------------------455// ---------------------------------------------------------------------------
428// Preprocess backward: splat gradients -> parameter gradients.456// Preprocess backward: splat gradients -> parameter gradients (dense).
429// Parameter gradient buffers are bulk-zeroed beforehand via `cudaMemsetAsync`;
430// culled or zero-gradient Gaussians return early.
457// Gaussians stored as culled (radius 0, i.e. no tile pairs) are untouched.
458// Uses warp-coalesced 16-byte stores (0 shared memory) to zero untouched SH
459// rows and streams active SH rows in 4 groups of 12 floats via `ShIo<ShT>`.
431460
432template <typename ShT>461template <typename ShT>
433__global__ void __launch_bounds__(kGaussianThreads)462__global__ void __launch_bounds__(kGaussianThreads)
434 PreprocessBackwardKernel(GaussianParams params, Frame frame,463 PreprocessBackwardKernel(GaussianParams params, Frame frame,
435 const Splat* __restrict__ splats,464 const Splat* __restrict__ splats,
436 const float* __restrict__ splat_grads,465 const float* __restrict__ splat_grads,
437 GaussianGrads grads) {466 GaussianGrads grads) {
467 constexpr int kShVec16 = ShIo<ShT>::kRowVec16; // 16-byte stores per row.
438 const int i = blockIdx.x * blockDim.x + threadIdx.x;468 const int i = blockIdx.x * blockDim.x + threadIdx.x;
469 const int lane = threadIdx.x % 32;
439 if (i >= params.n || splats[i].radius() == 0) return;470 bool touched = (i < params.n && splats[i].radius() != 0);
471 const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);
472 if (init_mask == 0u) {
473 if (i < params.n) {
474 grads.means[3 * i] = 0.0f;
475 grads.means[3 * i + 1] = 0.0f;
476 grads.means[3 * i + 2] = 0.0f;
477 grads.scales[3 * i] = 0.0f;
478 grads.scales[3 * i + 1] = 0.0f;
479 grads.scales[3 * i + 2] = 0.0f;
480 reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};
481 grads.opacities[i] = 0.0f;
482 }
483 const int warp_i0 =
484 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
485 const int warp_n = max(0, min(32, params.n - warp_i0));
486 uint4* warp_sh_dst =
487 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
488 const uint4 zero = {0u, 0u, 0u, 0u};
489 if (warp_n == 32) {
490#pragma unroll
491 for (int s = 0; s < kShVec16; ++s) {
492 warp_sh_dst[s * 32 + lane] = zero;
493 }
494 } else {
495 for (int k = lane; k < warp_n * kShVec16; k += 32) {
496 warp_sh_dst[k] = zero;
497 }
498 }
499 return;
500 }
501
502 float3 dmean = {0.0f, 0.0f, 0.0f};
503 float3 dscale = {0.0f, 0.0f, 0.0f};
504 float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};
505 float dopacity = 0.0f;
506 if (touched) {
440 float g[kSplatGradFloats];507 float g[kSplatGradFloats];
441 bool nonzero = false;508 bool nonzero = false;
442#pragma unroll509#pragma unroll
443 for (int q = 0; q < kSplatGradFloats; ++q) {510 for (int q = 0; q < kSplatGradFloats; ++q) {
444 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);511 g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
445 nonzero = nonzero || g[q] != 0.0f;512 nonzero = nonzero || g[q] != 0.0f;
446 }513 }
447 if (!nonzero) return;514 touched = nonzero;
515 if (touched) {
448 ProjectionGrad pg;516 ProjectionGrad pg;
449 pg.xy = {g[0], g[1]};517 pg.xy = {g[0], g[1]};
450 pg.conic = {g[2], g[3], g[4]};518 pg.conic = {g[2], g[3], g[4]};
451 pg.opacity = g[5];519 pg.opacity = g[5];
452 pg.rgb = {g[6], g[7], g[8]};520 pg.rgb = {g[6], g[7], g[8]};
453 float3 dmean, dscale;
454 float4 dquat;
521 dopacity = g[5];
455 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),522 ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
456 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);523 ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
524 }
525 }
526 if (i < params.n) {
457 grads.means[3 * i] = dmean.x;527 grads.means[3 * i] = dmean.x;
458 grads.means[3 * i + 1] = dmean.y;528 grads.means[3 * i + 1] = dmean.y;
459 grads.means[3 * i + 2] = dmean.z;529 grads.means[3 * i + 2] = dmean.z;
460 grads.scales[3 * i] = dscale.x;530 grads.scales[3 * i] = dscale.x;
461 grads.scales[3 * i + 1] = dscale.y;531 grads.scales[3 * i + 1] = dscale.y;
462 grads.scales[3 * i + 2] = dscale.z;532 grads.scales[3 * i + 2] = dscale.z;
463 reinterpret_cast<float4*>(grads.quats)[i] = dquat;533 reinterpret_cast<float4*>(grads.quats)[i] = dquat;
464 grads.opacities[i] = g[5];534 grads.opacities[i] = dopacity;
535 }
536 const uint32_t mask = __ballot_sync(warp::kFullMask, touched);
537 if (mask != warp::kFullMask) {
538 const int warp_i0 =
539 static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
540 const int warp_n = max(0, min(32, params.n - warp_i0));
541 uint4* warp_sh_dst =
542 reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
543 const uint4 zero = {0u, 0u, 0u, 0u};
544 if (warp_n == 32) {
545#pragma unroll
546 for (int s = 0; s < kShVec16; ++s) {
547 const int k = s * 32 + lane;
548 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
549 if (((mask >> row_in_warp) & 1u) == 0u) {
550 warp_sh_dst[k] = zero;
551 }
552 }
553 } else {
554 for (int k = lane; k < warp_n * kShVec16; k += 32) {
555 const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
556 if (((mask >> row_in_warp) & 1u) == 0u) {
557 warp_sh_dst[k] = zero;
558 }
559 }
560 }
561 }
465}562}
466563
467template <typename ShT>564template <typename ShT>
468cudaError_t LaunchPreprocessBackward(const GaussianParams& params,565cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
469 const Frame& frame, const Splat* splats,566 const Frame& frame, const Splat* splats,
470 const float* splat_grads,567 const float* splat_grads,
471 const GaussianGrads& grads,568 const GaussianGrads& grads,
472 cudaStream_t stream) {569 cudaStream_t stream) {
473 const int n = params.n;
474 // As Inria's PyTorch wrapper (`torch::zeros`): bulk-zero every parameter
475 // gradient up front; culled Gaussians then exit early in the kernel.
476 if (auto e = cudaMemsetAsync(grads.means, 0, sizeof(float) * 3 * n, stream);
477 e != cudaSuccess)
478 return e;
479 if (auto e = cudaMemsetAsync(grads.scales, 0, sizeof(float) * 3 * n, stream);
480 e != cudaSuccess)
481 return e;
482 if (auto e = cudaMemsetAsync(grads.quats, 0, sizeof(float) * 4 * n, stream);
483 e != cudaSuccess)
484 return e;
485 if (auto e = cudaMemsetAsync(grads.opacities, 0, sizeof(float) * n, stream);
486 e != cudaSuccess)
487 return e;
488 if (auto e =
489 cudaMemsetAsync(grads.sh, 0, sizeof(ShT) * 3 * kShCoeffs * n, stream);
490 e != cudaSuccess)
491 return e;
492 PreprocessBackwardKernel<ShT>570 PreprocessBackwardKernel<ShT>
493 <<<GaussianBlocks(n), kGaussianThreads, 0, stream>>>(571 <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(
494 params, frame, splats, splat_grads, grads);572 params, frame, splats, splat_grads, grads);
495 return cudaSuccess;573 return cudaSuccess;
496}574}
497575
↕ Expand 16 unchanged lines (lines 576–591) · in PreprocessBackwardKernel
498template __global__ void PreprocessBackwardKernel<float>(576template __global__ void PreprocessBackwardKernel<float>(
499 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,577 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
500 GaussianGrads);578 GaussianGrads);
501template __global__ void PreprocessBackwardKernel<ShBf16>(579template __global__ void PreprocessBackwardKernel<ShBf16>(
502 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,580 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
503 GaussianGrads);581 GaussianGrads);
504template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,582template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
505 const Frame&, const Splat*,583 const Frame&, const Splat*,
506 const float*,584 const float*,
507 const GaussianGrads&,585 const GaussianGrads&,
508 cudaStream_t);586 cudaStream_t);
509template cudaError_t LaunchPreprocessBackward<ShBf16>(587template cudaError_t LaunchPreprocessBackward<ShBf16>(
510 const GaussianParams&, const Frame&, const Splat*, const float*,588 const GaussianParams&, const Frame&, const Splat*, const float*,
511 const GaussianGrads&, cudaStream_t);589 const GaussianGrads&, cudaStream_t);
512590
513} // namespace gs591} // namespace gs
↕ 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 #include "backward.h"
1313
1414 #include "auxiliary.h"
1515
1616 namespace gs {
1717
1818 namespace sh {
19−// Evaluates SH color VJP loading the full 16-basis row `sh_in` and storing
19+// Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly
20−// `dsh` to `sh_out`.
20+// 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);
26− float3 coeffs[kShCoeffs];
27− ShIo<ShT>::LoadRow(sh_in, coeffs);
28− float3 raw = {0.5f, 0.5f, 0.5f};
26+ const float3 raw = EvalRow(sh_in, b);
29−#pragma unroll
30− for (int k = 0; k < kShCoeffs; ++k) {
31− raw += b[k] * coeffs[k];
32− }
3327 const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
3428 raw.y < 0.0f ? 0.0f : g_rgb.y,
3529 raw.z < 0.0f ? 0.0f : g_rgb.z};
36− float3 bgrad[kShCoeffs];
37− BasisGrad(d, bgrad);
38− float3 dsh[kShCoeffs];
30+ const float x = d.x, y = d.y, z = d.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;
3933 float3 ddir = {0.0f, 0.0f, 0.0f};
4034 #pragma unroll
35+ for (int grp = 0; grp < 4; ++grp) {
36+ const int k0 = 4 * grp;
37+ float v[12];
38+ ShIo<ShT>::LoadGroup(sh_in, grp, v);
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;
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;
43+ float out[12];
44+#pragma unroll
41− for (int k = 0; k < kShCoeffs; ++k) {
45+ for (int q = 0; q < 4; ++q) {
42− dsh[k] = b[k] * drgb;
46+ out[3 * q] = b[k0 + q] * drgb.x;
43− const float s = dot(coeffs[k], drgb);
47+ out[3 * q + 1] = b[k0 + q] * drgb.y;
44− ddir += s * bgrad[k];
48+ out[3 * q + 2] = b[k0 + q] * drgb.z;
4549⇥ }
46− ShIo<ShT>::StoreRow(sh_out, dsh);
50+ ShIo<ShT>::StoreGroup(sh_out, grp, out);
51+ if (grp == 0) {
52+ ddir.y += s1 * (-kC1);
53+ ddir.z += s2 * kC1;
54+ ddir.x += s3 * (-kC1);
55+ } else if (grp == 1) {
56+ ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};
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};
59+ ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};
60+ } else if (grp == 2) {
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};
63+ ddir += (s2 * kC3[1]) * float3{yz, xz, xy};
64+ ddir += (s3 * kC3[2]) *
65+ float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};
66+ } else {
67+ ddir += (s0 * kC3[3]) *
68+ float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};
69+ ddir += (s1 * kC3[4]) *
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};
72+ ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};
73+ }
74+ }
4775 return ddir;
4876 }
4977 } // namespace sh
5078
↕ Expand 373 unchanged lines (lines 79–451) · in RasterBackwardKernel
5179 // VJP of the projection for a touched Gaussian `i`; reads its SH row from
5280 // `sh_in` and writes dL/dsh directly to `sh_out`.
5381 // Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
5482 template <typename ShT>
5583 GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
5684 const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
5785 const ProjectionGrad& g, float3& dmean_out,
5886 float3& dscale_out, float4& dquat_out) {
5987 const float3 mean = LoadMean(p, i);
6088 const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
6189 p.scales[3 * i + 2]};
6290 const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
6391
6492 // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
6593 float3 dmean = {0.0f, 0.0f, 0.0f};
6694 {
6795 float inv_len;
6896 const float3 dir = ViewDir(cam, mean, inv_len);
6997 const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
7098 // d(v/|v|)/dv = (I - dir dir^T) / |v|.
7199 dmean += inv_len * (ddir - dot(dir, ddir) * dir);
72100 }
73101
74102 // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
75103 const float3 t = cam.ToView(mean);
76104 const float tz = t.z;
77105 const float limx = kFovClamp * cam.TanHalfFovX();
78106 const float limy = kFovClamp * cam.TanHalfFovY();
79107 const float txtz = t.x / tz, tytz = t.y / tz;
80108 const bool x_clamped = txtz < -limx || txtz > limx;
81109 const bool y_clamped = tytz < -limy || tytz > limy;
82110 const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
83111 const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
84112 const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
85113 const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
86114 const Mat3& w = cam.rotation;
87115 const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
88116 const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
89117
90118 const Mat3 r = QuatToRotation(quat);
91119 const Sym3 sigma = Covariance3D(r, scale);
92120 const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
93121 const float cov_a = dot(m0, s0) + kLowPass;
94122 const float cov_b = dot(m0, s1);
95123 const float cov_c = dot(m1, s1) + kLowPass;
96124
97125 // --- Conic -> 2D covariance: G_cov = -K G_K K.
98126 const float det = cov_a * cov_c - cov_b * cov_b;
99127 const float inv = 1.0f / det;
100128 const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
101129 const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
102130 // P = G_K K.
103131 const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
104132 const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
105133 // G_cov = -K P (symmetric).
106134 const float ca = -(ka * p00 + kb * p10);
107135 const float cb = -(ka * p01 + kb * p11);
108136 const float cc = -(kb * p01 + kc * p11);
109137
110138 // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
111139 const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
112140 const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
113141 // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
114142 Sym3 dsig;
115143 dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
116144 dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
117145 dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
118146 dsig.xy =
119147 ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
120148 dsig.xz =
121149 ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
122150 dsig.yz =
123151 ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
124152
125153 // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
126154 const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
127155 const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
128156 const float tz2 = tz * tz, tz3 = tz2 * tz;
129157 float3 dt;
130158 dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
131159 dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
132160 dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
133161 2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
134162 // --- Screen-space mean u = fx x / z + cx.
135163 dt.x += cam.fx / tz * g.xy.x;
136164 dt.y += cam.fy / tz * g.xy.y;
137165 dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
138166 dmean += MulT(w, dt);
139167 dmean_out = dmean;
140168
141169 // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
142170 const float sv[3] = {scale.x, scale.y, scale.z};
143171 float ds[3];
144172 Mat3 dr;
145173 #pragma unroll
146174 for (int j = 0; j < 3; ++j) {
147175 const float3 r_col = r.col(j);
148176 const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
149177 ds[j] = dot(da_col, r_col);
150178 dr.m[0][j] = da_col.x * sv[j];
151179 dr.m[1][j] = da_col.y * sv[j];
152180 dr.m[2][j] = da_col.z * sv[j];
153181 }
154182 dscale_out = {ds[0], ds[1], ds[2]};
155183
156184 // --- R(q) -> q.
157185 const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
158186 const auto& G = dr.m;
159187 dquat_out = {
160188 2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
161189 qy * G[2][0] + qx * G[2][1]),
162190 2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
163191 qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
164192 2.0f *
165193 (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
166194 qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
167195 2.0f *
168196 (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
169197 2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
170198 }
171199
172200 // Reverse traversal per pixel; accumulates splat gradients via Half-Warp
173201 // Two-Splat Folding and Halving Register Reduce-Scatter.
174202 __global__ void __launch_bounds__(RasterLayout<1>::kThreads, 4)
175203 RasterBackwardKernel(const int2* __restrict__ ranges,
176204 const int32_t* __restrict__ point_list,
177205 const Splat* __restrict__ splats, Frame frame,
178206 const float* __restrict__ final_t,
179207 const int32_t* __restrict__ n_contrib,
180208 const float* __restrict__ dl_dimage,
181209 float* __restrict__ splat_grads) {
182210 using L = RasterLayout<1>;
183211 constexpr int kRasterThreads = L::kThreads;
184212 constexpr int kWarps = L::kWarps;
185213 const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
186214 const int tid = threadIdx.x;
187215 const int warp = tid / 32;
188216 const int lane = tid % 32;
189217 const RasterPixels px = RasterPixels::For<1>(tid);
190218 const float fx = static_cast<float>(px.x);
191219 const float fy = static_cast<float>(px.y0);
192220 const int2 range = ranges[tile];
193221
194222 const bool inside = px.x < frame.width && px.y0 < frame.height;
195223 const int pixels = frame.width * frame.height;
196224 const int p = px.y0 * frame.width + px.x;
197225 float t = inside ? final_t[p] : 0.0f;
198226 const float3 dpix =
199227 inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
200228 dl_dimage[2 * pixels + p]}
201229 : float3{0.0f, 0.0f, 0.0f};
202230 float accum = 0.0f; // dot(color accumulated behind, dpix)
203231 const int last = inside ? n_contrib[p] : 0;
204232
205233 int warp_max_last = last;
206234 #pragma unroll
207235 for (int offset = 16; offset > 0; offset /= 2) {
208236 warp_max_last = max(
209237 warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
210238 }
211239 __shared__ int shared_max_last[kWarps];
212240 if (lane == 0) shared_max_last[warp] = warp_max_last;
213241 __syncthreads();
214242 int block_max_last = 0;
215243 #pragma unroll
216244 for (int w = 0; w < kWarps; ++w) {
217245 block_max_last = max(block_max_last, shared_max_last[w]);
218246 }
219247 if (block_max_last <= 0) return;
220248
221249 __shared__ float4 batch_geom0[kRasterThreads];
222250 __shared__ float2 batch_geom1[kRasterThreads];
223251 __shared__ float4 batch_rgb_o[kRasterThreads];
224252 __shared__ uint32_t batch_id[kRasterThreads];
225253 __shared__ uint8_t batch_mask[kRasterThreads];
226254 const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
227255 int scatter = -1;
228256 int scatter2 = -1;
229257 {
230258 const int sub_lane = lane & 15;
231259 const int row = sub_lane >> 3, col = sub_lane & 7;
232260 if (col == 0) {
233261 scatter2 = row;
234262 } else if (col == 2) {
235263 scatter2 = 2 + row;
236264 } else if (col == 4) {
237265 scatter2 = row == 0 ? 4 : 7;
238266 } else if (col == 6) {
239267 scatter2 = row == 0 ? 6 : 8;
240268 } else if (col == 1 && row == 0) {
241269 scatter2 = 5;
242270 }
243271 const int full_row = lane >> 3, full_col = lane & 7;
244272 if (full_col == 0) {
245273 scatter = full_row;
246274 } else if (full_col == 2) {
247275 scatter = full_row == 0 ? 4 : 5 + full_row;
248276 } else if (full_col == 4 && full_row == 0) {
249277 scatter = 5;
250278 }
251279 }
252280
253281 auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,
254282 float& ga_sum, float& ga_dy, float& ga_dy2,
255283 float3& drgb) -> bool {
256284 dx = g0.x - fx;
257285 ga_sum = 0.0f;
258286 ga_dy = 0.0f;
259287 ga_dy2 = 0.0f;
260288 drgb = {0.0f, 0.0f, 0.0f};
261289 if (entry < last) {
262290 const float dy = g0.y - fy;
263291 const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
264292 if (power <= 0.0f && power >= g1.y) {
265293 const float g_val = Exp2Approx(power);
266294 const float alpha = fminf(kMaxAlpha, ro.w * g_val);
267295 if (alpha >= kMinAlpha) {
268296 t = __fdividef(t, 1.0f - alpha);
269297 const float w = alpha * t;
270298 const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
271299 const float diff = color_dot - accum;
272300 const float dl_dalpha = t * diff;
273301 accum += alpha * diff;
274302 const float ga = g_val * dl_dalpha;
275303 ga_sum = ga;
276304 const float ga_y = ga * dy;
277305 ga_dy = ga_y;
278306 ga_dy2 = ga_y * dy;
279307 drgb = w * dpix;
280308 return true;
281309 }
282310 }
283311 }
284312 return false;
285313 };
286314 auto raw_conic = [](float4 g0, float2 g1) {
287315 return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
288316 kInvPower2Scale * g1.x};
289317 };
290318
291319 for (int end = block_max_last; end > 0; end -= kRasterThreads) {
292320 const int e = end - 1 - tid;
293321 if (e >= 0) {
294322 const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
295323 const float4 p0 = __ldg(&splats[id].p0);
296324 const float4 p1 = __ldg(&splats[id].p1);
297325 const float4 p2 = __ldg(&splats[id].p2);
298326 const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
299327 batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
300328 batch_geom1[tid] = {c2.z, p2.w};
301329 batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
302330 batch_id[tid] = id;
303331 batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
304332 } else {
305333 batch_mask[tid] = 0u;
306334 }
307335 __syncthreads();
308336 const int count = min(kRasterThreads, end);
309337 #pragma unroll
310338 for (int c = 0; c < kRasterThreads; c += 32) {
311339 if (c < count && end - 32 - c < warp_max_last) {
312340 uint32_t mask = __ballot_sync(
313341 warp::kFullMask,
314342 c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
315343 while ((mask & (mask - 1u)) != 0u) {
316344 const int bit0 = __ffs(mask) - 1;
317345 mask &= mask - 1u;
318346 const int bit1 = __ffs(mask) - 1;
319347 mask &= mask - 1u;
320348 const int j0 = c + bit0, j1 = c + bit1;
321349 const float4 g0_0 = batch_geom0[j0];
322350 const float2 g1_0 = batch_geom1[j0];
323351 const float4 ro_0 = batch_rgb_o[j0];
324352 const float4 g0_1 = batch_geom0[j1];
325353 const float2 g1_1 = batch_geom1[j1];
326354 const float4 ro_1 = batch_rgb_o[j1];
327355
328356 float dx0, ga_sum0, ga_dy0, ga_dy2_0;
329357 float3 drgb0;
330358 const bool hit0 = eval_splat(g0_0, g1_0, ro_0, end - 1 - j0, dx0,
331359 ga_sum0, ga_dy0, ga_dy2_0, drgb0);
332360
333361 float dx1, ga_sum1, ga_dy1, ga_dy2_1;
334362 float3 drgb1;
335363 const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,
336364 ga_sum1, ga_dy1, ga_dy2_1, drgb1);
337365 // As in s01: a pair that no pixel of the warp blends has zero
338366 // gradients, so skip its reduction and atomics entirely.
339367 if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;
340368
341369 const bool hi16 = (lane & 16) != 0;
342370 const float opacity = hi16 ? ro_1.w : ro_0.w;
343371 const float dx = hi16 ? dx1 : dx0;
344372 const float3 cn =
345373 hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);
346374
347375 auto fold16 = [&](float v0, float v1) {
348376 const float send = hi16 ? v0 : v1;
349377 const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);
350378 return (hi16 ? v1 : v0) + recv;
351379 };
352380 float ga_sum = fold16(ga_sum0, ga_sum1);
353381 float ga_dy = fold16(ga_dy0, ga_dy1);
354382 float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),
355383 fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),
356384 fold16(drgb0.z, drgb1.z)};
357385 warp::internal::HalvingStep<4, 8>(b4, lane);
358386 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
359387 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
360388
361389 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
362390 const float p_dx = dx * p_sum;
363391 const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
364392 : -(cn.y * p_dx + cn.z * p_dy);
365393 const float v_dx1 = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
366394 float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};
367395 warp::internal::HalvingStep<4, 4>(v4, lane);
368396 warp::internal::HalvingStep<2, 2>(v4, lane);
369397 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);
370398 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);
371399 const float v_sum =
372400 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);
373401 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);
374402 const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
375403 if (scatter2 >= 0 && out_grad != 0.0f) {
376404 const uint32_t target_id = batch_id[hi16 ? j1 : j0];
377405 atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,
378406 out_grad);
379407 }
380408 }
381409 while (mask != 0u) {
382410 const int bit = __ffs(mask) - 1;
383411 mask &= mask - 1u;
384412 const int j = c + bit;
385413 const float4 g0 = batch_geom0[j];
386414 const float2 g1 = batch_geom1[j];
387415 const float4 ro = batch_rgb_o[j];
388416 float dx, ga_sum, ga_dy, ga_dy2;
389417 float3 drgb;
390418 const bool hit =
391419 eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
392420 if (!__any_sync(warp::kFullMask, hit)) continue;
393421 const float opacity = ro.w;
394422 const float3 cn = raw_conic(g0, g1);
395423 float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};
396424 warp::internal::HalvingStep<4, 16>(b4, lane);
397425 warp::internal::HalvingStep<2, 8>(b4, lane);
398426 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);
399427 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);
400428 ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
401429 ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
402430 const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
403431 const float p_dx = dx * p_sum;
404432 float v_dx;
405433 if ((lane & 16) == 0) {
406434 v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
407435 : -(cn.y * p_dx + cn.z * p_dy);
408436 } else {
409437 v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
410438 }
411439 float v[3] = {v_dx, b4[0], ga_sum};
412440 warp::internal::HalvingStep<3, 4>(v, lane);
413441 warp::internal::HalvingStep<2, 2>(v, lane);
414442 const float out_grad =
415443 v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
416444 if (scatter >= 0 && out_grad != 0.0f) {
417445 atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
418446 out_grad);
419447 }
420448 }
421449 }
422450 }
423451 __syncthreads();
424452 }
425453 }
426454
427455 // ---------------------------------------------------------------------------
428−// Preprocess backward: splat gradients -> parameter gradients.
456+// Preprocess backward: splat gradients -> parameter gradients (dense).
429−// Parameter gradient buffers are bulk-zeroed beforehand via `cudaMemsetAsync`;
430−// culled or zero-gradient Gaussians return early.
457+// Gaussians stored as culled (radius 0, i.e. no tile pairs) are untouched.
458+// Uses warp-coalesced 16-byte stores (0 shared memory) to zero untouched SH
459+// rows and streams active SH rows in 4 groups of 12 floats via `ShIo<ShT>`.
431460
432461 template <typename ShT>
433462 __global__ void __launch_bounds__(kGaussianThreads)
434463 PreprocessBackwardKernel(GaussianParams params, Frame frame,
435464 const Splat* __restrict__ splats,
436465 const float* __restrict__ splat_grads,
437466 GaussianGrads grads) {
467+ constexpr int kShVec16 = ShIo<ShT>::kRowVec16; // 16-byte stores per row.
438468 const int i = blockIdx.x * blockDim.x + threadIdx.x;
469+ const int lane = threadIdx.x % 32;
439− if (i >= params.n || splats[i].radius() == 0) return;
470+ bool touched = (i < params.n && splats[i].radius() != 0);
471+ const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);
472+ if (init_mask == 0u) {
473+ if (i < params.n) {
474+ grads.means[3 * i] = 0.0f;
475+ grads.means[3 * i + 1] = 0.0f;
476+ grads.means[3 * i + 2] = 0.0f;
477+ grads.scales[3 * i] = 0.0f;
478+ grads.scales[3 * i + 1] = 0.0f;
479+ grads.scales[3 * i + 2] = 0.0f;
480+ reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};
481+ grads.opacities[i] = 0.0f;
482+ }
483+ const int warp_i0 =
484+ static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
485+ const int warp_n = max(0, min(32, params.n - warp_i0));
486+ uint4* warp_sh_dst =
487+ reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
488+ const uint4 zero = {0u, 0u, 0u, 0u};
489+ if (warp_n == 32) {
490+#pragma unroll
491+ for (int s = 0; s < kShVec16; ++s) {
492+ warp_sh_dst[s * 32 + lane] = zero;
493+ }
494+ } else {
495+ for (int k = lane; k < warp_n * kShVec16; k += 32) {
496+ warp_sh_dst[k] = zero;
497+ }
498+ }
499+ return;
500+ }
501+
502+ float3 dmean = {0.0f, 0.0f, 0.0f};
503+ float3 dscale = {0.0f, 0.0f, 0.0f};
504+ float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};
505+ float dopacity = 0.0f;
506+ if (touched) {
440507⇥ float g[kSplatGradFloats];
441508⇥ bool nonzero = false;
442509 #pragma unroll
443510⇥ for (int q = 0; q < kSplatGradFloats; ++q) {
444511⇥ g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
445512⇥ nonzero = nonzero || g[q] != 0.0f;
446513⇥ }
447− if (!nonzero) return;
514+ touched = nonzero;
515+ if (touched) {
448516⇥ ProjectionGrad pg;
449517⇥ pg.xy = {g[0], g[1]};
450518⇥ pg.conic = {g[2], g[3], g[4]};
451519⇥ pg.opacity = g[5];
452520⇥ pg.rgb = {g[6], g[7], g[8]};
453− float3 dmean, dscale;
454− float4 dquat;
521+ dopacity = g[5];
455522⇥ ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
456523⇥ ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
524+ }
525+ }
526+ if (i < params.n) {
457527⇥ grads.means[3 * i] = dmean.x;
458528⇥ grads.means[3 * i + 1] = dmean.y;
459529⇥ grads.means[3 * i + 2] = dmean.z;
460530⇥ grads.scales[3 * i] = dscale.x;
461531⇥ grads.scales[3 * i + 1] = dscale.y;
462532⇥ grads.scales[3 * i + 2] = dscale.z;
463533⇥ reinterpret_cast<float4*>(grads.quats)[i] = dquat;
464− grads.opacities[i] = g[5];
534+ grads.opacities[i] = dopacity;
535+ }
536+ const uint32_t mask = __ballot_sync(warp::kFullMask, touched);
537+ if (mask != warp::kFullMask) {
538+ const int warp_i0 =
539+ static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
540+ const int warp_n = max(0, min(32, params.n - warp_i0));
541+ uint4* warp_sh_dst =
542+ reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
543+ const uint4 zero = {0u, 0u, 0u, 0u};
544+ if (warp_n == 32) {
545+#pragma unroll
546+ for (int s = 0; s < kShVec16; ++s) {
547+ const int k = s * 32 + lane;
548+ const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
549+ if (((mask >> row_in_warp) & 1u) == 0u) {
550+ warp_sh_dst[k] = zero;
551+ }
552+ }
553+ } else {
554+ for (int k = lane; k < warp_n * kShVec16; k += 32) {
555+ const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
556+ if (((mask >> row_in_warp) & 1u) == 0u) {
557+ warp_sh_dst[k] = zero;
558+ }
559+ }
560+ }
561+ }
465562 }
466563
467564 template <typename ShT>
468565 cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
469566 const Frame& frame, const Splat* splats,
470567 const float* splat_grads,
471568 const GaussianGrads& grads,
472569 cudaStream_t stream) {
473− const int n = params.n;
474− // As Inria's PyTorch wrapper (`torch::zeros`): bulk-zero every parameter
475− // gradient up front; culled Gaussians then exit early in the kernel.
476− if (auto e = cudaMemsetAsync(grads.means, 0, sizeof(float) * 3 * n, stream);
477− e != cudaSuccess)
478− return e;
479− if (auto e = cudaMemsetAsync(grads.scales, 0, sizeof(float) * 3 * n, stream);
480− e != cudaSuccess)
481− return e;
482− if (auto e = cudaMemsetAsync(grads.quats, 0, sizeof(float) * 4 * n, stream);
483− e != cudaSuccess)
484− return e;
485− if (auto e = cudaMemsetAsync(grads.opacities, 0, sizeof(float) * n, stream);
486− e != cudaSuccess)
487− return e;
488− if (auto e =
489− cudaMemsetAsync(grads.sh, 0, sizeof(ShT) * 3 * kShCoeffs * n, stream);
490− e != cudaSuccess)
491− return e;
492570 PreprocessBackwardKernel<ShT>
493− <<<GaussianBlocks(n), kGaussianThreads, 0, stream>>>(
571+ <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(
494572 params, frame, splats, splat_grads, grads);
495573 return cudaSuccess;
496574 }
497575
↕ Expand 16 unchanged lines (lines 576–591) · in PreprocessBackwardKernel
498576 template __global__ void PreprocessBackwardKernel<float>(
499577 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
500578 GaussianGrads);
501579 template __global__ void PreprocessBackwardKernel<ShBf16>(
502580 GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
503581 GaussianGrads);
504582 template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
505583 const Frame&, const Splat*,
506584 const float*,
507585 const GaussianGrads&,
508586 cudaStream_t);
509587 template cudaError_t LaunchPreprocessBackward<ShBf16>(
510588 const GaussianParams&, const Frame&, const Splat*, const float*,
511589 const GaussianGrads&, cudaStream_t);
512590
513591 } // namespace gs
Unchanged files in s03_sh_streaming (identical to s02_reduce_scatter): csrc/config.h, csrc/forward.h, csrc/forward.cu.cc, csrc/backward.h, csrc/rasterizer_impl.h, csrc/rasterizer_impl.cu.cc