`s01_literature` → `s02_reduce_scatter`

Step 2: Warp Reduce-Scatter & Two-Splat Half-Warp Folding · +234 / −55 lines across 2 changed files

(a) Halving register reduce-scatter & two-splat folding primitives (csrc/auxiliary.h): Introduces warp::internal::HalvingStep so paired lanes split the active gradient array in half at each step, finishing with distinct threads holding distinct gradient sums.
(b) Two-splat half-warp folding in backward rasterization (csrc/backward.cu.cc): In RasterBackwardKernel, pairs consecutive active splats via while ((mask & (mask - 1u)) != 0u), folds Splat 0 into the lower half-warp (T0..T15) and Splat 1 into the upper half-warp (T16..T31) via fold16 (10 shuffles/splat), and issues all 18 global atomicAdds simultaneously from 18 distinct threads (with a 12-shuffle halving path for any odd remainder splat).
(c) Occupancy target & warp-level early exit (csrc/backward.cu.cc): RasterBackwardKernel is now compiled with __launch_bounds__(256, 4) (was 6), i.e. 4 resident blocks of 256 threads per SM, which lets the compiler use up to 64 registers per thread without spilling; and a pair whose splats no pixel of the warp blends skips the folding, reduction and atomics entirely via if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue; (the odd remainder splat keeps the same __any_sync skip).

Core files: csrc/auxiliary.h +78/−9 csrc/backward.cu.cc +156/−46

csrc/auxiliary.h +78 −9

s01_literature/csrc/auxiliary.h

s02_reduce_scatter/csrc/auxiliary.h

csrc/backward.cu.cc +156 −46

s01_literature/csrc/backward.cu.cc

s02_reduce_scatter/csrc/backward.cu.cc

Unchanged files in s02_reduce_scatter (identical to s01_literature): csrc/config.h, csrc/forward.h, csrc/forward.cu.cc, 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)
 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;
