`s04_raster_pruning` → `s05_branchless_predication`

Step 5: Branchless α=0 Predication & _rn-Pinned Separable Power2Column · +101 / −71 lines across 3 changed files

(a) Branchless α = 0 predication & RcpApprox in forward and backward rasterizers (csrc/forward.cu.cc, csrc/backward.cu.cc): Eliminates the per-thread if (!thread_done) / if (entry < thread_last) guards and the per-row if (done[k]) continue / if (power > 0.0f || power < g1.y) continue / if (alpha < kMinAlpha) continue branches inside the warp loop, substituting alpha = hit ? alpha_k : 0.0f with the 1-instruction RcpApprox(1.0f - alpha) (rcp.approx.ftz.f32) in the backward pass and w = blend ? alpha_k * t[k] : 0.0f in the forward pass, so the compiler unrolls 4 splats per iteration with zero inner-loop BSSY/BSYNC barriers. The per-splat min_power bound is no longer needed, shrinking batch_geom1 from float2 to float.
(b) IEEE _rn-pinned separable Power2Column (csrc/auxiliary.h): Precomputes the column-dependent terms Power2Column::At(g0.z, g0.w, dx) once per splat for the thread's pixel column and evaluates col.Row(c2, dy) for each of its rows with explicit __fmul_rn / __fmaf_rn intrinsics so forward and backward powers match bit-for-bit.

Core files: csrc/auxiliary.h +46/−15 csrc/forward.cu.cc +23/−26 csrc/backward.cu.cc +32/−30

csrc/auxiliary.h +46 −15

s04_raster_pruning/csrc/auxiliary.h

s05_branchless_predication/csrc/auxiliary.h

csrc/forward.cu.cc +23 −26

s04_raster_pruning/csrc/forward.cu.cc

s05_branchless_predication/csrc/forward.cu.cc

csrc/backward.cu.cc +32 −30

s04_raster_pruning/csrc/backward.cu.cc

s05_branchless_predication/csrc/backward.cu.cc