-// Standard warp butterfly reduction of an array of `N` per-lane values:
+// Reduce-scatter of an array of `N` <= 32 per-lane values over the warp.
-// 5 steps of `__shfl_xor_sync` for each element (`5 * N` shuffles total),
-// after which lane 0 (and every other lane) holds the warp sum of `v[0..N)`.
+//
+// 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>
+template <int N, int kOffset>
-GS_D void ButterflyReduce(float (&v)[N]) {
+GS_D void HalvingStep(float* v, int lane) {
+  constexpr int kHalf = (N + 1) / 2;
+  const bool upper = lane & kOffset;
 #pragma unroll
-  for (int offset = 16; offset >= 1; offset >>= 1) {
-#pragma unroll
-    for (int i = 0; i < N; ++i) {
+  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] += __shfl_xor_sync(kFullMask, v[i], offset);
+    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 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> {
   GS_D static void LoadRow(const float* row, float3 coeffs[kShCoeffs]) {
     const float4* p = reinterpret_cast<const float4*>(row);
     float raw[3 * kShCoeffs];
 #pragma unroll
     for (int k = 0; k < 12; ++k) {
       const float4 q = __ldg(p + k);
       raw[4 * k] = q.x;
       raw[4 * k + 1] = q.y;
       raw[4 * k + 2] = q.z;
       raw[4 * k + 3] = q.w;
     }
 #pragma unroll
     for (int c = 0; c < kShCoeffs; ++c) {
       coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
     }
   }
   GS_D static void StoreRow(float* row, const float3 dsh[kShCoeffs]) {
     float raw[3 * kShCoeffs];
 #pragma unroll
     for (int c = 0; c < kShCoeffs; ++c) {
       raw[3 * c] = dsh[c].x;
       raw[3 * c + 1] = dsh[c].y;
       raw[3 * c + 2] = dsh[c].z;
     }
     float4* p = reinterpret_cast<float4*>(row);
 #pragma unroll
     for (int k = 0; k < 12; ++k) {
       p[k] = {raw[4 * k], raw[4 * k + 1], raw[4 * k + 2], raw[4 * k + 3]};
     }
   }
 };
 template <>
 struct ShIo<ShBf16> {
   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 LoadRow(const ShBf16* row, float3 coeffs[kShCoeffs]) {
     const uint2* p = reinterpret_cast<const uint2*>(row);
     float raw[3 * kShCoeffs];
 #pragma unroll
     for (int k = 0; k < 12; ++k) {
       const uint2 w = __ldg(p + k);
       raw[4 * k] = Lo(w.x);
       raw[4 * k + 1] = Hi(w.x);
       raw[4 * k + 2] = Lo(w.y);
       raw[4 * k + 3] = Hi(w.y);
     }
 #pragma unroll
     for (int c = 0; c < kShCoeffs; ++c) {
       coeffs[c] = {raw[3 * c], raw[3 * c + 1], raw[3 * c + 2]};
     }
   }
   GS_D static void StoreRow(ShBf16* row, const float3 dsh[kShCoeffs]) {
     float raw[3 * kShCoeffs];
 #pragma unroll
     for (int c = 0; c < kShCoeffs; ++c) {
       raw[3 * c] = dsh[c].x;
       raw[3 * c + 1] = dsh[c].y;
       raw[3 * c + 2] = dsh[c].z;
     }
     uint2* p = reinterpret_cast<uint2*>(row);
 #pragma unroll
     for (int k = 0; k < 12; ++k) {
       p[k] = {Pack(raw[4 * k], raw[4 * k + 1]),
               Pack(raw[4 * k + 2], raw[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);
 }
 // Partial derivatives of every basis function w.r.t. (x, y, z) of the unit
 // direction (treated as free coordinates).
 GS_D void BasisGrad(float3 d, float3 g[kShCoeffs]) {
   const float x = d.x, y = d.y, z = d.z;
   const float xx = x * x, yy = y * y, zz = z * z;
   g[0] = {0.0f, 0.0f, 0.0f};
   g[1] = {0.0f, -kC1, 0.0f};
   g[2] = {0.0f, 0.0f, kC1};
   g[3] = {-kC1, 0.0f, 0.0f};
   g[4] = kC2[0] * float3{y, x, 0.0f};
   g[5] = kC2[1] * float3{0.0f, z, y};
   g[6] = kC2[2] * float3{-2.0f * x, -2.0f * y, 4.0f * z};
   g[7] = kC2[3] * float3{z, 0.0f, x};
   g[8] = kC2[4] * float3{2.0f * x, -2.0f * y, 0.0f};
   g[9] = kC3[0] * float3{6.0f * x * y, 3.0f * (xx - yy), 0.0f};
   g[10] = kC3[1] * float3{y * z, x * z, x * y};
   g[11] =
       kC3[2] * float3{-2.0f * x * y, 4.0f * zz - xx - 3.0f * yy, 8.0f * y * z};
   g[12] = kC3[3] * float3{-6.0f * x * z, -6.0f * y * z,
                           6.0f * zz - 3.0f * xx - 3.0f * yy};
   g[13] =
       kC3[4] * float3{4.0f * zz - 3.0f * xx - yy, -2.0f * x * y, 8.0f * x * z};
   g[14] = kC3[5] * float3{2.0f * x * z, -2.0f * y * z, xx - yy};
   g[15] = kC3[6] * float3{3.0f * (xx - yy), -6.0f * x * y, 0.0f};
 }
 // Pre-clamp color 0.5 + sum_k b_k sh_k, loading the full 16-basis SH row.
 template <typename ShT>
 GS_D float3 EvalRow(const ShT* __restrict__ sh_row, const float b[kShCoeffs]) {
   float3 coeffs[kShCoeffs];
   ShIo<ShT>::LoadRow(sh_row, coeffs);
   float3 c = {0.5f, 0.5f, 0.5f};
 #pragma unroll
   for (int k = 0; k < kShCoeffs; ++k) {
     c += b[k] * coeffs[k];
   }
   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, one pixel per thread, arranged into
 // eight 8x4 sub-tile warps (`RasterLayout<1>`). Splats are staged through
 // shared memory in batches of 256 with an 8-bit per-warp sub-tile mask.
 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);
   GS_HD static constexpr int WarpX(int w) { return (w % kWarpsX) * kWarpWidth; }
   GS_HD static constexpr int WarpY(int w) { return (w / kWarpsX) * kWarpRows; }
 };
 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 "backward.h"
 #include "auxiliary.h"
 namespace gs {
 namespace sh {
 // Evaluates SH color VJP loading the full 16-basis row `sh_in` and storing
 // `dsh` to `sh_out`.
 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);
   float3 coeffs[kShCoeffs];
   ShIo<ShT>::LoadRow(sh_in, coeffs);
   float3 raw = {0.5f, 0.5f, 0.5f};
 #pragma unroll
   for (int k = 0; k < kShCoeffs; ++k) {
     raw += b[k] * coeffs[k];
   }
   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};
   float3 bgrad[kShCoeffs];
   BasisGrad(d, bgrad);
   float3 dsh[kShCoeffs];
   float3 ddir = {0.0f, 0.0f, 0.0f};
 #pragma unroll
   for (int k = 0; k < kShCoeffs; ++k) {
     dsh[k] = b[k] * drgb;
     const float s = dot(coeffs[k], drgb);
     ddir += s * bgrad[k];
   }
   ShIo<ShT>::StoreRow(sh_out, dsh);
   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 standard
+// Reverse traversal per pixel; accumulates splat gradients via Half-Warp
-// 5-step full-warp butterfly reductions (45 shuffles) and 9 lane-0 atomics.
+// Two-Splat Folding and Halving Register Reduce-Scatter.
-__global__ void __launch_bounds__(RasterLayout<1>::kThreads, 6)
+__global__ void __launch_bounds__(RasterLayout<1>::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) {
   using L = RasterLayout<1>;
   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<1>(tid);
   const float fx = static_cast<float>(px.x);
   const float fy = static_cast<float>(px.y0);
   const int2 range = ranges[tile];
   const bool inside = px.x < frame.width && px.y0 < frame.height;
   const int pixels = frame.width * frame.height;
   const int p = px.y0 * frame.width + px.x;
   float t = inside ? final_t[p] : 0.0f;
   const float3 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};
   float accum = 0.0f;  // dot(color accumulated behind, dpix)
   const int last = inside ? n_contrib[p] : 0;
   int warp_max_last = 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;
   __shared__ float4 batch_geom0[kRasterThreads];
   __shared__ float2 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;
+  {
+    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;
+    }
+  }
+  auto eval_splat = [&](float4 g0, float2 g1, float4 ro, int entry, float& dx,
+                        float& ga_sum, float& ga_dy, float& ga_dy2,
+                        float3& drgb) -> bool {
+    dx = g0.x - fx;
+    ga_sum = 0.0f;
+    ga_dy = 0.0f;
+    ga_dy2 = 0.0f;
+    drgb = {0.0f, 0.0f, 0.0f};
+    if (entry < last) {
+      const float dy = g0.y - fy;
+      const float power = g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
+      if (power <= 0.0f && power >= g1.y) {
+        const float g_val = Exp2Approx(power);
+        const float alpha = fminf(kMaxAlpha, ro.w * g_val);
+        if (alpha >= kMinAlpha) {
+          t = __fdividef(t, 1.0f - alpha);
+          const float w = alpha * t;
+          const float color_dot = ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
+          const float diff = color_dot - accum;
+          const float dl_dalpha = t * diff;
+          accum += alpha * diff;
+          const float ga = g_val * dl_dalpha;
+          ga_sum = ga;
+          const float ga_y = ga * dy;
+          ga_dy = ga_y;
+          ga_dy2 = ga_y * dy;
+          drgb = w * dpix;
+          return true;
+        }
+      }
+    }
+    return false;
+  };
   auto raw_conic = [](float4 g0, float2 g1) {
     return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
                   kInvPower2Scale * g1.x};
   };
   for (int end = block_max_last; end > 0; end -= kRasterThreads) {
     const int e = end - 1 - tid;
     if (e >= 0) {
       const uint32_t id = static_cast<uint32_t>(point_list[range.x + e]);
       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_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
       batch_id[tid] = id;
       batch_mask[tid] = static_cast<uint8_t>(WarpMask<1>(p0, p1));
     } 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);
+        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 float4 ro_0 = batch_rgb_o[j0];
+          const float4 g0_1 = batch_geom0[j1];
+          const float2 g1_1 = batch_geom1[j1];
+          const float4 ro_1 = batch_rgb_o[j1];
+          float dx0, ga_sum0, ga_dy0, ga_dy2_0;
+          float3 drgb0;
+          const bool hit0 = eval_splat(g0_0, g1_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;
+          const bool hit1 = eval_splat(g0_1, g1_1, ro_1, end - 1 - j1, dx1,
+                                       ga_sum1, ga_dy1, ga_dy2_1, drgb1);
+          // As in s01: a pair that no pixel of the warp blends has zero
+          // gradients, so skip its reduction and atomics entirely.
+          if (!__any_sync(warp::kFullMask, hit0 || hit1)) continue;
+          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);
+          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 int entry = end - 1 - j;
-          float gr[kSplatGradFloats] = {};
-          bool hit = false;
-          if (entry < last) {
-            const float4 g0 = batch_geom0[j];
+          const float4 g0 = batch_geom0[j];
-            const float2 g1 = batch_geom1[j];
+          const float2 g1 = batch_geom1[j];
-            const float dx = g0.x - fx;
-            const float dy = g0.y - fy;
-            const float power =
-                g0.z * dx * dx + g1.x * dy * dy + g0.w * dx * dy;
-            if (power <= 0.0f && power >= g1.y) {
-              const float4 ro = batch_rgb_o[j];
+          const float4 ro = batch_rgb_o[j];
-              const float g_val = Exp2Approx(power);
-              const float alpha = fminf(kMaxAlpha, ro.w * g_val);
-              if (alpha >= kMinAlpha) {
-                hit = true;
-                t = __fdividef(t, 1.0f - alpha);
-                const float w = alpha * t;
-                const float color_dot =
-                    ro.x * dpix.x + ro.y * dpix.y + ro.z * dpix.z;
-                const float diff = color_dot - accum;
-                const float dl_dalpha = t * diff;
-                accum += alpha * diff;
-                const float ga = g_val * dl_dalpha;
+          float dx, ga_sum, ga_dy, ga_dy2;
-                const float dl_dpower = ro.w * ga;
+          float3 drgb;
-                const float p_dx = dx * dl_dpower;
+          const bool hit =
+              eval_splat(g0, g1, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
+          if (!__any_sync(warp::kFullMask, hit)) continue;
-                const float p_dy = dy * dl_dpower;
+          const float opacity = ro.w;
-                const float3 cn = raw_conic(g0, g1);
+          const float3 cn = raw_conic(g0, g1);
+          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) {
-                gr[0] = -(cn.x * p_dx + cn.y * p_dy);
+            v_dx = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
-                gr[1] = -(cn.y * p_dx + cn.z * p_dy);
+                                   : -(cn.y * p_dx + cn.z * p_dy);
+          } else {
-                gr[2] = -0.5f * dx * p_dx;
+            v_dx = (lane & 8) == 0 ? -0.5f * dx * p_dx : -dx * p_dy;
-                gr[3] = -dx * p_dy;
-                gr[4] = -0.5f * dy * p_dy;
-                gr[5] = ga;
-                gr[6] = w * dpix.x;
-                gr[7] = w * dpix.y;
-                gr[8] = w * dpix.z;
-              }
+          }
-            }
-          }
+          float v[3] = {v_dx, b4[0], ga_sum};
+          warp::internal::HalvingStep<3, 4>(v, lane);
+          warp::internal::HalvingStep<2, 2>(v, lane);
+          const float out_grad =
-          if (!__any_sync(warp::kFullMask, hit)) continue;
+              v[0] + __shfl_xor_sync(warp::kFullMask, v[0], 1);
-          warp::ButterflyReduce(gr);
-          if (lane == 0) {
+          if (scatter >= 0 && out_grad != 0.0f) {
-            float* dst = splat_grads + kSplatGradFloats * batch_id[j];
+            atomicAdd(splat_grads + kSplatGradFloats * batch_id[j] + scatter,
-#pragma unroll
-            for (int q = 0; q < kSplatGradFloats; ++q) {
-              atomicAdd(dst + q, gr[q]);
-            }
+                      out_grad);
           }
         }
       }
     }
     __syncthreads();
   }
 }
 // ---------------------------------------------------------------------------
 // Preprocess backward: splat gradients -> parameter gradients.
 // Parameter gradient buffers are bulk-zeroed beforehand via `cudaMemsetAsync`;
 // culled or zero-gradient Gaussians return early.
 template <typename ShT>
 __global__ void __launch_bounds__(kGaussianThreads)
     PreprocessBackwardKernel(GaussianParams params, Frame frame,
                              const Splat* __restrict__ splats,
                              const float* __restrict__ splat_grads,
                              GaussianGrads grads) {
   const int i = blockIdx.x * blockDim.x + threadIdx.x;
   if (i >= params.n || splats[i].radius() == 0) return;
   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;
   }
   if (!nonzero) return;
   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]};
   float3 dmean, dscale;
   float4 dquat;
   ProjectVjp(params, frame.LoadCamera(), i, ShRow<ShT>(params.sh, i),
              ShRow<ShT>(grads.sh, i), pg, dmean, dscale, dquat);
   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] = g[5];
 }
 template <typename ShT>
 cudaError_t LaunchPreprocessBackward(const GaussianParams& params,
                                      const Frame& frame, const Splat* splats,
                                      const float* splat_grads,
                                      const GaussianGrads& grads,
                                      cudaStream_t stream) {
   const int n = params.n;
   // As Inria's PyTorch wrapper (`torch::zeros`): bulk-zero every parameter
   // gradient up front; culled Gaussians then exit early in the kernel.
   if (auto e = cudaMemsetAsync(grads.means, 0, sizeof(float) * 3 * n, stream);
       e != cudaSuccess)
     return e;
   if (auto e = cudaMemsetAsync(grads.scales, 0, sizeof(float) * 3 * n, stream);
       e != cudaSuccess)
     return e;
   if (auto e = cudaMemsetAsync(grads.quats, 0, sizeof(float) * 4 * n, stream);
       e != cudaSuccess)
     return e;
   if (auto e = cudaMemsetAsync(grads.opacities, 0, sizeof(float) * n, stream);
       e != cudaSuccess)
     return e;
   if (auto e =
           cudaMemsetAsync(grads.sh, 0, sizeof(ShT) * 3 * kShCoeffs * n, stream);
       e != cudaSuccess)
     return e;
   PreprocessBackwardKernel<ShT>
       <<<GaussianBlocks(n), kGaussianThreads, 0, stream>>>(
           params, frame, splats, splat_grads, grads);
   return cudaSuccess;
 }
 template __global__ void PreprocessBackwardKernel<float>(
     GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
     GaussianGrads);
 template __global__ void PreprocessBackwardKernel<ShBf16>(
     GaussianParams, Frame, const Splat* __restrict__, const float* __restrict__,
     GaussianGrads);
 template cudaError_t LaunchPreprocessBackward<float>(const GaussianParams&,
                                                      const Frame&, const Splat*,
                                                      const float*,
                                                      const GaussianGrads&,
                                                      cudaStream_t);
 template cudaError_t LaunchPreprocessBackward<ShBf16>(
     const GaussianParams&, const Frame&, const Splat*, const float*,
     const GaussianGrads&, cudaStream_t);
 }  // namespace gs

s01_literature → s02_reduce_scatter

`s01_literature` → `s02_reduce_scatter`