Unchanged files in s05_branchless_predication (identical to s04_raster_pruning): csrc/config.h, csrc/forward.h, csrc/backward.h, csrc/rasterizer_impl.h, csrc/rasterizer_impl.cu.cc

 /*
  * Copyright (C) 2023, Inria
  * GRAPHDECO research group, https://team.inria.fr/graphdeco
  * All rights reserved.
  *
  * This software is free for non-commercial, research and evaluation use
  * under the terms of the LICENSE.md file.
  *
  * For inquiries contact  george.drettakis@inria.fr
  */
 #ifndef CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
 #define CUDA_RASTERIZER_AUXILIARY_H_INCLUDED
 #include <cuda_runtime.h>
 #include <cmath>
 #include <cstdint>
 #include "config.h"
 #include "stdio.h"
 #define BLOCK_SIZE (BLOCK_X * BLOCK_Y)
 #define NUM_WARPS (BLOCK_SIZE / 32)
 // Spherical harmonics coefficients
 __device__ const float SH_C0 = 0.28209479177387814f;
 __device__ const float SH_C1 = 0.4886025119029199f;
 __device__ const float SH_C2[] = {1.0925484305920792f, -1.0925484305920792f,
                                   0.31539156525252005f, -1.0925484305920792f,
                                   0.5462742152960396f};
 __device__ const float SH_C3[] = {-0.5900435899266435f, 2.890611442640554f,
                                   -0.4570457994644658f, 0.3731763325901154f,
                                   -0.4570457994644658f, 1.445305721320277f,
                                   -0.5900435899266435f};
 namespace gs {
 #define GS_HD __host__ __device__ __forceinline__
 #define GS_D __device__ __forceinline__
 GS_HD float3 operator+(float3 a, float3 b) {
   return {a.x + b.x, a.y + b.y, a.z + b.z};
 }
 GS_HD float3 operator-(float3 a, float3 b) {
   return {a.x - b.x, a.y - b.y, a.z - b.z};
 }
 GS_HD float3 operator*(float s, float3 a) {
   return {s * a.x, s * a.y, s * a.z};
 }
 GS_HD float3& operator+=(float3& a, float3 b) {
   a.x += b.x;
   a.y += b.y;
   a.z += b.z;
   return a;
 }
 GS_HD float dot(float3 a, float3 b) {
   return a.x * b.x + a.y * b.y + a.z * b.z;
 }
 // Row-major 3x3 matrix.
 struct Mat3 {
   float m[3][3];
   GS_HD float3 row(int i) const { return {m[i][0], m[i][1], m[i][2]}; }
   GS_HD float3 col(int j) const { return {m[0][j], m[1][j], m[2][j]}; }
 };
 GS_HD float3 Mul(const Mat3& a, float3 v) {
   return {dot(a.row(0), v), dot(a.row(1), v), dot(a.row(2), v)};
 }
 // a^T v.
 GS_HD float3 MulT(const Mat3& a, float3 v) {
   return {dot(a.col(0), v), dot(a.col(1), v), dot(a.col(2), v)};
 }
 // Reinterprets bits; used to pack small integers into float records.
 GS_HD float AsFloat(uint32_t u) {
   union {
     uint32_t u;
     float f;
   } c{u};
   return c.f;
 }
 GS_HD uint32_t AsUint(float f) {
   union {
     float f;
     uint32_t u;
   } c{f};
   return c.u;
 }
 // Number of floats in the packed camera buffer built by `camera.py`:
 // [R (9, row-major world-to-camera), t (3), fx, fy, cx, cy, position (3), 0].
 inline constexpr int kCameraFloats = 20;
 // Pinhole camera with OpenCV axes (+x right, +y down, +z forward). Pixel (i, j)
 // is sampled at coordinate (j, i), as in the Inria rasterizer.
 struct Camera {
   Mat3 rotation;  // World-to-camera rotation.
   float3 translation;
   float fx, fy, cx, cy;
   float3 position;  // Camera center in world space.
   int width, height;
   GS_HD static Camera Load(const float* p, int width, int height) {
     Camera c;
     for (int i = 0; i < 3; ++i)
       for (int j = 0; j < 3; ++j) c.rotation.m[i][j] = p[3 * i + j];
     c.translation = {p[9], p[10], p[11]};
     c.fx = p[12];
     c.fy = p[13];
     c.cx = p[14];
     c.cy = p[15];
     c.position = {p[16], p[17], p[18]};
     c.width = width;
     c.height = height;
     return c;
   }
   GS_HD float3 ToView(float3 p) const { return Mul(rotation, p) + translation; }
   // Tangent of the half field of view, as the Inria rasterizer defines it.
   GS_HD float TanHalfFovX() const { return width / (2.0f * fx); }
   GS_HD float TanHalfFovY() const { return height / (2.0f * fy); }
 };
 // Screen tiling. Tiles are `kTileSize` x `kTileSize` pixels; one thread block
 // rasterizes one tile.
 inline constexpr int kTileSize = 16;
 inline constexpr int kTilePixels = kTileSize * kTileSize;
 struct TileGrid {
   int tiles_x, tiles_y;
   GS_HD static TileGrid For(int width, int height) {
     return {(width + kTileSize - 1) / kTileSize,
             (height + kTileSize - 1) / kTileSize};
   }
   GS_HD int num_tiles() const { return tiles_x * tiles_y; }
 };
 // Smallest alpha that contributes to a pixel; smaller ones are skipped.
 inline constexpr float kMinAlpha = 1.0f / 255.0f;
 inline constexpr float kMaxAlpha = 0.99f;
 // Transmittance below which a pixel stops accumulating.
 inline constexpr float kMinTransmittance = 1e-4f;
 // One splat as stored in memory: three float4 so the rasterizer issues three
 // 16-byte loads. Layout:
 //   p0 = (x, y, conic.a, conic.b)
 //   p1 = (conic.c, opacity, limit, radius)   radius 0 => culled
-//   p2 = (r, g, b, min_power)
+//   p2 = (r, g, b, 0)
 inline constexpr int kSplatFloats = 12;
 struct alignas(16) Splat {
   float4 p0, p1, p2;
   GS_HD float2 xy() const { return {p0.x, p0.y}; }
   GS_HD float3 conic() const { return {p0.z, p0.w, p1.x}; }
   GS_HD float opacity() const { return p1.y; }
   GS_HD float limit() const { return p1.z; }
   GS_HD int radius() const { return static_cast<int>(p1.w); }
   GS_HD float3 rgb() const { return {p2.x, p2.y, p2.z}; }
 };
 // Gradient w.r.t. a splat's rasterized quantities, accumulated by the raster
 // backward pass: (dx, dy, da, db, dc, dopacity, dr, dg, db).
 inline constexpr int kSplatGradFloats = 9;
 // Half-open tile rectangle [x0, x1) x [y0, y1).
 struct TileRect {
   int x0, y0, x1, y1;
   GS_HD int area() const { return (x1 - x0) * (y1 - y0); }
 };
 // Tile rectangle of a splat of integer bounding `radius`, identical to the
 // reference rasterizer.
 GS_HD TileRect RectFor(float2 xy, int radius, const TileGrid& grid) {
   const float r = static_cast<float>(radius);
   auto lo = [](float v, int n) {
     return min(n, max(0, static_cast<int>(v / kTileSize)));
   };
   auto hi = [](float v, int n) {
     return min(n, max(0, static_cast<int>((v + kTileSize - 1) / kTileSize)));
   };
   return {lo(xy.x - r, grid.tiles_x), lo(xy.y - r, grid.tiles_y),
           hi(xy.x + r, grid.tiles_x), hi(xy.y + r, grid.tiles_y)};
 }
 // How splats are assigned to tiles.
 enum class Culling : int {
   // Every tile of the bounding rectangle (reference behaviour).
   kRect = 0,
   // Tiles of the rectangle where some pixel can reach alpha >= 1/255. Gives
   // the same image: skipped (splat, tile) pairs never contribute.
   kEllipse = 1,
 };
 // alpha = o exp(-q/2) >= 1/255  <=>  q <= 2 ln(255 o). Writes the bound on q
 // with margins that keep tests conservative under float rounding of the
 // rasterizer's power. Returns false if no pixel can reach 1/255.
 GS_HD bool AlphaLimit(float opacity, float& limit) {
   const float o255 = opacity * 255.0f;
   if (!(o255 >= 0.999f)) return false;
   limit = 2.0f * logf(fmaxf(o255, 1.0f)) * 1.001f + 0.02f;
   return true;
 }
 // Pad (pixels) on ellipse extents, against rounding.
 inline constexpr float kSpanPadPx = 0.01f;
 // Precomputed splat-level ellipse constants for fast band intersection tests
 // of {q(d) <= limit}, where q(d) = a dx^2 + 2 b dx dy + c dy^2.
 struct EllipseBounds {
   float a, b, c;
   float det;
   float inv_a;
   float a_limit;
   float y_ext;
   float dy_peak;
   GS_D static bool Init(float3 conic, float opacity, float cached_limit,
                         EllipseBounds& out) {
     float limit = cached_limit;
     if (limit <= 0.0f && !AlphaLimit(opacity, limit)) return false;
     out.a = conic.x;
     out.b = conic.y;
     out.c = conic.z;
     out.det = out.a * out.c - out.b * out.b;
     if (!(out.det > 0.0f && out.a > 0.0f && out.c > 0.0f)) return false;
     out.inv_a = 1.0f / out.a;
     out.a_limit = out.a * limit;
     const float limit_over_det = limit / out.det;
     out.y_ext = sqrtf(out.a * limit_over_det);
     out.dy_peak = out.b * sqrtf(limit_over_det / out.c);
     return true;
   }
   // Span [x0_px, x1_px] of integer pixel columns hit across integer pixel rows
   // Y in [y0_px, y1_px] (inclusive). Returns false if no integer pixel hits.
   GS_D bool SpanForPixelRows(float2 xy, int y0_px, int y1_px, int& x0_px,
                              int& x1_px) const {
     const int ylo_i =
         max(y0_px, static_cast<int>(ceilf(xy.y - y_ext - kSpanPadPx)));
     const int yhi_i =
         min(y1_px, static_cast<int>(floorf(xy.y + y_ext + kSpanPadPx)));
     if (ylo_i > yhi_i) return false;
     const float ylo = static_cast<float>(ylo_i) - xy.y;
     const float yhi = static_cast<float>(yhi_i) - xy.y;
     // Boundary x(dy) = (-b dy +- sqrt(a limit - det dy^2)) / a; the right
     // (left) branch is concave (convex) and peaks at dy = -+dy_peak.
     const float dyr = fminf(yhi, fmaxf(ylo, -dy_peak));
     const float dyl = fminf(yhi, fmaxf(ylo, dy_peak));
     const float dyr2 = dyr * dyr, dyl2 = dyl * dyl;
     const float rr = sqrtf(fmaxf(0.0f, a_limit - det * dyr2));
     const float rl =
         (dyr2 == dyl2) ? rr : sqrtf(fmaxf(0.0f, a_limit - det * dyl2));
     const float hi = (-b * dyr + rr) * inv_a;
     const float lo = (-b * dyl - rl) * inv_a;
     x0_px = static_cast<int>(ceilf(xy.x + lo - kSpanPadPx));
     x1_px = static_cast<int>(floorf(xy.x + hi + kSpanPadPx));
     return x0_px <= x1_px;
   }
 };
 // Visits every tile row `ty` with the half-open span [tx0, tx1) of tiles that
 // splat `s` (not culled) is binned to. Rows may be empty (tx0 >= tx1).
 template <Culling kCulling, typename Fn>
 GS_D void ForEachTileRow(const Splat& s, const TileGrid& grid, int width,
                          int height, Fn&& fn) {
   const float2 xy = s.xy();
   const TileRect rect = RectFor(xy, s.radius(), grid);
   if constexpr (kCulling == Culling::kRect) {
     for (int ty = rect.y0; ty < rect.y1; ++ty) fn(ty, rect.x0, rect.x1);
   } else {
     EllipseBounds eb;
     if (!EllipseBounds::Init(s.conic(), s.opacity(), s.limit(), eb)) return;
     const int min_y =
         max(0, static_cast<int>(ceilf(xy.y - eb.y_ext - kSpanPadPx)));
     const int max_y =
         min(height - 1, static_cast<int>(floorf(xy.y + eb.y_ext + kSpanPadPx)));
     if (min_y > max_y) return;
     const int ty0 = max(rect.y0, min_y / kTileSize);
     const int ty1 = min(rect.y1, (max_y / kTileSize) + 1);
     for (int ty = ty0; ty < ty1; ++ty) {
       int x0_px, x1_px;
       if (!eb.SpanForPixelRows(xy, ty * kTileSize,
                                min((ty + 1) * kTileSize, height) - 1, x0_px,
                                x1_px)) {
         continue;
       }
       x0_px = max(0, x0_px);
       x1_px = min(width - 1, x1_px);
       if (x0_px > x1_px) continue;
       const int tx0 = max(rect.x0, x0_px / kTileSize);
       const int tx1 = min(rect.x1, (x1_px / kTileSize) + 1);
       if (tx0 < tx1) fn(ty, tx0, tx1);
     }
   }
 }
 namespace warp {
 inline constexpr unsigned kFullMask = 0xffffffffu;
 // Reduce-scatter of an array of `N` <= 32 per-lane values over the warp.
 //
 // A butterfly where each step halves the payload: at offset `o`, lanes with
 // bit `o` clear keep the low half of the array and send the high half, the
 // others do the opposite. After log2(32) steps every lane holds a single
 // value, v[0], which is the warp total of element `ScatterIndex(N, lane)`
 // (or junk if that is -1). For N = 9 this takes 12 shuffles instead of the
 // 45 of nine independent butterflies, and the totals end in distinct lanes,
 // so writing them out is one instruction.
 namespace internal {
 template <int N, int kOffset>
 GS_D void HalvingStep(float* v, int lane) {
   constexpr int kHalf = (N + 1) / 2;
   const bool upper = lane & kOffset;
 #pragma unroll
   for (int i = 0; i < kHalf; ++i) {
     const float lo = v[i];
     const float hi = i + kHalf < N ? v[i + kHalf] : 0.0f;
     const float send = upper ? lo : hi;
     const float keep = upper ? hi : lo;
     v[i] = keep + __shfl_xor_sync(kFullMask, send, kOffset);
   }
 }
 template <int N, int kOffset>
 GS_D void ReduceScatterImpl(float* v, int lane) {
   if constexpr (kOffset > 0) {
     HalvingStep<N, kOffset>(v, lane);
     ReduceScatterImpl<(N + 1) / 2, kOffset / 2>(v, lane);
   }
 }
 }  // namespace internal
 template <int N>
 GS_D void ReduceScatter(float (&v)[N], int lane) {
   static_assert(N >= 1 && N <= 32);
   internal::ReduceScatterImpl<N, 16>(v, lane);
 }
 template <int N>
 GS_D void ReduceScatter16(float (&v)[N], int lane) {
   static_assert(N >= 1 && N <= 16);
   internal::ReduceScatterImpl<N, 8>(v, lane);
 }
 // Element whose warp total lane `lane` holds after `ReduceScatter<N>`, or -1.
 GS_HD constexpr int ScatterIndex(int n, int lane) {
   // Payload sizes and halves per step, offset 16 first.
   int sizes[5] = {};
   int size = n;
   for (int s = 0; s < 5; ++s) {
     sizes[s] = size;
     size = (size + 1) / 2;
   }
   int index = 0;  // Slot in the final (size 1) array.
   for (int s = 4; s >= 0; --s) {
     const int half = (sizes[s] + 1) / 2;
     if (lane & (16 >> s)) index += half;
     if (index >= sizes[s]) return -1;
   }
   return index;
 }
 // Element whose 16-lane half-warp total lane `lane` holds after
 // `ReduceScatter16<N>`, or -1.
 GS_HD constexpr int ScatterIndex16(int n, int lane) {
   const int half_lane = lane & 15;
   int sizes[4] = {};
   int size = n;
   for (int s = 0; s < 4; ++s) {
     sizes[s] = size;
     size = (size + 1) / 2;
   }
   int index = 0;
   for (int s = 3; s >= 0; --s) {
     const int half = (sizes[s] + 1) / 2;
     if (half_lane & (8 >> s)) index += half;
     if (index >= sizes[s]) return -1;
   }
   return index;
 }
 }  // namespace warp
 inline constexpr int kShCoeffs = 16;  // Degree 3.
 inline constexpr float kNearPlane = 0.2f;
 inline constexpr float kLowPass = 0.3f;
 inline constexpr float kFovClamp = 1.3f;
 // Pointers to the trainable parameters, all activated, row-major.
 struct GaussianParams {
   const float* means;      // (N, 3)
   const float* scales;     // (N, 3)
   const float* quats;      // (N, 4) as (w, x, y, z), not normalized.
   const float* opacities;  // (N,)
   const void* sh;          // (N, 16, 3) of `float` or `ShBf16` (see `ShIo`).
   int n;
 };
 // Gradients w.r.t. `GaussianParams`, same layout and SH storage type.
 struct GaussianGrads {
   float* means;
   float* scales;
   float* quats;
   float* opacities;
   void* sh;
 };
 // Screen-space quantities needed by rasterization. `conic` is the inverse 2D
 // covariance (a, b, c) with power = -0.5 (a dx^2 + c dy^2) - b dx dy.
 struct Projection {
   float2 xy;
   float depth;
   float3 conic;
   float opacity;
   float limit;
   float3 rgb;
   int radius;  // 0 means culled.
 };
 // ----------------------------------------------------------------------------
 // SH coefficient storage. A row holds the 16 x 3 coefficients of one Gaussian
 // and is accessed in 4 groups of 4 coefficients (12 values). The coefficients
 // may be stored as f32 or as bf16 (raw bits, `ShBf16`); math is always f32.
 using ShBf16 = uint16_t;  // bfloat16 bits.
 template <typename T>
 struct ShIo;
 template <>
 struct ShIo<float> {
   static constexpr int kRowVec16 = 12;  // 16-byte vectors per row.
   GS_D static void LoadGroup(const float* row, int grp, float v[12]) {
     const float4* p = reinterpret_cast<const float4*>(row) + 3 * grp;
     const float4 a = __ldg(p), b = __ldg(p + 1), c = __ldg(p + 2);
     v[0] = a.x, v[1] = a.y, v[2] = a.z, v[3] = a.w;
     v[4] = b.x, v[5] = b.y, v[6] = b.z, v[7] = b.w;
     v[8] = c.x, v[9] = c.y, v[10] = c.z, v[11] = c.w;
   }
   GS_D static void StoreGroup(float* row, int grp, const float v[12]) {
     float4* p = reinterpret_cast<float4*>(row) + 3 * grp;
     p[0] = {v[0], v[1], v[2], v[3]};
     p[1] = {v[4], v[5], v[6], v[7]};
     p[2] = {v[8], v[9], v[10], v[11]};
   }
 };
 template <>
 struct ShIo<ShBf16> {
   static constexpr int kRowVec16 = 6;
   GS_D static float Lo(uint32_t w) { return __uint_as_float(w << 16); }
   GS_D static float Hi(uint32_t w) { return __uint_as_float(w & 0xffff0000u); }
   // Round to nearest even (inputs are finite).
   GS_D static uint32_t Round(float f) {
     const uint32_t u = __float_as_uint(f);
     return (u + 0x7fffu + ((u >> 16) & 1u)) >> 16;
   }
   GS_D static uint32_t Pack(float lo, float hi) {
     return Round(lo) | (Round(hi) << 16);
   }
   GS_D static void LoadGroup(const ShBf16* row, int grp, float v[12]) {
     // 24-byte groups: three 8-byte loads.
     const uint2* p = reinterpret_cast<const uint2*>(row) + 3 * grp;
 #pragma unroll
     for (int k = 0; k < 3; ++k) {
       const uint2 w = __ldg(p + k);
       v[4 * k] = Lo(w.x), v[4 * k + 1] = Hi(w.x);
       v[4 * k + 2] = Lo(w.y), v[4 * k + 3] = Hi(w.y);
     }
   }
   GS_D static void StoreGroup(ShBf16* row, int grp, const float v[12]) {
     uint2* p = reinterpret_cast<uint2*>(row) + 3 * grp;
 #pragma unroll
     for (int k = 0; k < 3; ++k) {
       p[k] = {Pack(v[4 * k], v[4 * k + 1]), Pack(v[4 * k + 2], v[4 * k + 3])};
     }
   }
 };
 template <typename ShT>
 GS_D const ShT* ShRow(const void* sh, int i) {
   return static_cast<const ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
 }
 template <typename ShT>
 GS_D ShT* ShRow(void* sh, int i) {
   return static_cast<ShT*>(sh) + static_cast<size_t>(i) * 3 * kShCoeffs;
 }
 // ----------------------------------------------------------------------------
 // Spherical harmonics (degree 3, Inria sign convention).
 namespace sh {
 inline constexpr float kC0 = 0.28209479177387814f;
 inline constexpr float kC1 = 0.4886025119029199f;
 inline constexpr float kC2[5] = {1.0925484305920792f, -1.0925484305920792f,
                                  0.31539156525252005f, -1.0925484305920792f,
                                  0.5462742152960396f};
 inline constexpr float kC3[7] = {-0.5900435899266435f, 2.890611442640554f,
                                  -0.4570457994644658f, 0.3731763325901154f,
                                  -0.4570457994644658f, 1.445305721320277f,
                                  -0.5900435899266435f};
 // Basis values for a unit direction.
 GS_D void Basis(float3 d, float b[kShCoeffs]) {
   const float x = d.x, y = d.y, z = d.z;
   const float xx = x * x, yy = y * y, zz = z * z;
   b[0] = kC0;
   b[1] = -kC1 * y;
   b[2] = kC1 * z;
   b[3] = -kC1 * x;
   b[4] = kC2[0] * x * y;
   b[5] = kC2[1] * y * z;
   b[6] = kC2[2] * (2.0f * zz - xx - yy);
   b[7] = kC2[3] * x * z;
   b[8] = kC2[4] * (xx - yy);
   b[9] = kC3[0] * y * (3.0f * xx - yy);
   b[10] = kC3[1] * x * y * z;
   b[11] = kC3[2] * y * (4.0f * zz - xx - yy);
   b[12] = kC3[3] * z * (2.0f * zz - 3.0f * xx - 3.0f * yy);
   b[13] = kC3[4] * x * (4.0f * zz - xx - yy);
   b[14] = kC3[5] * z * (xx - yy);
   b[15] = kC3[6] * x * (xx - 3.0f * yy);
 }
 // Pre-clamp color 0.5 + sum_k b_k sh_k, reading directly from an SH row.
 template <typename ShT>
 GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
   float3 c = {0.5f, 0.5f, 0.5f};
 #pragma unroll
   for (int grp = 0; grp < 4; ++grp) {
     const int k0 = 4 * grp;
     float v[12];
     ShIo<ShT>::LoadGroup(sh_row, grp, v);
     c += b[k0] * float3{v[0], v[1], v[2]};
     c += b[k0 + 1] * float3{v[3], v[4], v[5]};
     c += b[k0 + 2] * float3{v[6], v[7], v[8]};
     c += b[k0 + 3] * float3{v[9], v[10], v[11]};
   }
   return c;
 }
 }  // namespace sh
 // ----------------------------------------------------------------------------
 // Covariance helpers.
 // Standard rotation matrix of the (unnormalized) quaternion (w, x, y, z).
 GS_D Mat3 QuatToRotation(float4 q) {
   const float w = q.x, x = q.y, y = q.z, z = q.w;
   Mat3 r;
   r.m[0][0] = 1.0f - 2.0f * (y * y + z * z);
   r.m[0][1] = 2.0f * (x * y - w * z);
   r.m[0][2] = 2.0f * (x * z + w * y);
   r.m[1][0] = 2.0f * (x * y + w * z);
   r.m[1][1] = 1.0f - 2.0f * (x * x + z * z);
   r.m[1][2] = 2.0f * (y * z - w * x);
   r.m[2][0] = 2.0f * (x * z - w * y);
   r.m[2][1] = 2.0f * (y * z + w * x);
   r.m[2][2] = 1.0f - 2.0f * (x * x + y * y);
   return r;
 }
 // Symmetric 3x3 matrix.
 struct Sym3 {
   float xx, xy, xz, yy, yz, zz;
   GS_D float3 Mul(float3 v) const {
     return {xx * v.x + xy * v.y + xz * v.z, xy * v.x + yy * v.y + yz * v.z,
             xz * v.x + yz * v.y + zz * v.z};
   }
 };
 // Sigma = A A^T with A = R diag(s).
 GS_D Sym3 Covariance3D(const Mat3& r, float3 s) {
   Mat3 a;
   const float sv[3] = {s.x, s.y, s.z};
 #pragma unroll
   for (int i = 0; i < 3; ++i)
 #pragma unroll
     for (int j = 0; j < 3; ++j) a.m[i][j] = r.m[i][j] * sv[j];
   return {dot(a.row(0), a.row(0)), dot(a.row(0), a.row(1)),
           dot(a.row(0), a.row(2)), dot(a.row(1), a.row(1)),
           dot(a.row(1), a.row(2)), dot(a.row(2), a.row(2))};
 }
 // Intermediate values of the projection shared by forward and VJP.
 struct ProjectionState {
   float3 t;          // View-space mean.
   float tx_c, ty_c;  // Clamped x, y used by the Jacobian.
   bool x_clamped, y_clamped;
   float3 m0, m1;  // Rows of M = J W (2x3).
   Sym3 sigma;     // World-space 3D covariance.
   float a, b, c;  // Dilated 2D covariance.
 };
 // Computes 2D covariance given view-space mean `t` (with t.z > kNearPlane).
 GS_D void ProjectCovarianceFromView(const Camera& cam, float3 t, float3 scale,
                                     float4 quat, ProjectionState& st) {
   st.t = t;
   const float tz = t.z;
   const float limx = kFovClamp * cam.TanHalfFovX();
   const float limy = kFovClamp * cam.TanHalfFovY();
   const float txtz = t.x / tz, tytz = t.y / tz;
   st.x_clamped = txtz < -limx || txtz > limx;
   st.y_clamped = tytz < -limy || tytz > limy;
   st.tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
   st.ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
   const float j00 = cam.fx / tz, j02 = -cam.fx * st.tx_c / (tz * tz);
   const float j11 = cam.fy / tz, j12 = -cam.fy * st.ty_c / (tz * tz);
   const Mat3& w = cam.rotation;
   st.m0 = j00 * w.row(0) + j02 * w.row(2);
   st.m1 = j11 * w.row(1) + j12 * w.row(2);
   st.sigma = Covariance3D(QuatToRotation(quat), scale);
   const float3 s0 = st.sigma.Mul(st.m0), s1 = st.sigma.Mul(st.m1);
   st.a = dot(st.m0, s0) + kLowPass;
   st.b = dot(st.m0, s1);
   st.c = dot(st.m1, s1) + kLowPass;
 }
 // Computes everything up to the 2D covariance. Returns false if culled by the
 // near plane.
 GS_D bool ProjectCovariance(const Camera& cam, float3 mean, float3 scale,
                             float4 quat, ProjectionState& st) {
   const float3 t = cam.ToView(mean);
   if (t.z <= kNearPlane) return false;
   ProjectCovarianceFromView(cam, t, scale, quat, st);
   return true;
 }
 // Integer bounding radius (pixels) of the 2D covariance, as in the reference.
 GS_D int BoundingRadius(float a, float b, float c) {
   const float det = a * c - b * b;
   const float mid = 0.5f * (a + c);
   const float lambda = mid + sqrtf(fmaxf(0.1f, mid * mid - det));
   return static_cast<int>(ceilf(3.0f * sqrtf(lambda)));
 }
 GS_D float2 ProjectMean(const Camera& cam, float3 t) {
   return {cam.fx * (t.x / t.z) + cam.cx, cam.fy * (t.y / t.z) + cam.cy};
 }
 GS_D float3 LoadMean(const GaussianParams& p, int i) {
   return {p.means[3 * i], p.means[3 * i + 1], p.means[3 * i + 2]};
 }
 // Unit direction from the camera center to `mean`, and 1 / distance.
 GS_D float3 ViewDir(const Camera& cam, float3 mean, float& inv_len) {
   const float3 v = mean - cam.position;
   inv_len = rsqrtf(dot(v, v));
   return inv_len * v;
 }
 // Gradient of the loss w.r.t. one splat's screen-space quantities.
 struct ProjectionGrad {
   float2 xy;
   float3 conic;  // w.r.t. the scalars (a, b, c) of `Projection::conic`.
   float opacity;
   float3 rgb;
 };
 // Image geometry passed by value to every kernel.
 struct Frame {
   const float* camera;  // Device pointer to the packed camera.
   int width, height;
   TileGrid grid;
   GS_D Camera LoadCamera() const { return Camera::Load(camera, width, height); }
 };
 // ---------------------------------------------------------------------------
 // Per-Gaussian kernels run one thread per Gaussian.
 inline constexpr int kGaussianThreads = 128;
 inline int GaussianBlocks(int64_t n) {
   return static_cast<int>((n + kGaussianThreads - 1) / kGaussianThreads);
 }
-// Base-2 exponent scale: power2 = kPower2Scale * q = (-log2(e)/2) * q.
-inline constexpr float kPower2Scale = -0.72134752044448170f;  // -log2(e)/2
-inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
-GS_HD float3 Power2Conic(float a, float b, float c) {
-  return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
-}
-GS_D float Exp2Approx(float x) {
-  float y;
-  asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
-  return y;
-}
 // ---------------------------------------------------------------------------
 // Rasterization. One block per 16x16 tile. Each thread shades `kRows`
 // vertically adjacent pixels, so that per-splat costs (shared-memory loads,
 // barriers, the warp reduction in the backward pass) are paid once per
 // `kRows` pixels; more rows per thread also means fewer resident threads and
 // more waiting on the slowest pixel. A warp covers a compact
 // kWarpWidth x kWarpRows region (as square as possible, so per-warp culling
 // is tight); warps tile the 16x16 tile row-major. Splats are staged through
 // shared memory, one per thread per batch.
 inline constexpr int kRasterRowChoices[] = {1, 2, 4};
+// Per-pixel splat evaluation shared by the forward and backward rasterizers.
+// They must agree bit for bit on every alpha (the backward recovers
+// transmittance by dividing alphas back out), so the rounding of every step
+// is pinned with explicit `_rn` intrinsics.
+//
+// The falloff is evaluated in log2 units, power2 = A dx^2 + B dx dy + C dy^2
+// with (A, B, C) = -log2(e)/2 (a, 2b, c), so that the Gaussian is 2^power2:
+// two FMAs and one MUFU.EX2 per pixel. A pixel receives the splat iff
+// power2 <= 0 and alpha >= kMinAlpha (the latter implies power2 >= the
+// splat's culling bound, see `AlphaLimit`).
+inline constexpr float kPower2Scale = -0.72134752044448170f;  // -log2(e)/2
+inline constexpr float kInvPower2Scale = 1.0f / kPower2Scale;
+GS_HD float3 Power2Conic(float a, float b, float c) {
+  return {kPower2Scale * a, 2.0f * kPower2Scale * b, kPower2Scale * c};
+}
+// Splat terms of one pixel column at dx = splat.x - pixel.x.
+struct Power2Column {
+  float pa, pb;  // A dx^2, B dx.
+  GS_D static Power2Column At(float a2, float b2, float dx) {
+    return {__fmul_rn(__fmul_rn(a2, dx), dx), __fmul_rn(b2, dx)};
+  }
+  GS_D float Row(float c2, float dy) const {
+    return __fmaf_rn(__fmaf_rn(c2, dy, pb), dy, pa);
+  }
+};
+// 2^x and 1/x without the denormal range fix-ups of exp2f/__fdividef (their
+// inputs and results here are never denormal where it matters).
+GS_D float Exp2Approx(float x) {
+  float y;
+  asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+  return y;
+}
+GS_D float RcpApprox(float x) {
+  float y;
+  asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+  return y;
+}
+GS_D float SplatAlpha(float opacity, float gaussian) {
+  return fminf(kMaxAlpha, __fmul_rn(opacity, gaussian));
+}
 template <int kRows>
 struct RasterLayout {
   static_assert(kRows == 1 || kRows == 2 || kRows == 4);
   static constexpr int kThreads = kTilePixels / kRows;
   static constexpr int kWarps = kThreads / 32;
   static constexpr int kWarpWidth = kRows == 4 ? kTileSize : 8;
   static constexpr int kWarpRows = (32 / kWarpWidth) * kRows;
   static constexpr int kWarpsX = kTileSize / kWarpWidth;
   static_assert(kWarpsX * (kTileSize / kWarpRows) == kWarps);
   // Tile-relative origin of warp `w`'s region.
   GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
   GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
 };
 // Pixel coordinates of this thread's first pixel; the others follow in y.
 struct RasterPixels {
   int x, y0;
   template <int kRows>
   GS_D static RasterPixels For(int tid) {
     using L = RasterLayout<kRows>;
     const int warp = tid / 32, lane = tid % 32;
     return {static_cast<int>(blockIdx.x) * kTileSize + L::WarpX(warp) +
                 lane % L::kWarpWidth,
             static_cast<int>(blockIdx.y) * kTileSize + L::WarpY(warp) +
                 (lane / L::kWarpWidth) * kRows};
   }
 };
 // Bit w is set if splat (p0, p1) can reach alpha >= 1/255 somewhere in the
 // pixel region of warp w of this block's tile.
 template <int kRows>
 GS_D uint32_t WarpMask(float4 p0, float4 p1) {
   using L = RasterLayout<kRows>;
   constexpr int kWarpsY = L::kWarps / L::kWarpsX;
   float limit = p1.z;
   if (limit <= 0.0f && !AlphaLimit(p1.y, limit)) return 0;
   const float a = p0.z, b = p0.w, c = p1.x;
   const float det = a * c - b * b;
   if (!(det > 0.0f && a > 0.0f && c > 0.0f)) return 0;
   const float l_over_det = limit / det;
   const float rx = sqrtf(c * l_over_det);
   const float ry = sqrtf(a * l_over_det);
   const float2 xy = {p0.x, p0.y};
   const int x0 = static_cast<int>(ceilf(xy.x - rx - kSpanPadPx));
   const int x1 = static_cast<int>(floorf(xy.x + rx + kSpanPadPx));
   const int y0 = static_cast<int>(ceilf(xy.y - ry - kSpanPadPx));
   const int y1 = static_cast<int>(floorf(xy.y + ry + kSpanPadPx));
   if (x0 > x1 || y0 > y1) return 0;
   const int tx = blockIdx.x * kTileSize, ty = blockIdx.y * kTileSize;
   uint32_t mx = 0u;
 #pragma unroll
   for (int wx = 0; wx < L::kWarpsX; ++wx) {
     const int wx0 = tx + wx * L::kWarpWidth;
     if (x0 <= wx0 + L::kWarpWidth - 1 && x1 >= wx0) mx |= (1u << wx);
   }
   if (mx == 0u) return 0;
   uint32_t my = 0u;
 #pragma unroll
   for (int wy = 0; wy < kWarpsY; ++wy) {
     const int wy0 = ty + wy * L::kWarpRows;
     if (y0 <= wy0 + L::kWarpRows - 1 && y1 >= wy0) my |= (1u << wy);
   }
   if (my == 0u) return 0;
   if ((mx & (mx - 1u)) == 0u || (my & (my - 1u)) == 0u) {
     uint32_t mask = 0u;
 #pragma unroll
     for (int wy = 0; wy < kWarpsY; ++wy) {
       if (!((my >> wy) & 1u)) continue;
 #pragma unroll
       for (int wx = 0; wx < L::kWarpsX; ++wx) {
         if ((mx >> wx) & 1u) mask |= (1u << (wy * L::kWarpsX + wx));
       }
     }
     return mask;
   }
   EllipseBounds eb;
   eb.a = a;
   eb.b = b;
   eb.c = c;
   eb.det = det;
   eb.inv_a = 1.0f / a;
   eb.a_limit = a * limit;
   eb.y_ext = ry;
   eb.dy_peak = b * sqrtf(l_over_det / c);
   uint32_t mask = 0u;
 #pragma unroll
   for (int wy = 0; wy < kWarpsY; ++wy) {
     if (!((my >> wy) & 1u)) continue;
     const int wy0 = ty + wy * L::kWarpRows;
     int x0_px, x1_px;
     if (!eb.SpanForPixelRows(xy, wy0, wy0 + L::kWarpRows - 1, x0_px, x1_px)) {
       continue;
     }
 #pragma unroll
     for (int wx = 0; wx < L::kWarpsX; ++wx) {
       const int wx0 = tx + wx * L::kWarpWidth;
       if (x0_px <= wx0 + L::kWarpWidth - 1 && x1_px >= wx0) {
         mask |= 1u << (wy * L::kWarpsX + wx);
       }
     }
   }
   return mask;
 }
 }  // namespace gs
 #endif

 /*
  * Copyright (C) 2023, Inria
  * GRAPHDECO research group, https://team.inria.fr/graphdeco
  * All rights reserved.
  *
  * This software is free for non-commercial, research and evaluation use
  * under the terms of the LICENSE.md file.
  *
  * For inquiries contact  george.drettakis@inria.fr
  */
 #include "forward.h"
 #include "auxiliary.h"
 namespace gs {
 // Screen-space geometry of Gaussian `i`: everything in `Projection` except
 // `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
 // tile-rectangle cull).
 template <Culling kCulling>
 GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
                                 float3 mean, int i) {
   Projection out;
   out.radius = 0;
   const float3 t = cam.ToView(mean);
   if (t.z <= kNearPlane) return out;
   const float opacity = p.opacities[i];
   float limit = 0.0f;
   if (!AlphaLimit(opacity, limit)) {
     if constexpr (kCulling == Culling::kEllipse) return out;
   }
   const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
                         p.scales[3 * i + 2]};
   const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
   ProjectionState st;
   ProjectCovarianceFromView(cam, t, scale, quat, st);
   const float det = st.a * st.c - st.b * st.b;
   if (det == 0.0f) return out;
   const float inv = 1.0f / det;
   out.conic = {st.c * inv, -st.b * inv, st.a * inv};
   out.xy = ProjectMean(cam, st.t);
   out.depth = st.t.z;
   out.opacity = opacity;
   out.limit = limit;
   out.radius = BoundingRadius(st.a, st.b, st.c);
   return out;
 }
 // Color seen along `dir`, clamped at zero, reading from an SH row.
 template <typename ShT>
 GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
   float basis[kShCoeffs];
   sh::Basis(dir, basis);
   const float3 raw = sh::EvalRow(sh_row, basis);
   return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
 }
 // ---------------------------------------------------------------------------
 // Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
 // evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
 // depth sort via per-thread `atomicAdd`.
 template <Culling kCulling, typename ShT>
 __global__ void __launch_bounds__(kGaussianThreads, 8)
     PreprocessKernel(GaussianParams params, Frame frame,
                      Splat* __restrict__ splats,
                      int32_t* __restrict__ tile_counts,
                      uint32_t* __restrict__ depth_keys,
                      int32_t* __restrict__ ids,
                      int32_t* __restrict__ num_visible) {
   const int i = blockIdx.x * blockDim.x + threadIdx.x;
   if (i >= params.n) return;
   const Camera cam = frame.LoadCamera();
   // Geometry first: only Gaussians that reach a tile need their SH row.
   const float3 mean = LoadMean(params, i);
   const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
   Splat s;
   int count = 0;
   if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
-    const float min_power =
-        pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;
     s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
     s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
-    s.p2 = {0.0f, 0.0f, 0.0f, min_power};
+    s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
     ForEachTileRow<kCulling>(
         s, frame.grid, frame.width, frame.height,
         [&](int, int tx0, int tx1) { count += tx1 - tx0; });
   }
   if (count == 0) {
     s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
     splats[i] = s;
     tile_counts[i] = 0;
     return;
   }
   const int vis_idx = atomicAdd(num_visible, 1);
   float inv_len;
   const float3 rgb =
       ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
   s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
   splats[i] = s;
   tile_counts[i] = count;
   depth_keys[vis_idx] = AsUint(pr.depth);
   ids[vis_idx] = i;
 }
 template <int kRows>
 __global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
     RasterForwardKernel(const int2* __restrict__ ranges,
                         int32_t* __restrict__ point_list,
                         const Splat* __restrict__ splats, Frame frame,
                         float* __restrict__ image, float* __restrict__ final_t,
                         int32_t* __restrict__ n_contrib) {
   constexpr int kRasterRows = kRows;
   constexpr int kRasterThreads = RasterLayout<kRows>::kThreads;
   constexpr int kWarps = RasterLayout<kRows>::kWarps;
   const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
   const int tid = threadIdx.x;
   const int warp = tid / 32;
   const int lane = tid % 32;
   const RasterPixels px = RasterPixels::For<kRows>(tid);
   const float fx = static_cast<float>(px.x);
   const int2 range = ranges[tile];
   float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
   int last[kRasterRows];
   bool done[kRasterRows];
 #pragma unroll
   for (int k = 0; k < kRasterRows; ++k) {
     t[k] = 1.0f;
     r[k] = g[k] = b[k] = 0.0f;
     last[k] = 0;
     done[k] = px.x >= frame.width || px.y0 + k >= frame.height;
   }
   auto all_done = [&] {
     bool d = true;
 #pragma unroll
     for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
     return d;
   };
   __shared__ float4 batch_geom0[kRasterThreads];
-  __shared__ float2 batch_geom1[kRasterThreads];
+  __shared__ float batch_geom1[kRasterThreads];
   __shared__ float4 batch_rgb_o[kRasterThreads];
   __shared__ uint8_t batch_mask[kRasterThreads];
   __shared__ uint32_t warp_contrib[kWarps][kWarps];
   const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
   bool thread_done = all_done();
   int tile_compact = 0;
   for (int start = range.x; start < range.y; start += kRasterThreads) {
     const int idx = start + tid;
     int32_t id = 0;
     if (idx < range.y) {
       id = point_list[idx];
       const float4 p0 = __ldg(&splats[id].p0);
       const float4 p1 = __ldg(&splats[id].p1);
       const float4 p2 = __ldg(&splats[id].p2);
       const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
       batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
-      batch_geom1[tid] = {c2.z, p2.w};
+      batch_geom1[tid] = c2.z;
       batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
       batch_mask[tid] = static_cast<uint8_t>(WarpMask<kRows>(p0, p1));
     } else {
       batch_mask[tid] = 0u;
     }
     __syncthreads();
     const int count = min(kRasterThreads, range.y - start);
     bool warp_done = __all_sync(warp::kFullMask, thread_done);
 #pragma unroll
     for (int c = 0; c < kRasterThreads; c += 32) {
       uint32_t contrib = 0u;
       if (c < count && !warp_done) {
         uint32_t mask = __ballot_sync(
             warp::kFullMask,
             c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
         while (mask != 0u) {
           const int bit = __ffs(mask) - 1;
           mask &= mask - 1u;
           const int j = c + bit;
           bool blended = false;
-          if (!thread_done) {
+          {
+            // Predicated like the backward pass: rows that are done or not
+            // reached get alpha = 0 and keep their state.
             const float4 g0 = batch_geom0[j];
-            const float2 g1 = batch_geom1[j];
+            const float c2 = batch_geom1[j];
-            const float dx = g0.x - fx;
+            const float4 ro = batch_rgb_o[j];
-            const float cx_dx2 = g0.z * dx * dx;
-            const float cy_dx = g0.w * dx;
+            const Power2Column col =
+                Power2Column::At(g0.z, g0.w, __fsub_rn(g0.x, fx));
 #pragma unroll
             for (int k = 0; k < kRasterRows; ++k) {
-              if (done[k]) continue;
-              const float dy = g0.y - static_cast<float>(px.y0 + k);
+              const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
-              const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
+              const float power = col.Row(c2, dy);
-              if (power > 0.0f || power < g1.y) continue;
-              const float4 ro = batch_rgb_o[j];
-              const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
+              const float alpha_k = SplatAlpha(ro.w, Exp2Approx(power));
+              const bool hit =
-              if (alpha < kMinAlpha) continue;
+                  !done[k] && power <= 0.0f && alpha_k >= kMinAlpha;
-              const float next_t = t[k] * (1.0f - alpha);
+              const float next_t = t[k] * (1.0f - alpha_k);
-              if (next_t < kMinTransmittance) {
+              const bool stop = hit && next_t < kMinTransmittance;
+              const bool blend = hit && !stop;
-                done[k] = true;
+              done[k] = done[k] || stop;
-                continue;
-              }
-              const float w = alpha * t[k];
+              const float w = blend ? alpha_k * t[k] : 0.0f;
               r[k] += w * ro.x;
               g[k] += w * ro.y;
               b[k] += w * ro.z;
-              t[k] = next_t;
+              t[k] = blend ? next_t : t[k];
-              last[k] = -1 - j;
+              last[k] = blend ? -1 - j : last[k];
-              blended = true;
+              blended = blended || blend;
             }
             thread_done = all_done();
           }
           if (__any_sync(warp::kFullMask, blended)) {
             contrib |= (1u << bit);
           }
           if (__all_sync(warp::kFullMask, thread_done)) {
             warp_done = true;
             break;
           }
         }
         warp_done = __all_sync(warp::kFullMask, thread_done);
       }
       if (lane == 0) {
         warp_contrib[warp][c / 32] = contrib;
       }
     }
     const bool block_done =
         (__syncthreads_count(thread_done) == kRasterThreads);
     uint32_t chunk_any[kWarps];
     int chunk_prefix[kWarps];
     int batch_active = 0;
 #pragma unroll
     for (int u = 0; u < kWarps; ++u) {
       uint32_t any_u = 0u;
 #pragma unroll
       for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
         any_u |= warp_contrib[w_idx][u];
       }
       chunk_any[u] = any_u;
       chunk_prefix[u] = tile_compact + batch_active;
       batch_active += __popc(any_u);
     }
     if ((chunk_any[warp] >> lane) & 1u) {
       uint32_t cmask = 0u;
 #pragma unroll
       for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
         cmask |= ((warp_contrib[w_idx][warp] >> lane) & 1u) << w_idx;
       }
       const int dst =
           chunk_prefix[warp] + __popc(chunk_any[warp] & ((1u << lane) - 1u));
       point_list[range.x + dst] =
           static_cast<int32_t>((cmask << 24) | static_cast<uint32_t>(id));
     }
 #pragma unroll
     for (int k = 0; k < kRasterRows; ++k) {
       if (last[k] < 0) {
         const int j = -1 - last[k];
         const int u = j >> 5;
         const int b_idx = j & 31;
         last[k] =
             chunk_prefix[u] + __popc(chunk_any[u] & ((1u << b_idx) - 1u)) + 1;
       }
     }
     tile_compact += batch_active;
     if (block_done) break;
   }
   const int pixels = frame.width * frame.height;
 #pragma unroll
   for (int k = 0; k < kRasterRows; ++k) {
     const int y = px.y0 + k;
     if (px.x >= frame.width || y >= frame.height) continue;
     const int p = y * frame.width + px.x;
     image[0 * pixels + p] = r[k];
     image[1 * pixels + p] = g[k];
     image[2 * pixels + p] = b[k];
     final_t[p] = t[k];
     n_contrib[p] = last[k];
   }
 }
 template __global__ void PreprocessKernel<Culling::kRect, float>(
     GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
     uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
 template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
     GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
     uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
 template __global__ void PreprocessKernel<Culling::kEllipse, float>(
     GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
     uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
 template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
     GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
     uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
 template __global__ void RasterForwardKernel<1>(
     const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
     Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
 template __global__ void RasterForwardKernel<2>(
     const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
     Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
 template __global__ void RasterForwardKernel<4>(
     const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
     Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
 }  // namespace gs

↕ Expand 76 unchanged lines (lines 1–76) · in `PreprocessKernel`
1	1		/*
2	2		* Copyright (C) 2023, Inria
3	3		* GRAPHDECO research group, https://team.inria.fr/graphdeco
4	4		* All rights reserved.
5	5		*
6	6		* This software is free for non-commercial, research and evaluation use
7	7		* under the terms of the LICENSE.md file.
8	8		*
9	9		* For inquiries contact george.drettakis@inria.fr
10	10		*/
11	11
12	12		#include "forward.h"
13	13
14	14		#include "auxiliary.h"
15	15
16	16		namespace gs {
17	17
18	18		// Screen-space geometry of Gaussian `i`: everything in `Projection` except
19	19		// `rgb`. `radius` = 0 marks a culled Gaussian (the caller also applies the
20	20		// tile-rectangle cull).
21	21		template <Culling kCulling>
22	22		GS_D Projection ProjectGeometry(const GaussianParams& p, const Camera& cam,
23	23		float3 mean, int i) {
24	24		Projection out;
25	25		out.radius = 0;
26	26		const float3 t = cam.ToView(mean);
27	27		if (t.z <= kNearPlane) return out;
28	28		const float opacity = p.opacities[i];
29	29		float limit = 0.0f;
30	30		if (!AlphaLimit(opacity, limit)) {
31	31		if constexpr (kCulling == Culling::kEllipse) return out;
32	32		}
33	33		const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
34	34		p.scales[3 * i + 2]};
35	35		const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
36	36		ProjectionState st;
37	37		ProjectCovarianceFromView(cam, t, scale, quat, st);
38	38		const float det = st.a * st.c - st.b * st.b;
39	39		if (det == 0.0f) return out;
40	40		const float inv = 1.0f / det;
41	41		out.conic = {st.c * inv, -st.b * inv, st.a * inv};
42	42		out.xy = ProjectMean(cam, st.t);
43	43		out.depth = st.t.z;
44	44		out.opacity = opacity;
45	45		out.limit = limit;
46	46		out.radius = BoundingRadius(st.a, st.b, st.c);
47	47		return out;
48	48		}
49	49
50	50		// Color seen along `dir`, clamped at zero, reading from an SH row.
51	51		template <typename ShT>
52	52		GS_D float3 ShadeColorRow(const ShT* sh_row, float3 dir) {
53	53		float basis[kShCoeffs];
54	54		sh::Basis(dir, basis);
55	55		const float3 raw = sh::EvalRow(sh_row, basis);
56	56		return {fmaxf(raw.x, 0.0f), fmaxf(raw.y, 0.0f), fmaxf(raw.z, 0.0f)};
57	57		}
58	58
59	59		// ---------------------------------------------------------------------------
60	60		// Preprocess: project every Gaussian, count its tiles via `ForEachTileRow`,
61	61		// evaluate SH only for visible Gaussians, and emit depth keys for the 32-bit
62	62		// depth sort via per-thread `atomicAdd`.
63	63
64	64		template <Culling kCulling, typename ShT>
65	65		__global__ void __launch_bounds__(kGaussianThreads, 8)
66	66		PreprocessKernel(GaussianParams params, Frame frame,
67	67		Splat* __restrict__ splats,
68	68		int32_t* __restrict__ tile_counts,
69	69		uint32_t* __restrict__ depth_keys,
70	70		int32_t* __restrict__ ids,
71	71		int32_t* __restrict__ num_visible) {
72	72		const int i = blockIdx.x * blockDim.x + threadIdx.x;
73	73		if (i >= params.n) return;
74	74		const Camera cam = frame.LoadCamera();
75	75		// Geometry first: only Gaussians that reach a tile need their SH row.
76	76		const float3 mean = LoadMean(params, i);
77	77		const Projection pr = ProjectGeometry<kCulling>(params, cam, mean, i);
78	78		Splat s;
79	79		int count = 0;
80	80		if (pr.radius > 0 && RectFor(pr.xy, pr.radius, frame.grid).area() > 0) {
81		−	const float min_power =
82		−	pr.limit > 0.0f ? (kPower2Scale * pr.limit - 1e-3f) : -8.01f;
83	81		s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
84	82		s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
85		−	s.p2 = {0.0f, 0.0f, 0.0f, min_power};
	83	+	s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
86	84		ForEachTileRow<kCulling>(
87	85		s, frame.grid, frame.width, frame.height,
88	86		[&](int, int tx0, int tx1) { count += tx1 - tx0; });
89	87		}
↕ Expand 49 unchanged lines (lines 88–136) · in `RasterForwardKernel`
90	88		if (count == 0) {
91	89		s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
92	90		splats[i] = s;
93	91		tile_counts[i] = 0;
94	92		return;
95	93		}
96	94		const int vis_idx = atomicAdd(num_visible, 1);
97	95		float inv_len;
98	96		const float3 rgb =
99	97		ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
100	98		s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
101	99		splats[i] = s;
102	100		tile_counts[i] = count;
103	101		depth_keys[vis_idx] = AsUint(pr.depth);
104	102		ids[vis_idx] = i;
105	103		}
106	104
107	105		template <int kRows>
108	106		__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
109	107		RasterForwardKernel(const int2* __restrict__ ranges,
110	108		int32_t* __restrict__ point_list,
111	109		const Splat* __restrict__ splats, Frame frame,
112	110		float* __restrict__ image, float* __restrict__ final_t,
113	111		int32_t* __restrict__ n_contrib) {
114	112		constexpr int kRasterRows = kRows;
115	113		constexpr int kRasterThreads = RasterLayout<kRows>::kThreads;
116	114		constexpr int kWarps = RasterLayout<kRows>::kWarps;
117	115		const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
118	116		const int tid = threadIdx.x;
119	117		const int warp = tid / 32;
120	118		const int lane = tid % 32;
121	119		const RasterPixels px = RasterPixels::For<kRows>(tid);
122	120		const float fx = static_cast<float>(px.x);
123	121		const int2 range = ranges[tile];
124	122
125	123		float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
126	124		int last[kRasterRows];
127	125		bool done[kRasterRows];
128	126		#pragma unroll
129	127		for (int k = 0; k < kRasterRows; ++k) {
130	128		t[k] = 1.0f;
131	129		r[k] = g[k] = b[k] = 0.0f;
132	130		last[k] = 0;
133	131		done[k] = px.x >= frame.width \|\| px.y0 + k >= frame.height;
134	132		}
135	133		auto all_done = [&] {
136	134		bool d = true;
137	135		#pragma unroll
138	136		for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
139	137		return d;
140	138		};
141	139
142	140		__shared__ float4 batch_geom0[kRasterThreads];
143		−	__shared__ float2 batch_geom1[kRasterThreads];
	141	+	__shared__ float batch_geom1[kRasterThreads];
144	142		__shared__ float4 batch_rgb_o[kRasterThreads];
145	143		__shared__ uint8_t batch_mask[kRasterThreads];
146	144		__shared__ uint32_t warp_contrib[kWarps][kWarps];
147	145		const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
↕ Expand 8 unchanged lines (lines 146–153) · in `RasterForwardKernel`
148	146		bool thread_done = all_done();
149	147		int tile_compact = 0;
150	148		for (int start = range.x; start < range.y; start += kRasterThreads) {
151	149		const int idx = start + tid;
152	150		int32_t id = 0;
153	151		if (idx < range.y) {
154	152		id = point_list[idx];
155	153		const float4 p0 = __ldg(&splats[id].p0);
156	154		const float4 p1 = __ldg(&splats[id].p1);
157	155		const float4 p2 = __ldg(&splats[id].p2);
158	156		const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
159	157		batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
160		−	batch_geom1[tid] = {c2.z, p2.w};
	158	+	batch_geom1[tid] = c2.z;
161	159		batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
162	160		batch_mask[tid] = static_cast<uint8_t>(WarpMask<kRows>(p0, p1));
163	161		} else {
164	162		batch_mask[tid] = 0u;
↕ Expand 12 unchanged lines (lines 163–174) · in `RasterForwardKernel`
165	163		}
166	164		__syncthreads();
167	165		const int count = min(kRasterThreads, range.y - start);
168	166		bool warp_done = __all_sync(warp::kFullMask, thread_done);
169	167		#pragma unroll
170	168		for (int c = 0; c < kRasterThreads; c += 32) {
171	169		uint32_t contrib = 0u;
172	170		if (c < count && !warp_done) {
173	171		uint32_t mask = __ballot_sync(
174	172		warp::kFullMask,
175	173		c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
176	174		while (mask != 0u) {
177	175		const int bit = __ffs(mask) - 1;
178	176		mask &= mask - 1u;
179	177		const int j = c + bit;
180	178		bool blended = false;
181		−	if (!thread_done) {
	179	+	{
	180	+	// Predicated like the backward pass: rows that are done or not
	181	+	// reached get alpha = 0 and keep their state.
182	182		const float4 g0 = batch_geom0[j];
183		−	const float2 g1 = batch_geom1[j];
	183	+	const float c2 = batch_geom1[j];
184		−	const float dx = g0.x - fx;
	184	+	const float4 ro = batch_rgb_o[j];
185		−	const float cx_dx2 = g0.z * dx * dx;
186		−	const float cy_dx = g0.w * dx;
	185	+	const Power2Column col =
	186	+	Power2Column::At(g0.z, g0.w, __fsub_rn(g0.x, fx));
187	187		#pragma unroll
188	188		for (int k = 0; k < kRasterRows; ++k) {
189		−	if (done[k]) continue;
190		−	const float dy = g0.y - static_cast<float>(px.y0 + k);
	189	+	const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
191		−	const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
	190	+	const float power = col.Row(c2, dy);
192		−	if (power > 0.0f \|\| power < g1.y) continue;
193		−	const float4 ro = batch_rgb_o[j];
194		−	const float alpha = fminf(kMaxAlpha, ro.w * Exp2Approx(power));
	191	+	const float alpha_k = SplatAlpha(ro.w, Exp2Approx(power));
	192	+	const bool hit =
195		−	if (alpha < kMinAlpha) continue;
	193	+	!done[k] && power <= 0.0f && alpha_k >= kMinAlpha;
196		−	const float next_t = t[k] * (1.0f - alpha);
	194	+	const float next_t = t[k] * (1.0f - alpha_k);
197		−	if (next_t < kMinTransmittance) {
	195	+	const bool stop = hit && next_t < kMinTransmittance;
	196	+	const bool blend = hit && !stop;
198		−	done[k] = true;
	197	+	done[k] = done[k] \|\| stop;
199		−	continue;
200		−	}
201		−	const float w = alpha * t[k];
	198	+	const float w = blend ? alpha_k * t[k] : 0.0f;
202	199		r[k] += w * ro.x;
203	200		g[k] += w * ro.y;
204	201		b[k] += w * ro.z;
205		−	t[k] = next_t;
	202	+	t[k] = blend ? next_t : t[k];
206		−	last[k] = -1 - j;
	203	+	last[k] = blend ? -1 - j : last[k];
207		−	blended = true;
	204	+	blended = blended \|\| blend;
208	205		}
209	206		thread_done = all_done();
210	207		}
211	208		if (__any_sync(warp::kFullMask, blended)) {
↕ Expand 91 unchanged lines (lines 209–299) · in `RasterForwardKernel`
212	209		contrib \|= (1u << bit);
213	210		}
214	211		if (__all_sync(warp::kFullMask, thread_done)) {
215	212		warp_done = true;
216	213		break;
217	214		}
218	215		}
219	216		warp_done = __all_sync(warp::kFullMask, thread_done);
220	217		}
221	218		if (lane == 0) {
222	219		warp_contrib[warp][c / 32] = contrib;
223	220		}
224	221		}
225	222		const bool block_done =
226	223		(__syncthreads_count(thread_done) == kRasterThreads);
227	224		uint32_t chunk_any[kWarps];
228	225		int chunk_prefix[kWarps];
229	226		int batch_active = 0;
230	227		#pragma unroll
231	228		for (int u = 0; u < kWarps; ++u) {
232	229		uint32_t any_u = 0u;
233	230		#pragma unroll
234	231		for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
235	232		any_u \|= warp_contrib[w_idx][u];
236	233		}
237	234		chunk_any[u] = any_u;
238	235		chunk_prefix[u] = tile_compact + batch_active;
239	236		batch_active += __popc(any_u);
240	237		}
241	238		if ((chunk_any[warp] >> lane) & 1u) {
242	239		uint32_t cmask = 0u;
243	240		#pragma unroll
244	241		for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
245	242		cmask \|= ((warp_contrib[w_idx][warp] >> lane) & 1u) << w_idx;
246	243		}
247	244		const int dst =
248	245		chunk_prefix[warp] + __popc(chunk_any[warp] & ((1u << lane) - 1u));
249	246		point_list[range.x + dst] =
250	247		static_cast<int32_t>((cmask << 24) \| static_cast<uint32_t>(id));
251	248		}
252	249		#pragma unroll
253	250		for (int k = 0; k < kRasterRows; ++k) {
254	251		if (last[k] < 0) {
255	252		const int j = -1 - last[k];
256	253		const int u = j >> 5;
257	254		const int b_idx = j & 31;
258	255		last[k] =
259	256		chunk_prefix[u] + __popc(chunk_any[u] & ((1u << b_idx) - 1u)) + 1;
260	257		}
261	258		}
262	259		tile_compact += batch_active;
263	260		if (block_done) break;
264	261		}
265	262		const int pixels = frame.width * frame.height;
266	263		#pragma unroll
267	264		for (int k = 0; k < kRasterRows; ++k) {
268	265		const int y = px.y0 + k;
269	266		if (px.x >= frame.width \|\| y >= frame.height) continue;
270	267		const int p = y * frame.width + px.x;
271	268		image[0 * pixels + p] = r[k];
272	269		image[1 * pixels + p] = g[k];
273	270		image[2 * pixels + p] = b[k];
274	271		final_t[p] = t[k];
275	272		n_contrib[p] = last[k];
276	273		}
277	274		}
278	275
279	276		template __global__ void PreprocessKernel<Culling::kRect, float>(
280	277		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
281	278		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
282	279		template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
283	280		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
284	281		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
285	282		template __global__ void PreprocessKernel<Culling::kEllipse, float>(
286	283		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
287	284		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
288	285		template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
289	286		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
290	287		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
291	288
292	289		template __global__ void RasterForwardKernel<1>(
293	290		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
294	291		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
295	292		template __global__ void RasterForwardKernel<2>(
296	293		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
297	294		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
298	295		template __global__ void RasterForwardKernel<4>(
299	296		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
300	297		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
301	298
302	299		} // namespace gs

 /*
  * Copyright (C) 2023, Inria
  * GRAPHDECO research group, https://team.inria.fr/graphdeco
  * All rights reserved.
  *
  * This software is free for non-commercial, research and evaluation use
  * under the terms of the LICENSE.md file.
  *
  * For inquiries contact  george.drettakis@inria.fr
  */
 #include "backward.h"
 #include "auxiliary.h"
 namespace gs {
 namespace sh {
 // Evaluates SH color VJP reading the row `sh_in` and writing dL/dsh directly
 // to the row `sh_out` (same storage type), group by group.
 template <typename ShT>
 GS_D float3 VjpRow(float3 d, float3 g_rgb, const ShT* __restrict__ sh_in,
                    ShT* __restrict__ sh_out) {
   float b[kShCoeffs];
   Basis(d, b);
   const float3 raw = EvalRow(sh_in, b);
   const float3 drgb = {raw.x < 0.0f ? 0.0f : g_rgb.x,
                        raw.y < 0.0f ? 0.0f : g_rgb.y,
                        raw.z < 0.0f ? 0.0f : g_rgb.z};
   const float x = d.x, y = d.y, z = d.z;
   const float xx = x * x, yy = y * y, zz = z * z;
   const float xy = x * y, xz = x * z, yz = y * z;
   float3 ddir = {0.0f, 0.0f, 0.0f};
 #pragma unroll
   for (int grp = 0; grp < 4; ++grp) {
     const int k0 = 4 * grp;
     float v[12];
     ShIo<ShT>::LoadGroup(sh_in, grp, v);
     const float s0 = v[0] * drgb.x + v[1] * drgb.y + v[2] * drgb.z;
     const float s1 = v[3] * drgb.x + v[4] * drgb.y + v[5] * drgb.z;
     const float s2 = v[6] * drgb.x + v[7] * drgb.y + v[8] * drgb.z;
     const float s3 = v[9] * drgb.x + v[10] * drgb.y + v[11] * drgb.z;
     float out[12];
 #pragma unroll
     for (int q = 0; q < 4; ++q) {
       out[3 * q] = b[k0 + q] * drgb.x;
       out[3 * q + 1] = b[k0 + q] * drgb.y;
       out[3 * q + 2] = b[k0 + q] * drgb.z;
     }
     ShIo<ShT>::StoreGroup(sh_out, grp, out);
     if (grp == 0) {
       ddir.y += s1 * (-kC1);
       ddir.z += s2 * kC1;
       ddir.x += s3 * (-kC1);
     } else if (grp == 1) {
       ddir += (s0 * kC2[0]) * float3{y, x, 0.0f};
       ddir += (s1 * kC2[1]) * float3{0.0f, z, y};
       ddir += (s2 * kC2[2]) * float3{-2.0f * x, -2.0f * y, 4.0f * z};
       ddir += (s3 * kC2[3]) * float3{z, 0.0f, x};
     } else if (grp == 2) {
       ddir += (s0 * kC2[4]) * float3{2.0f * x, -2.0f * y, 0.0f};
       ddir += (s1 * kC3[0]) * float3{6.0f * xy, 3.0f * (xx - yy), 0.0f};
       ddir += (s2 * kC3[1]) * float3{yz, xz, xy};
       ddir += (s3 * kC3[2]) *
               float3{-2.0f * xy, 4.0f * zz - xx - 3.0f * yy, 8.0f * yz};
     } else {
       ddir += (s0 * kC3[3]) *
               float3{-6.0f * xz, -6.0f * yz, 6.0f * zz - 3.0f * xx - 3.0f * yy};
       ddir += (s1 * kC3[4]) *
               float3{4.0f * zz - 3.0f * xx - yy, -2.0f * xy, 8.0f * xz};
       ddir += (s2 * kC3[5]) * float3{2.0f * xz, -2.0f * yz, xx - yy};
       ddir += (s3 * kC3[6]) * float3{3.0f * (xx - yy), -6.0f * xy, 0.0f};
     }
   }
   return ddir;
 }
 }  // namespace sh
 // VJP of the projection for a touched Gaussian `i`; reads its SH row from
 // `sh_in` and writes dL/dsh directly to `sh_out`.
 // Writes non-SH parameter gradients to `dmean_out`, `dscale_out`, `dquat_out`.
 template <typename ShT>
 GS_D void ProjectVjp(const GaussianParams& p, const Camera& cam, int i,
                      const ShT* __restrict__ sh_in, ShT* __restrict__ sh_out,
                      const ProjectionGrad& g, float3& dmean_out,
                      float3& dscale_out, float4& dquat_out) {
   const float3 mean = LoadMean(p, i);
   const float3 scale = {p.scales[3 * i], p.scales[3 * i + 1],
                         p.scales[3 * i + 2]};
   const float4 quat = reinterpret_cast<const float4*>(p.quats)[i];
   // --- Color: rgb = max(0.5 + sum_k B_k(dir) sh_k, 0).
   float3 dmean = {0.0f, 0.0f, 0.0f};
   {
     float inv_len;
     const float3 dir = ViewDir(cam, mean, inv_len);
     const float3 ddir = sh::VjpRow(dir, g.rgb, sh_in, sh_out);
     // d(v/|v|)/dv = (I - dir dir^T) / |v|.
     dmean += inv_len * (ddir - dot(dir, ddir) * dir);
   }
   // --- Forward covariance intermediates (reusing R, s0, s1 in VJP).
   const float3 t = cam.ToView(mean);
   const float tz = t.z;
   const float limx = kFovClamp * cam.TanHalfFovX();
   const float limy = kFovClamp * cam.TanHalfFovY();
   const float txtz = t.x / tz, tytz = t.y / tz;
   const bool x_clamped = txtz < -limx || txtz > limx;
   const bool y_clamped = tytz < -limy || tytz > limy;
   const float tx_c = fminf(limx, fmaxf(-limx, txtz)) * tz;
   const float ty_c = fminf(limy, fmaxf(-limy, tytz)) * tz;
   const float j00 = cam.fx / tz, j02 = -cam.fx * tx_c / (tz * tz);
   const float j11 = cam.fy / tz, j12 = -cam.fy * ty_c / (tz * tz);
   const Mat3& w = cam.rotation;
   const float3 m0 = j00 * w.row(0) + j02 * w.row(2);
   const float3 m1 = j11 * w.row(1) + j12 * w.row(2);
   const Mat3 r = QuatToRotation(quat);
   const Sym3 sigma = Covariance3D(r, scale);
   const float3 s0 = sigma.Mul(m0), s1 = sigma.Mul(m1);
   const float cov_a = dot(m0, s0) + kLowPass;
   const float cov_b = dot(m0, s1);
   const float cov_c = dot(m1, s1) + kLowPass;
   // --- Conic -> 2D covariance: G_cov = -K G_K K.
   const float det = cov_a * cov_c - cov_b * cov_b;
   const float inv = 1.0f / det;
   const float ka = cov_c * inv, kb = -cov_b * inv, kc = cov_a * inv;
   const float ga = g.conic.x, gb = 0.5f * g.conic.y, gc = g.conic.z;
   // P = G_K K.
   const float p00 = ga * ka + gb * kb, p01 = ga * kb + gb * kc;
   const float p10 = gb * ka + gc * kb, p11 = gb * kb + gc * kc;
   // G_cov = -K P (symmetric).
   const float ca = -(ka * p00 + kb * p10);
   const float cb = -(ka * p01 + kb * p11);
   const float cc = -(kb * p01 + kc * p11);
   // --- cov = M Sigma M^T: dM = 2 G_cov M Sigma, dSigma = M^T G_cov M.
   const float3 dm0 = 2.0f * (ca * s0 + cb * s1);
   const float3 dm1 = 2.0f * (cb * s0 + cc * s1);
   // dSigma = ca m0 m0^T + cb (m0 m1^T + m1 m0^T) + cc m1 m1^T.
   Sym3 dsig;
   dsig.xx = ca * m0.x * m0.x + 2.0f * cb * m0.x * m1.x + cc * m1.x * m1.x;
   dsig.yy = ca * m0.y * m0.y + 2.0f * cb * m0.y * m1.y + cc * m1.y * m1.y;
   dsig.zz = ca * m0.z * m0.z + 2.0f * cb * m0.z * m1.z + cc * m1.z * m1.z;
   dsig.xy =
       ca * m0.x * m0.y + cb * (m0.x * m1.y + m1.x * m0.y) + cc * m1.x * m1.y;
   dsig.xz =
       ca * m0.x * m0.z + cb * (m0.x * m1.z + m1.x * m0.z) + cc * m1.x * m1.z;
   dsig.yz =
       ca * m0.y * m0.z + cb * (m0.y * m1.z + m1.y * m0.z) + cc * m1.y * m1.z;
   // --- M = J W: dJ = dM W^T (only the nonzero entries of J).
   const float dj00 = dot(dm0, w.row(0)), dj02 = dot(dm0, w.row(2));
   const float dj11 = dot(dm1, w.row(1)), dj12 = dot(dm1, w.row(2));
   const float tz2 = tz * tz, tz3 = tz2 * tz;
   float3 dt;
   dt.x = x_clamped ? 0.0f : -cam.fx / tz2 * dj02;
   dt.y = y_clamped ? 0.0f : -cam.fy / tz2 * dj12;
   dt.z = -cam.fx / tz2 * dj00 - cam.fy / tz2 * dj11 +
          2.0f * cam.fx * tx_c / tz3 * dj02 + 2.0f * cam.fy * ty_c / tz3 * dj12;
   // --- Screen-space mean u = fx x / z + cx.
   dt.x += cam.fx / tz * g.xy.x;
   dt.y += cam.fy / tz * g.xy.y;
   dt.z -= (cam.fx * t.x * g.xy.x + cam.fy * t.y * g.xy.y) / tz2;
   dmean += MulT(w, dt);
   dmean_out = dmean;
   // --- Sigma = A A^T, A = R diag(s): dA = 2 dSigma A.
   const float sv[3] = {scale.x, scale.y, scale.z};
   float ds[3];
   Mat3 dr;
 #pragma unroll
   for (int j = 0; j < 3; ++j) {
     const float3 r_col = r.col(j);
     const float3 da_col = (2.0f * sv[j]) * dsig.Mul(r_col);
     ds[j] = dot(da_col, r_col);
     dr.m[0][j] = da_col.x * sv[j];
     dr.m[1][j] = da_col.y * sv[j];
     dr.m[2][j] = da_col.z * sv[j];
   }
   dscale_out = {ds[0], ds[1], ds[2]};
   // --- R(q) -> q.
   const float qw = quat.x, qx = quat.y, qy = quat.z, qz = quat.w;
   const auto& G = dr.m;
   dquat_out = {
       2.0f * (-qz * G[0][1] + qy * G[0][2] + qz * G[1][0] - qx * G[1][2] -
               qy * G[2][0] + qx * G[2][1]),
       2.0f * (qy * G[0][1] + qz * G[0][2] + qy * G[1][0] - 2.0f * qx * G[1][1] -
               qw * G[1][2] + qz * G[2][0] + qw * G[2][1] - 2.0f * qx * G[2][2]),
       2.0f *
           (-2.0f * qy * G[0][0] + qx * G[0][1] + qw * G[0][2] + qx * G[1][0] +
            qz * G[1][2] - qw * G[2][0] + qz * G[2][1] - 2.0f * qy * G[2][2]),
       2.0f *
           (-2.0f * qz * G[0][0] - qw * G[0][1] + qx * G[0][2] + qw * G[1][0] -
            2.0f * qz * G[1][1] + qy * G[1][2] + qx * G[2][0] + qy * G[2][1])};
 }
 // Reverse traversal per pixel; accumulates splat gradients with atomics.
 template <int kRows>
 __global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 4)
     RasterBackwardKernel(const int2* __restrict__ ranges,
                          const int32_t* __restrict__ point_list,
                          const Splat* __restrict__ splats, Frame frame,
                          const float* __restrict__ final_t,
                          const int32_t* __restrict__ n_contrib,
                          const float* __restrict__ dl_dimage,
                          float* __restrict__ splat_grads,
                          uint8_t* __restrict__ splat_touched) {
   using L = RasterLayout<kRows>;
   constexpr int kRasterRows = kRows;
   constexpr int kRasterThreads = L::kThreads;
   constexpr int kWarps = L::kWarps;
   const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
   const int tid = threadIdx.x;
   const int warp = tid / 32;
   const int lane = tid % 32;
   const RasterPixels px = RasterPixels::For<kRows>(tid);
   const float fx = static_cast<float>(px.x);
   const int2 range = ranges[tile];
   struct PixelState {
     float t;
     float3 dpix;
     float accum;  // dot(color accumulated behind, dpix)
     int last;
   };
   PixelState st[kRasterRows];
   const int pixels = frame.width * frame.height;
   int thread_last = 0;
 #pragma unroll
   for (int k = 0; k < kRasterRows; ++k) {
     const int y = px.y0 + k;
     const bool inside = px.x < frame.width && y < frame.height;
     const int p = y * frame.width + px.x;
     st[k].t = inside ? final_t[p] : 0.0f;
     st[k].dpix =
         inside ? float3{dl_dimage[0 * pixels + p], dl_dimage[1 * pixels + p],
                         dl_dimage[2 * pixels + p]}
                : float3{0.0f, 0.0f, 0.0f};
     st[k].accum = 0.0f;
     st[k].last = inside ? n_contrib[p] : 0;
     thread_last = max(thread_last, st[k].last);
   }
   int warp_max_last = thread_last;
 #pragma unroll
   for (int offset = 16; offset > 0; offset /= 2) {
     warp_max_last = max(
         warp_max_last, __shfl_xor_sync(warp::kFullMask, warp_max_last, offset));
   }
   __shared__ int shared_max_last[kWarps];
   if (lane == 0) shared_max_last[warp] = warp_max_last;
   __syncthreads();
   int block_max_last = 0;
 #pragma unroll
   for (int w = 0; w < kWarps; ++w) {
     block_max_last = max(block_max_last, shared_max_last[w]);
   }
-  if (block_max_last <= 0) return;
+  if (block_max_last <= 0) {
+    return;
+  }
   __shared__ float4 batch_geom0[kRasterThreads];
-  __shared__ float2 batch_geom1[kRasterThreads];
+  __shared__ float batch_geom1[kRasterThreads];
   __shared__ float4 batch_rgb_o[kRasterThreads];
   __shared__ uint32_t batch_id[kRasterThreads];
   __shared__ uint8_t batch_mask[kRasterThreads];
   const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
   int scatter = -1;
   int scatter2 = -1;
   if constexpr (L::kWarpWidth == 8) {
     const int sub_lane = lane & 15;
     const int row = sub_lane >> 3, col = sub_lane & 7;
     if (col == 0) {
       scatter2 = row;
     } else if (col == 2) {
       scatter2 = 2 + row;
     } else if (col == 4) {
       scatter2 = row == 0 ? 4 : 7;
     } else if (col == 6) {
       scatter2 = row == 0 ? 6 : 8;
     } else if (col == 1 && row == 0) {
       scatter2 = 5;
     }
     const int full_row = lane >> 3, full_col = lane & 7;
     if (full_col == 0) {
       scatter = full_row;
     } else if (full_col == 2) {
       scatter = full_row == 0 ? 4 : 5 + full_row;
     } else if (full_col == 4 && full_row == 0) {
       scatter = 5;
     }
   } else {
     scatter = warp::ScatterIndex(kSplatGradFloats, lane);
   }
+  // Evaluates one splat at this thread's pixels and accumulates the per-thread
+  // partial gradient sums. Branch-free: lanes of a warp almost never agree on
+  // which pixels a splat reaches, so predication is cheaper than branching.
+  // Pixels the splat does not reach get alpha = 0, which leaves t and accum
+  // exactly unchanged (t * rcp(1) == t) and contributes zero gradient.
-  auto eval_splat = [&](int j, float4 g0, float2 g1, float4 ro, int entry,
+  auto eval_splat = [&](float4 g0, float c2, float4 ro, int entry, float& dx,
-                        float& dx, float& ga_sum, float& ga_dy, float& ga_dy2,
+                        float& ga_sum, float& ga_dy, float& ga_dy2,
                         float3& drgb) {
-    dx = g0.x - fx;
+    dx = __fsub_rn(g0.x, fx);
+    const Power2Column col = Power2Column::At(g0.z, g0.w, dx);
     ga_sum = 0.0f;
     ga_dy = 0.0f;
     ga_dy2 = 0.0f;
     drgb = {0.0f, 0.0f, 0.0f};
-    if (entry < thread_last) {
-      const float cx_dx2 = g0.z * dx * dx;
-      const float cy_dx = g0.w * dx;
 #pragma unroll
-      for (int k = 0; k < kRasterRows; ++k) {
+    for (int k = 0; k < kRasterRows; ++k) {
-        PixelState& ps = st[k];
+      PixelState& ps = st[k];
-        if (entry >= ps.last) continue;
-        const float dy = g0.y - static_cast<float>(px.y0 + k);
+      const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
-        const float power = cx_dx2 + g1.x * dy * dy + cy_dx * dy;
+      const float power = col.Row(c2, dy);
-        if (power > 0.0f || power < g1.y) continue;
-        const float g_val = Exp2Approx(power);
+      const float g_val = Exp2Approx(power);
-        const float alpha = fminf(kMaxAlpha, ro.w * g_val);
+      const float alpha_k = SplatAlpha(ro.w, g_val);
+      const bool hit = entry < ps.last && power <= 0.0f && alpha_k >= kMinAlpha;
-        if (alpha < kMinAlpha) continue;
+      const float alpha = hit ? alpha_k : 0.0f;
-        ps.t = __fdividef(ps.t, 1.0f - alpha);
+      ps.t = ps.t * RcpApprox(1.0f - alpha);
-        const float w = alpha * ps.t;
+      const float w = alpha * ps.t;
-        const float color_dot =
+      const float color_dot =
-            ro.x * ps.dpix.x + ro.y * ps.dpix.y + ro.z * ps.dpix.z;
+          ro.x * ps.dpix.x + ro.y * ps.dpix.y + ro.z * ps.dpix.z;
-        const float diff = color_dot - ps.accum;
+      const float diff = color_dot - ps.accum;
-        const float dl_dalpha = ps.t * diff;
-        ps.accum += alpha * diff;
+      ps.accum += alpha * diff;
-        const float ga = g_val * dl_dalpha;
+      const float ga = hit ? g_val * (ps.t * diff) : 0.0f;
-        ga_sum += ga;
+      ga_sum += ga;
-        const float ga_y = ga * dy;
+      const float ga_y = ga * dy;
-        ga_dy += ga_y;
+      ga_dy += ga_y;
-        ga_dy2 += ga_y * dy;
+      ga_dy2 += ga_y * dy;
-        drgb += w * ps.dpix;
+      drgb += w * ps.dpix;
-      }
+    }
-    }
   };
+  // Conic (a, b, c) from the staged log2-scaled one.
-  auto raw_conic = [](float4 g0, float2 g1) {
+  auto raw_conic = [](float4 g0, float c2) {
     return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
-                  kInvPower2Scale * g1.x};
+                  kInvPower2Scale * c2};
   };
   for (int end = block_max_last; end > 0; end -= kRasterThreads) {
     const int e = end - 1 - tid;
     if (e >= 0) {
       const uint32_t packed = static_cast<uint32_t>(point_list[range.x + e]);
       const uint32_t id = packed & 0x00ffffffu;
       splat_touched[id] = 1u;
       const float4 p0 = __ldg(&splats[id].p0);
       const float4 p1 = __ldg(&splats[id].p1);
       const float4 p2 = __ldg(&splats[id].p2);
       const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
       batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
-      batch_geom1[tid] = {c2.z, p2.w};
+      batch_geom1[tid] = c2.z;
       batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
       batch_id[tid] = id;
       batch_mask[tid] = static_cast<uint8_t>(packed >> 24);
     } else {
       batch_mask[tid] = 0u;
     }
     __syncthreads();
     const int count = min(kRasterThreads, end);
 #pragma unroll
     for (int c = 0; c < kRasterThreads; c += 32) {
       if (c < count && end - 32 - c < warp_max_last) {
         uint32_t mask = __ballot_sync(
             warp::kFullMask,
             c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
         if constexpr (L::kWarpWidth == 8) {
           while ((mask & (mask - 1u)) != 0u) {
             const int bit0 = __ffs(mask) - 1;
             mask &= mask - 1u;
             const int bit1 = __ffs(mask) - 1;
             mask &= mask - 1u;
             const int j0 = c + bit0, j1 = c + bit1;
             const float4 g0_0 = batch_geom0[j0];
-            const float2 g1_0 = batch_geom1[j0];
+            const float c2_0 = batch_geom1[j0];
             const float4 ro_0 = batch_rgb_o[j0];
             const float4 g0_1 = batch_geom0[j1];
-            const float2 g1_1 = batch_geom1[j1];
+            const float c2_1 = batch_geom1[j1];
             const float4 ro_1 = batch_rgb_o[j1];
             float dx0, ga_sum0, ga_dy0, ga_dy2_0;
             float3 drgb0;
-            eval_splat(j0, g0_0, g1_0, ro_0, end - 1 - j0, dx0, ga_sum0, ga_dy0,
+            eval_splat(g0_0, c2_0, ro_0, end - 1 - j0, dx0, ga_sum0, ga_dy0,
                        ga_dy2_0, drgb0);
             float dx1, ga_sum1, ga_dy1, ga_dy2_1;
             float3 drgb1;
-            eval_splat(j1, g0_1, g1_1, ro_1, end - 1 - j1, dx1, ga_sum1, ga_dy1,
+            eval_splat(g0_1, c2_1, ro_1, end - 1 - j1, dx1, ga_sum1, ga_dy1,
                        ga_dy2_1, drgb1);
             const bool hi16 = (lane & 16) != 0;
             const float opacity = hi16 ? ro_1.w : ro_0.w;
             const float dx = hi16 ? dx1 : dx0;
             const float3 cn =
-                hi16 ? raw_conic(g0_1, g1_1) : raw_conic(g0_0, g1_0);
+                hi16 ? raw_conic(g0_1, c2_1) : raw_conic(g0_0, c2_0);
             auto fold16 = [&](float v0, float v1) {
               const float send = hi16 ? v0 : v1;
               const float recv = __shfl_xor_sync(warp::kFullMask, send, 16);
               return (hi16 ? v1 : v0) + recv;
             };
             float ga_sum = fold16(ga_sum0, ga_sum1);
             float ga_dy = fold16(ga_dy0, ga_dy1);
             float b4[4] = {-0.5f * opacity * fold16(ga_dy2_0, ga_dy2_1),
                            fold16(drgb0.x, drgb1.x), fold16(drgb0.y, drgb1.y),
                            fold16(drgb0.z, drgb1.z)};
             warp::internal::HalvingStep<4, 8>(b4, lane);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
             ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
             const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
             const float p_dx = dx * p_sum;
             const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
                                                 : -(cn.y * p_dx + cn.z * p_dy);
             const float v_dx1 =
                 (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
             float v4[4] = {v_dx0, v_dx1, b4[0], b4[1]};
             warp::internal::HalvingStep<4, 4>(v4, lane);
             warp::internal::HalvingStep<2, 2>(v4, lane);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 4);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 2);
             const float v_sum =
                 v4[0] + __shfl_xor_sync(warp::kFullMask, v4[0], 1);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 1);
             const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
             if (scatter2 >= 0 && out_grad != 0.0f) {
               const uint32_t target_id = batch_id[hi16 ? j1 : j0];
               atomicAdd(splat_grads + kSplatGradFloats * target_id + scatter2,
                         out_grad);
             }
           }
         }
         while (mask != 0u) {
           const int bit = __ffs(mask) - 1;
           mask &= mask - 1u;
           const int j = c + bit;
           const float4 g0 = batch_geom0[j];
-          const float2 g1 = batch_geom1[j];
+          const float c2 = batch_geom1[j];
           const float4 ro = batch_rgb_o[j];
           float dx, ga_sum, ga_dy, ga_dy2;
           float3 drgb;
-          eval_splat(j, g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2,
+          eval_splat(g0, c2, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
-                     drgb);
           const float opacity = ro.w;
-          const float3 cn = raw_conic(g0, g1);
+          const float3 cn = raw_conic(g0, c2);
           float out_grad = 0.0f;
           if constexpr (L::kWarpWidth == 8) {
             float b4[4] = {-0.5f * opacity * ga_dy2, drgb.x, drgb.y, drgb.z};
             warp::internal::HalvingStep<4, 16>(b4, lane);
             warp::internal::HalvingStep<2, 8>(b4, lane);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 16);
             ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 16);
             ga_sum += __shfl_xor_sync(warp::kFullMask, ga_sum, 8);
             ga_dy += __shfl_xor_sync(warp::kFullMask, ga_dy, 8);
             const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
             const float p_dx = dx * p_sum;
             float v_dx;
             if ((lane & 16) == 0) {
               v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
                                      : -(cn.y * p_dx + cn.z * p_dy);
             } else {
               v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
             }
             float v[3] = {v_dx, b4[0], ga_sum};
             warp::internal::HalvingStep<3, 4>(v, lane);
             warp::internal::HalvingStep<2, 2>(v, lane);
             out_grad = v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
           } else {
             const float p_sum = opacity * ga_sum, p_dy = opacity * ga_dy;
             const float p_dx = dx * p_sum;
             float gr[kSplatGradFloats] = {-(cn.x * p_dx + cn.y * p_dy),  // xy
                                           -(cn.y * p_dx + cn.z * p_dy),
                                           -0.5f * dx * p_dx,  // conic
                                           -dx * p_dy,
                                           -0.5f * opacity * ga_dy2,
                                           ga_sum,  // opacity
                                           drgb.x,
                                           drgb.y,
                                           drgb.z};  // rgb
             warp::ReduceScatter(gr, lane);
             out_grad = gr[0];
           }
           if (scatter >= 0 && out_grad != 0.0f) {
             atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
                       out_grad);
           }
         }
       }
     }
     __syncthreads();
   }
 }
 // ---------------------------------------------------------------------------
 // Preprocess backward: splat gradients -> parameter gradients (dense).
 // Uses the 1-byte `splat_touched` flag written by `RasterBackwardKernel` to
 // skip reading `splat_grads` and `params.sh` for untouched Gaussians, and
 // zeros untouched SH rows with warp-coalesced float4 stores (0 shared memory).
 template <typename ShT>
 __global__ void __launch_bounds__(kGaussianThreads)
     PreprocessBackwardKernel(GaussianParams params, Frame frame,
                              const float* __restrict__ splat_grads,
                              const uint8_t* __restrict__ splat_touched,
                              GaussianGrads grads) {
   constexpr int kShVec16 = ShIo<ShT>::kRowVec16;  // 16-byte stores per row.
   const int i = blockIdx.x * blockDim.x + threadIdx.x;
   const int lane = threadIdx.x % 32;
   bool touched = (i < params.n && splat_touched[i] != 0u);
   const uint32_t init_mask = __ballot_sync(warp::kFullMask, touched);
   if (init_mask == 0u) {
     if (i < params.n) {
       grads.means[3 * i] = 0.0f;
       grads.means[3 * i + 1] = 0.0f;
       grads.means[3 * i + 2] = 0.0f;
       grads.scales[3 * i] = 0.0f;
       grads.scales[3 * i + 1] = 0.0f;
       grads.scales[3 * i + 2] = 0.0f;
       reinterpret_cast<float4*>(grads.quats)[i] = {0.0f, 0.0f, 0.0f, 0.0f};
       grads.opacities[i] = 0.0f;
     }
     const int warp_i0 =
         static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
     const int warp_n = max(0, min(32, params.n - warp_i0));
     uint4* warp_sh_dst =
         reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
     const uint4 zero = {0u, 0u, 0u, 0u};
     if (warp_n == 32) {
 #pragma unroll
       for (int s = 0; s < kShVec16; ++s) {
         warp_sh_dst[s * 32 + lane] = zero;
       }
     } else {
       for (int k = lane; k < warp_n * kShVec16; k += 32) {
         warp_sh_dst[k] = zero;
       }
     }
     return;
   }
   float3 dmean = {0.0f, 0.0f, 0.0f};
   float3 dscale = {0.0f, 0.0f, 0.0f};
   float4 dquat = {0.0f, 0.0f, 0.0f, 0.0f};
   float dopacity = 0.0f;
   if (touched) {
     float g[kSplatGradFloats];
     bool nonzero = false;
 #pragma unroll
     for (int q = 0; q < kSplatGradFloats; ++q) {
       g[q] = __ldg(&splat_grads[kSplatGradFloats * i + q]);
       nonzero = nonzero || g[q] != 0.0f;
     }
     touched = nonzero;
     if (touched) {
       ProjectionGrad pg;
       pg.xy = {g[0], g[1]};
       pg.conic = {g[2], g[3], g[4]};
       pg.opacity = g[5];
       pg.rgb = {g[6], g[7], g[8]};
       dopacity = g[5];
       ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
                  ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
     }
   }
   if (i < params.n) {
     grads.means[3 * i] = dmean.x;
     grads.means[3 * i + 1] = dmean.y;
     grads.means[3 * i + 2] = dmean.z;
     grads.scales[3 * i] = dscale.x;
     grads.scales[3 * i + 1] = dscale.y;
     grads.scales[3 * i + 2] = dscale.z;
     reinterpret_cast<float4*>(grads.quats)[i] = dquat;
     grads.opacities[i] = dopacity;
   }
   const uint32_t mask = __ballot_sync(warp::kFullMask, touched);
   if (mask != warp::kFullMask) {
     const int warp_i0 =
         static_cast<int>(blockIdx.x * blockDim.x) + (threadIdx.x / 32) * 32;
     const int warp_n = max(0, min(32, params.n - warp_i0));
     uint4* warp_sh_dst =
         reinterpret_cast<uint4*>(ShRow<ShT>(grads.sh, warp_i0));
     const uint4 zero = {0u, 0u, 0u, 0u};
     if (warp_n == 32) {
 #pragma unroll
       for (int s = 0; s < kShVec16; ++s) {
         const int k = s * 32 + lane;
         const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
         if (((mask >> row_in_warp) & 1u) == 0u) {
           warp_sh_dst[k] = zero;
         }
       }
     } else {
       for (int k = lane; k < warp_n * kShVec16; k += 32) {
         const unsigned row_in_warp = static_cast<unsigned>(k) / kShVec16;
         if (((mask >> row_in_warp) & 1u) == 0u) {
           warp_sh_dst[k] = zero;
         }
       }
     }
   }
 }
 template <typename ShT>
 cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
                                      const Frame& frame,
                                      const float* splat_grads,
                                      const uint8_t* splat_touched,
                                      const GaussianGrads& grads,
                                      cudaStream_t stream) {
   PreprocessBackwardKernel<ShT>
       <<<GaussianBlocks(params.n), kGaussianThreads, 0, stream>>>(
           params, frame, splat_grads, splat_touched, grads);
   return cudaSuccess;
 }
 template __global__ void RasterBackwardKernel<1>(
     const int2* __restrict__, const int32_t* __restrict__,
     const Splat* __restrict__, Frame, const float* __restrict__,
     const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
     uint8_t* __restrict__);
 template __global__ void RasterBackwardKernel<2>(
     const int2* __restrict__, const int32_t* __restrict__,
     const Splat* __restrict__, Frame, const float* __restrict__,
     const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
     uint8_t* __restrict__);
 template __global__ void RasterBackwardKernel<4>(
     const int2* __restrict__, const int32_t* __restrict__,
     const Splat* __restrict__, Frame, const float* __restrict__,
     const int32_t* __restrict__, const float* __restrict__, float* __restrict__,
     uint8_t* __restrict__);
 template __global__ void PreprocessBackwardKernel<float>(
     GaussianParams, Frame, const float* __restrict__,
     const uint8_t* __restrict__, GaussianGrads);
 template __global__ void PreprocessBackwardKernel<ShBf16>(
     GaussianParams, Frame, const float* __restrict__,
     const uint8_t* __restrict__, GaussianGrads);
 template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
                                                      const Frame&, const float*,
                                                      const uint8_t*,
                                                      const GaussianGrads&,
                                                      cudaStream_t);
 template cudaError_t LaunchPreprocessBackward<ShBf16>(
     const GaussianParams&, const Frame&, const float*, const uint8_t*,
     const GaussianGrads&, cudaStream_t);
 }  // namespace gs

s04_raster_pruning → s05_branchless_predication

`s04_raster_pruning` → `s05_branchless_predication`