`s05_branchless_predication` → `s06_multirow_and_pipe_math`

Step 6: Multi-Row 8×8 Square Warps, Per-Warp Rings & Pipe-Balanced Math · +148 / −121 lines across 4 changed files

(a) 2-pixel/thread 8×8 square warps & per-warp __syncwarp() staging slices (csrc/forward.cu.cc, csrc/backward.cu.cc, csrc/rasterizer_impl.h): Switches to kDefaultForwardRows = 2, kDefaultBackwardRows = 2 (4 warps of 128 threads per 16×16 tile, each warp covering a square 8×8 sub-tile) and decouples the 4 warps in RasterBackwardKernel into independent 32-splat shared-memory slices (warp_base = warp * 32) synchronized with __syncwarp() and zero __syncthreads(); its occupancy target goes from __launch_bounds__(..., 4) to 6 blocks of 128 threads.
(b) Folded log₂(opacity), zero-multiply row moments & sign-bit transmittance termination (csrc/auxiliary.h, csrc/forward.cu.cc, csrc/backward.cu.cc): Folds log2f(opacity) into Power2Column::At (eliminating the per-pixel opacity multiply before MUFU.EX2), evaluates 2-row vertical moments with zero multiplies (m1 = ga[1], m2 = ga[1]), and encodes pixel termination in the sign bit of t[k].

Core files: csrc/auxiliary.h +12/−12 csrc/forward.cu.cc +51/−40 csrc/backward.cu.cc +83/−67 csrc/rasterizer_impl.h +2/−2

csrc/auxiliary.h +12 −12

s05_branchless_predication/csrc/auxiliary.h

s06_multirow_and_pipe_math/csrc/auxiliary.h

csrc/forward.cu.cc +51 −40

s05_branchless_predication/csrc/forward.cu.cc

s06_multirow_and_pipe_math/csrc/forward.cu.cc

csrc/backward.cu.cc +83 −67

s05_branchless_predication/csrc/backward.cu.cc

s06_multirow_and_pipe_math/csrc/backward.cu.cc

csrc/rasterizer_impl.h +2 −2

s05_branchless_predication/csrc/rasterizer_impl.h

s06_multirow_and_pipe_math/csrc/rasterizer_impl.h

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

↕ Expand 17 unchanged lines (lines 1–17)
1	/*	1	/*
2	* Copyright (C) 2023, Inria	2	* Copyright (C) 2023, Inria
3	* GRAPHDECO research group, https://team.inria.fr/graphdeco	3	* GRAPHDECO research group, https://team.inria.fr/graphdeco
4	* All rights reserved.	4	* All rights reserved.
5	*	5	*
6	* This software is free for non-commercial, research and evaluation use	6	* This software is free for non-commercial, research and evaluation use
7	* under the terms of the LICENSE.md file.	7	* under the terms of the LICENSE.md file.
8	*	8	*
9	* For inquiries contact george.drettakis@inria.fr	9	* For inquiries contact george.drettakis@inria.fr
10	*/	10	*/
11		11
12	#pragma once	12	#pragma once
13		13
14	#include <cstddef>	14	#include <cstddef>
15	#include <cstdint>	15	#include <cstdint>
16	#include <iterator>	16	#include <iterator>
17		17
18	#include "auxiliary.h"	18	#include "auxiliary.h"
19		19
20	namespace gs {	20	namespace gs {
21		21
22	inline constexpr int kDefaultForwardRows = 1;	22	inline constexpr int kDefaultForwardRows = 2;
23	inline constexpr int kDefaultBackwardRows = 1;	23	inline constexpr int kDefaultBackwardRows = 2;
24		24
25	// Number of bits needed to represent tile ids in [0, num_tiles).	25	// Number of bits needed to represent tile ids in [0, num_tiles).
26	inline int TileBits(int num_tiles) {	26	inline int TileBits(int num_tiles) {
27	int bits = 1;	27	int bits = 1;
↕ Expand 41 unchanged lines (lines 28–68) · in `TileRangesKernel`
28	while ((1 << bits) < num_tiles) ++bits;	28	while ((1 << bits) < num_tiles) ++bits;
29	return bits;	29	return bits;
30	}	30	}
31		31
32	struct GatherCountsIter {	32	struct GatherCountsIter {
33	const int32_t* counts;	33	const int32_t* counts;
34	const int32_t* order;	34	const int32_t* order;
35	__host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {	35	__host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
36	return counts[order[idx]];	36	return counts[order[idx]];
37	}	37	}
38	__host__ __device__ __forceinline__ GatherCountsIter	38	__host__ __device__ __forceinline__ GatherCountsIter
39	operator+(ptrdiff_t offset) const {	39	operator+(ptrdiff_t offset) const {
40	return {counts, order + offset};	40	return {counts, order + offset};
41	}	41	}
42	__host__ __device__ __forceinline__ int32_t operator*() const {	42	__host__ __device__ __forceinline__ int32_t operator*() const {
43	return counts[*order];	43	return counts[*order];
44	}	44	}
45	};	45	};
46		46
47	template <Culling kCulling>	47	template <Culling kCulling>
48	__global__ void EmitPairsKernel(const Splat* __restrict__ splats,	48	__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
49	const int32_t* __restrict__ sorted_ids,	49	const int32_t* __restrict__ sorted_ids,
50	const int32_t* __restrict__ offsets, int n_vis,	50	const int32_t* __restrict__ offsets, int n_vis,
51	Frame frame, uint16_t* __restrict__ keys,	51	Frame frame, uint16_t* __restrict__ keys,
52	int32_t* __restrict__ values);	52	int32_t* __restrict__ values);
53		53
54	__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,	54	__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
55	int num_pairs, int2* __restrict__ ranges);	55	int num_pairs, int2* __restrict__ ranges);
56		56
57	} // namespace gs	57	} // namespace gs
58		58
59	namespace std {	59	namespace std {
60	template <>	60	template <>
61	struct iterator_traits<gs::GatherCountsIter> {	61	struct iterator_traits<gs::GatherCountsIter> {
62	using difference_type = ptrdiff_t;	62	using difference_type = ptrdiff_t;
63	using value_type = int32_t;	63	using value_type = int32_t;
64	using pointer = const int32_t*;	64	using pointer = const int32_t*;
65	using reference = int32_t;	65	using reference = int32_t;
66	using iterator_category = random_access_iterator_tag;	66	using iterator_category = random_access_iterator_tag;
67	};	67	};
68	} // namespace std	68	} // namespace std

↕ Expand 17 unchanged lines (lines 1–17)
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		#pragma once
13	13
14	14		#include <cstddef>
15	15		#include <cstdint>
16	16		#include <iterator>
17	17
18	18		#include "auxiliary.h"
19	19
20	20		namespace gs {
21	21
22		−	inline constexpr int kDefaultForwardRows = 1;
	22	+	inline constexpr int kDefaultForwardRows = 2;
23		−	inline constexpr int kDefaultBackwardRows = 1;
	23	+	inline constexpr int kDefaultBackwardRows = 2;
24	24
25	25		// Number of bits needed to represent tile ids in [0, num_tiles).
26	26		inline int TileBits(int num_tiles) {
27	27		int bits = 1;
↕ Expand 41 unchanged lines (lines 28–68) · in `TileRangesKernel`
28	28		while ((1 << bits) < num_tiles) ++bits;
29	29		return bits;
30	30		}
31	31
32	32		struct GatherCountsIter {
33	33		const int32_t* counts;
34	34		const int32_t* order;
35	35		__host__ __device__ __forceinline__ int32_t operator[](ptrdiff_t idx) const {
36	36		return counts[order[idx]];
37	37		}
38	38		__host__ __device__ __forceinline__ GatherCountsIter
39	39		operator+(ptrdiff_t offset) const {
40	40		return {counts, order + offset};
41	41		}
42	42		__host__ __device__ __forceinline__ int32_t operator*() const {
43	43		return counts[*order];
44	44		}
45	45		};
46	46
47	47		template <Culling kCulling>
48	48		__global__ void EmitPairsKernel(const Splat* __restrict__ splats,
49	49		const int32_t* __restrict__ sorted_ids,
50	50		const int32_t* __restrict__ offsets, int n_vis,
51	51		Frame frame, uint16_t* __restrict__ keys,
52	52		int32_t* __restrict__ values);
53	53
54	54		__global__ void TileRangesKernel(const uint16_t* __restrict__ sorted_tiles,
55	55		int num_pairs, int2* __restrict__ ranges);
56	56
57	57		} // namespace gs
58	58
59	59		namespace std {
60	60		template <>
61	61		struct iterator_traits<gs::GatherCountsIter> {
62	62		using difference_type = ptrdiff_t;
63	63		using value_type = int32_t;
64	64		using pointer = const int32_t*;
65	65		using reference = int32_t;
66	66		using iterator_category = random_access_iterator_tag;
67	67		};
68	68		} // namespace std

 /*
  * 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, 0)
+//   p2 = (r, g, b, log2(opacity))
 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);
 }
 // ---------------------------------------------------------------------------
 // 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:
+// with (A, B, C) = -log2(e)/2 (a, 2b, c), so that the Gaussian is 2^power2.
+// log2(opacity) (stored in `Splat::p2.w`) is folded into the per-column
+// constant, so 2^power is directly the unclamped alpha = opacity * Gaussian:
-// two FMAs and one MUFU.EX2 per pixel. A pixel receives the splat iff
+// two FMAs and one MUFU.EX2 per pixel, no multiply. A pixel receives the
-// power2 <= 0 and alpha >= kMinAlpha (the latter implies power2 >= the
+// splat iff alpha >= kMinAlpha (which implies power2 >= the splat's culling
-// splat's culling bound, see `AlphaLimit`).
+// 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.
+  float pa, pb;  // A dx^2 + log2(opacity), B dx.
-  GS_D static Power2Column At(float a2, float b2, float dx) {
+  GS_D static Power2Column At(float a2, float b2, float log2_opacity,
+                              float dx) {
-    return {__fmul_rn(__fmul_rn(a2, dx), dx), __fmul_rn(b2, dx)};
+    return {__fmaf_rn(__fmul_rn(a2, dx), dx, log2_opacity), __fmul_rn(b2, dx)};
   }
+  // log2 of the unclamped alpha at dy = splat.y - pixel.y.
   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) {
     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, 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));
+  // Visible splats have opacity >= ~1/255 (see `AlphaLimit`).
-  s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
+  s.p2 = {rgb.x, rgb.y, rgb.z, log2f(s.p1.y)};
   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];
+  // t[k] >= 0: live pixel with transmittance t. t[k] < 0: the pixel is done
+  // (transmittance would drop below kMinTransmittance) and |t| is final.
+  // Done pixels then fail `next_t >= kMinTransmittance`, so they need no
+  // separate flag.
   float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
   int last[kRasterRows];
-  bool done[kRasterRows];
 #pragma unroll
   for (int k = 0; k < kRasterRows; ++k) {
+    const bool outside = px.x >= frame.width || px.y0 + k >= frame.height;
-    t[k] = 1.0f;
+    t[k] = outside ? -1.0f : 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;
+    uint32_t sign = 0xffffffffu;
 #pragma unroll
-    for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
+    for (int k = 0; k < kRasterRows; ++k) sign &= __float_as_uint(t[k]);
-    return d;
+    return (sign >> 31) != 0u;
   };
-  __shared__ float4 batch_geom0[kRasterThreads];
-  __shared__ float batch_geom1[kRasterThreads];
+  // Staged splat j: batch[0][j] = (x, y, A, B), batch[1][j] = (r, g, b,
+  // log2(opacity)), batch[2][j].x = C (see `Power2Conic`). One array so all
+  // three loads share an address.
-  __shared__ float4 batch_rgb_o[kRasterThreads];
+  __shared__ float4 batch[3][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[0][tid] = {p0.x, p0.y, c2.x, c2.y};
-      batch_geom1[tid] = c2.z;
+      batch[1][tid] = p2;
-      batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
+      batch[2][tid].x = c2.z;
       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) {
+      // Bit `j - c` is set if some pixel of this lane blended splat j.
-      uint32_t contrib = 0u;
+      uint32_t lane_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;
+          const uint32_t low = mask & (0u - mask);
-          mask &= mask - 1u;
+          mask ^= low;
-          const int j = c + bit;
+          const int j = c + (31 - __clz(low));
+          // Predicated like the backward pass: rows that are done or not
+          // reached keep their state.
+          const float4 g0 = batch[0][j];
+          const float4 ro = batch[1][j];
+          const float c2 = batch[2][j].x;
+          const Power2Column col =
+              Power2Column::At(g0.z, g0.w, ro.w, __fsub_rn(g0.x, fx));
           bool blended = false;
-          {
-            // 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 float c2 = batch_geom1[j];
-            const float4 ro = batch_rgb_o[j];
-            const Power2Column col =
-                Power2Column::At(g0.z, g0.w, __fsub_rn(g0.x, fx));
 #pragma unroll
-            for (int k = 0; k < kRasterRows; ++k) {
+          for (int k = 0; k < kRasterRows; ++k) {
-              const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
+            const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
-              const float power = col.Row(c2, dy);
+            const float alpha_raw = Exp2Approx(col.Row(c2, dy));
-              const float alpha_k = SplatAlpha(ro.w, Exp2Approx(power));
+            const float alpha_k = fminf(kMaxAlpha, alpha_raw);
-              const bool hit =
-                  !done[k] && power <= 0.0f && alpha_k >= kMinAlpha;
+            const bool hit = alpha_raw >= kMinAlpha;
+            // w = alpha t and next_t = t (1 - alpha) = t - alpha t.
+            const float w_k = __fmul_rn(alpha_k, t[k]);
-              const float next_t = t[k] * (1.0f - alpha_k);
+            const float next_t = __fmaf_rn(-alpha_k, t[k], t[k]);
-              const bool stop = hit && next_t < kMinTransmittance;
+            const bool blend = hit && next_t >= kMinTransmittance;
-              const bool blend = hit && !stop;
-              done[k] = done[k] || stop;
-              const float w = blend ? alpha_k * t[k] : 0.0f;
+            const float w = blend ? w_k : 0.0f;
-              r[k] += w * ro.x;
+            r[k] += w * ro.x;
-              g[k] += w * ro.y;
+            g[k] += w * ro.y;
-              b[k] += w * ro.z;
+            b[k] += w * ro.z;
+            // A hit that is not blended stops the pixel (t -> -|t|).
-              t[k] = blend ? next_t : t[k];
+            t[k] = blend ? next_t : (hit ? -fabsf(t[k]) : t[k]);
-              last[k] = blend ? -1 - j : last[k];
+            last[k] = blend ? -1 - j : last[k];
-              blended = blended || blend;
+            blended = blended || blend;
-            }
+          }
+          lane_contrib |= blended ? low : 0u;
-            thread_done = all_done();
+          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 defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+      const uint32_t contrib = __reduce_or_sync(warp::kFullMask, lane_contrib);
+#else
+      uint32_t contrib = lane_contrib;
+#pragma unroll
+      for (int offset = 16; offset > 0; offset /= 2) {
+        contrib |= __shfl_xor_sync(warp::kFullMask, contrib, offset);
+      }
+#endif
       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];
+    final_t[p] = fabsf(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 93 unchanged lines (lines 1–93) · 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	81		s.p0 = {pr.xy.x, pr.xy.y, pr.conic.x, pr.conic.y};
82	82		s.p1 = {pr.conic.z, pr.opacity, pr.limit, static_cast<float>(pr.radius)};
83	83		s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
84	84		ForEachTileRow<kCulling>(
85	85		s, frame.grid, frame.width, frame.height,
86	86		[&](int, int tx0, int tx1) { count += tx1 - tx0; });
87	87		}
88	88		if (count == 0) {
89	89		s.p0 = s.p1 = s.p2 = {0.0f, 0.0f, 0.0f, 0.0f};
90	90		splats[i] = s;
91	91		tile_counts[i] = 0;
92	92		return;
93	93		}
94	94		const int vis_idx = atomicAdd(num_visible, 1);
95	95		float inv_len;
96	96		const float3 rgb =
97	97		ShadeColorRow(ShRow<ShT>(params.sh, i), ViewDir(cam, mean, inv_len));
	98	+	// Visible splats have opacity >= ~1/255 (see `AlphaLimit`).
98		−	s.p2 = {rgb.x, rgb.y, rgb.z, s.p2.w};
	99	+	s.p2 = {rgb.x, rgb.y, rgb.z, log2f(s.p1.y)};
99	100		splats[i] = s;
100	101		tile_counts[i] = count;
101	102		depth_keys[vis_idx] = AsUint(pr.depth);
102	103		ids[vis_idx] = i;
↕ Expand 16 unchanged lines (lines 104–119) · in `RasterForwardKernel`
103	104		}
104	105
105	106		template <int kRows>
106	107		__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 8)
107	108		RasterForwardKernel(const int2* __restrict__ ranges,
108	109		int32_t* __restrict__ point_list,
109	110		const Splat* __restrict__ splats, Frame frame,
110	111		float* __restrict__ image, float* __restrict__ final_t,
111	112		int32_t* __restrict__ n_contrib) {
112	113		constexpr int kRasterRows = kRows;
113	114		constexpr int kRasterThreads = RasterLayout<kRows>::kThreads;
114	115		constexpr int kWarps = RasterLayout<kRows>::kWarps;
115	116		const int tile = blockIdx.y * frame.grid.tiles_x + blockIdx.x;
116	117		const int tid = threadIdx.x;
117	118		const int warp = tid / 32;
118	119		const int lane = tid % 32;
119	120		const RasterPixels px = RasterPixels::For<kRows>(tid);
120	121		const float fx = static_cast<float>(px.x);
121	122		const int2 range = ranges[tile];
122	123
	124	+	// t[k] >= 0: live pixel with transmittance t. t[k] < 0: the pixel is done
	125	+	// (transmittance would drop below kMinTransmittance) and \|t\| is final.
	126	+	// Done pixels then fail `next_t >= kMinTransmittance`, so they need no
	127	+	// separate flag.
123	128		float t[kRasterRows], r[kRasterRows], g[kRasterRows], b[kRasterRows];
124	129		int last[kRasterRows];
125		−	bool done[kRasterRows];
126	130		#pragma unroll
127	131		for (int k = 0; k < kRasterRows; ++k) {
	132	+	const bool outside = px.x >= frame.width \|\| px.y0 + k >= frame.height;
128		−	t[k] = 1.0f;
	133	+	t[k] = outside ? -1.0f : 1.0f;
129	134		r[k] = g[k] = b[k] = 0.0f;
130	135		last[k] = 0;
131		−	done[k] = px.x >= frame.width \|\| px.y0 + k >= frame.height;
132	136		}
133	137		auto all_done = [&] {
134		−	bool d = true;
	138	+	uint32_t sign = 0xffffffffu;
135	139		#pragma unroll
136		−	for (int k = 0; k < kRasterRows; ++k) d = d && done[k];
	140	+	for (int k = 0; k < kRasterRows; ++k) sign &= __float_as_uint(t[k]);
137		−	return d;
	141	+	return (sign >> 31) != 0u;
138	142		};
139	143
140		−	__shared__ float4 batch_geom0[kRasterThreads];
141		−	__shared__ float batch_geom1[kRasterThreads];
	144	+	// Staged splat j: batch[0][j] = (x, y, A, B), batch[1][j] = (r, g, b,
	145	+	// log2(opacity)), batch[2][j].x = C (see `Power2Conic`). One array so all
	146	+	// three loads share an address.
142		−	__shared__ float4 batch_rgb_o[kRasterThreads];
	147	+	__shared__ float4 batch[3][kRasterThreads];
143	148		__shared__ uint8_t batch_mask[kRasterThreads];
144	149		__shared__ uint32_t warp_contrib[kWarps][kWarps];
145	150		const uint8_t warp_bit = static_cast<uint8_t>(1u << warp);
146	151		bool thread_done = all_done();
↕ Expand 6 unchanged lines (lines 152–157) · in `RasterForwardKernel`
147	152		int tile_compact = 0;
148	153		for (int start = range.x; start < range.y; start += kRasterThreads) {
149	154		const int idx = start + tid;
150	155		int32_t id = 0;
151	156		if (idx < range.y) {
152	157		id = point_list[idx];
153	158		const float4 p0 = __ldg(&splats[id].p0);
154	159		const float4 p1 = __ldg(&splats[id].p1);
155	160		const float4 p2 = __ldg(&splats[id].p2);
156	161		const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
157		−	batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
	162	+	batch[0][tid] = {p0.x, p0.y, c2.x, c2.y};
158		−	batch_geom1[tid] = c2.z;
	163	+	batch[1][tid] = p2;
159		−	batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
	164	+	batch[2][tid].x = c2.z;
160	165		batch_mask[tid] = static_cast<uint8_t>(WarpMask<kRows>(p0, p1));
161	166		} else {
162	167		batch_mask[tid] = 0u;
163	168		}
164	169		__syncthreads();
165	170		const int count = min(kRasterThreads, range.y - start);
166	171		bool warp_done = __all_sync(warp::kFullMask, thread_done);
167	172		#pragma unroll
168	173		for (int c = 0; c < kRasterThreads; c += 32) {
	174	+	// Bit `j - c` is set if some pixel of this lane blended splat j.
169		−	uint32_t contrib = 0u;
	175	+	uint32_t lane_contrib = 0u;
170	176		if (c < count && !warp_done) {
171	177		uint32_t mask = __ballot_sync(
172	178		warp::kFullMask,
173	179		c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
174	180		while (mask != 0u) {
175		−	const int bit = __ffs(mask) - 1;
	181	+	const uint32_t low = mask & (0u - mask);
176		−	mask &= mask - 1u;
	182	+	mask ^= low;
177		−	const int j = c + bit;
	183	+	const int j = c + (31 - __clz(low));
	184	+	// Predicated like the backward pass: rows that are done or not
	185	+	// reached keep their state.
	186	+	const float4 g0 = batch[0][j];
	187	+	const float4 ro = batch[1][j];
	188	+	const float c2 = batch[2][j].x;
	189	+	const Power2Column col =
	190	+	Power2Column::At(g0.z, g0.w, ro.w, __fsub_rn(g0.x, fx));
178	191		bool blended = false;
179		−	{
180		−	// Predicated like the backward pass: rows that are done or not
181		−	// reached get alpha = 0 and keep their state.
182		−	const float4 g0 = batch_geom0[j];
183		−	const float c2 = batch_geom1[j];
184		−	const float4 ro = batch_rgb_o[j];
185		−	const Power2Column col =
186		−	Power2Column::At(g0.z, g0.w, __fsub_rn(g0.x, fx));
187	192		#pragma unroll
188	193	⇥	for (int k = 0; k < kRasterRows; ++k) {
189	194	⇥	const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
190		−	const float power = col.Row(c2, dy);
	195	+	const float alpha_raw = Exp2Approx(col.Row(c2, dy));
191		−	const float alpha_k = SplatAlpha(ro.w, Exp2Approx(power));
	196	+	const float alpha_k = fminf(kMaxAlpha, alpha_raw);
192		−	const bool hit =
193		−	!done[k] && power <= 0.0f && alpha_k >= kMinAlpha;
	197	+	const bool hit = alpha_raw >= kMinAlpha;
	198	+	// w = alpha t and next_t = t (1 - alpha) = t - alpha t.
	199	+	const float w_k = __fmul_rn(alpha_k, t[k]);
194		−	const float next_t = t[k] * (1.0f - alpha_k);
	200	+	const float next_t = __fmaf_rn(-alpha_k, t[k], t[k]);
195		−	const bool stop = hit && next_t < kMinTransmittance;
	201	+	const bool blend = hit && next_t >= kMinTransmittance;
196		−	const bool blend = hit && !stop;
197		−	done[k] = done[k] \|\| stop;
198		−	const float w = blend ? alpha_k * t[k] : 0.0f;
	202	+	const float w = blend ? w_k : 0.0f;
199	203	⇥	r[k] += w * ro.x;
200	204	⇥	g[k] += w * ro.y;
201	205	⇥	b[k] += w * ro.z;
	206	+	// A hit that is not blended stops the pixel (t -> -\|t\|).
202		−	t[k] = blend ? next_t : t[k];
	207	+	t[k] = blend ? next_t : (hit ? -fabsf(t[k]) : t[k]);
203	208	⇥	last[k] = blend ? -1 - j : last[k];
204	209	⇥	blended = blended \|\| blend;
205	210	⇥	}
	211	+	lane_contrib \|= blended ? low : 0u;
206	212	⇥	thread_done = all_done();
207		−	}
208		−	if (__any_sync(warp::kFullMask, blended)) {
209		−	contrib \|= (1u << bit);
210		−	}
211	213		if (__all_sync(warp::kFullMask, thread_done)) {
212	214		warp_done = true;
213	215		break;
214	216		}
215	217		}
216	218		warp_done = __all_sync(warp::kFullMask, thread_done);
217	219		}
	220	+	#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
	221	+	const uint32_t contrib = __reduce_or_sync(warp::kFullMask, lane_contrib);
	222	+	#else
	223	+	uint32_t contrib = lane_contrib;
	224	+	#pragma unroll
	225	+	for (int offset = 16; offset > 0; offset /= 2) {
	226	+	contrib \|= __shfl_xor_sync(warp::kFullMask, contrib, offset);
	227	+	}
	228	+	#endif
218	229		if (lane == 0) {
219	230		warp_contrib[warp][c / 32] = contrib;
220	231		}
221	232		}
↕ Expand 45 unchanged lines (lines 233–277) · in `RasterForwardKernel`
222	233		const bool block_done =
223	234		(__syncthreads_count(thread_done) == kRasterThreads);
224	235		uint32_t chunk_any[kWarps];
225	236		int chunk_prefix[kWarps];
226	237		int batch_active = 0;
227	238		#pragma unroll
228	239		for (int u = 0; u < kWarps; ++u) {
229	240		uint32_t any_u = 0u;
230	241		#pragma unroll
231	242		for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
232	243		any_u \|= warp_contrib[w_idx][u];
233	244		}
234	245		chunk_any[u] = any_u;
235	246		chunk_prefix[u] = tile_compact + batch_active;
236	247		batch_active += __popc(any_u);
237	248		}
238	249		if ((chunk_any[warp] >> lane) & 1u) {
239	250		uint32_t cmask = 0u;
240	251		#pragma unroll
241	252		for (int w_idx = 0; w_idx < kWarps; ++w_idx) {
242	253		cmask \|= ((warp_contrib[w_idx][warp] >> lane) & 1u) << w_idx;
243	254		}
244	255		const int dst =
245	256		chunk_prefix[warp] + __popc(chunk_any[warp] & ((1u << lane) - 1u));
246	257		point_list[range.x + dst] =
247	258		static_cast<int32_t>((cmask << 24) \| static_cast<uint32_t>(id));
248	259		}
249	260		#pragma unroll
250	261		for (int k = 0; k < kRasterRows; ++k) {
251	262		if (last[k] < 0) {
252	263		const int j = -1 - last[k];
253	264		const int u = j >> 5;
254	265		const int b_idx = j & 31;
255	266		last[k] =
256	267		chunk_prefix[u] + __popc(chunk_any[u] & ((1u << b_idx) - 1u)) + 1;
257	268		}
258	269		}
259	270		tile_compact += batch_active;
260	271		if (block_done) break;
261	272		}
262	273		const int pixels = frame.width * frame.height;
263	274		#pragma unroll
264	275		for (int k = 0; k < kRasterRows; ++k) {
265	276		const int y = px.y0 + k;
266	277		if (px.x >= frame.width \|\| y >= frame.height) continue;
267	278		const int p = y * frame.width + px.x;
268	279		image[0 * pixels + p] = r[k];
269	280		image[1 * pixels + p] = g[k];
270	281		image[2 * pixels + p] = b[k];
271		−	final_t[p] = t[k];
	282	+	final_t[p] = fabsf(t[k]);
272	283		n_contrib[p] = last[k];
273	284		}
274	285		}
275	286
↕ Expand 24 unchanged lines (lines 287–310) · in `RasterForwardKernel`
276	287		template __global__ void PreprocessKernel<Culling::kRect, float>(
277	288		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
278	289		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
279	290		template __global__ void PreprocessKernel<Culling::kRect, ShBf16>(
280	291		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
281	292		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
282	293		template __global__ void PreprocessKernel<Culling::kEllipse, float>(
283	294		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
284	295		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
285	296		template __global__ void PreprocessKernel<Culling::kEllipse, ShBf16>(
286	297		GaussianParams, Frame, Splat* __restrict__, int32_t* __restrict__,
287	298		uint32_t* __restrict__, int32_t* __restrict__, int32_t* __restrict__);
288	299
289	300		template __global__ void RasterForwardKernel<1>(
290	301		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
291	302		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
292	303		template __global__ void RasterForwardKernel<2>(
293	304		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
294	305		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
295	306		template __global__ void RasterForwardKernel<4>(
296	307		const int2* __restrict__, int32_t* __restrict__, const Splat* __restrict__,
297	308		Frame, float* __restrict__, float* __restrict__, int32_t* __restrict__);
298	309
299	310		} // 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)
+__global__ void __launch_bounds__(RasterLayout<kRows>::kThreads, 6)
     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;
-  }
+  // Each warp stages its own splats in its 32-entry slice of these arrays,
+  // so warps never wait for each other (no block barriers).
   __shared__ float4 batch_geom0[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);
+  const uint32_t warp_bit = 1u << warp;
+  const int warp_base = warp * 32;
   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.
+  //
+  // Per pixel, ga = alpha_raw T diff = dL/dpower (log2 units, up to the
+  // constant), opacity included: the opacity gradient divides it back out.
+  // Returns ga_sum = sum ga, ga_dy = sum ga dy, ga_dy2 = sum ga dy^2. Since the
+  // rows are dy_k = dy_0 - k, the last two are formed from the moments
+  // m1 = sum k ga_k and m2 = sum k^2 ga_k: sum ga dy = dy_0 ga_sum - m1 and
+  // sum ga dy^2 = dy_0 (dy_0 ga_sum - 2 m1) + m2.
   auto eval_splat = [&](float4 g0, float c2, float4 ro, int entry, float& dx,
                         float& ga_sum, float& ga_dy, float& ga_dy2,
                         float3& drgb) {
     dx = __fsub_rn(g0.x, fx);
-    const Power2Column col = Power2Column::At(g0.z, g0.w, dx);
+    const Power2Column col = Power2Column::At(g0.z, g0.w, ro.w, dx);
-    ga_sum = 0.0f;
+    float ga[kRasterRows];
-    ga_dy = 0.0f;
+    float dy0 = 0.0f;
-    ga_dy2 = 0.0f;
     drgb = {0.0f, 0.0f, 0.0f};
 #pragma unroll
     for (int k = 0; k < kRasterRows; ++k) {
       PixelState& ps = st[k];
       const float dy = __fsub_rn(g0.y, static_cast<float>(px.y0 + k));
+      if (k == 0) dy0 = dy;
-      const float power = col.Row(c2, dy);
+      const float alpha_raw = Exp2Approx(col.Row(c2, dy));
-      const float g_val = Exp2Approx(power);
-      const float alpha_k = SplatAlpha(ro.w, g_val);
-      const bool hit = entry < ps.last && power <= 0.0f && alpha_k >= kMinAlpha;
+      const bool hit = entry < ps.last && alpha_raw >= kMinAlpha;
-      const float alpha = hit ? alpha_k : 0.0f;
+      const float alpha = hit ? fminf(kMaxAlpha, alpha_raw) : 0.0f;
       ps.t = ps.t * RcpApprox(1.0f - alpha);
       const float w = alpha * ps.t;
       const float color_dot =
           ro.x * ps.dpix.x + ro.y * ps.dpix.y + ro.z * ps.dpix.z;
       const float diff = color_dot - ps.accum;
       ps.accum += alpha * diff;
-      const float ga = hit ? g_val * (ps.t * diff) : 0.0f;
+      ga[k] = hit ? alpha_raw * (ps.t * diff) : 0.0f;
-      ga_sum += ga;
-      const float ga_y = ga * dy;
-      ga_dy += ga_y;
-      ga_dy2 += ga_y * dy;
       drgb += w * ps.dpix;
+    }
+    ga_sum = ga[0];
+#pragma unroll
+    for (int k = 1; k < kRasterRows; ++k) ga_sum += ga[k];
+    if constexpr (kRasterRows == 1) {
+      ga_dy = ga_sum * dy0;
+      ga_dy2 = ga_dy * dy0;
+    } else {
+      float m1 = ga[1], m2 = ga[1];
+#pragma unroll
+      for (int k = 2; k < kRasterRows; ++k) {
+        m1 = fmaf(static_cast<float>(k), ga[k], m1);
+        m2 = fmaf(static_cast<float>(k * k), ga[k], m2);
+      }
+      ga_dy = fmaf(dy0, ga_sum, -m1);
+      ga_dy2 = fmaf(dy0, ga_dy - m1, m2);
     }
   };
   // Conic (a, b, c) from the staged log2-scaled one.
   auto raw_conic = [](float4 g0, float c2) {
     return float3{kInvPower2Scale * g0.z, 0.5f * kInvPower2Scale * g0.w,
                   kInvPower2Scale * c2};
   };
-  for (int end = block_max_last; end > 0; end -= kRasterThreads) {
+  for (int end = warp_max_last; end > 0; end -= 32) {
+    // Lane l stages entry end - 1 - l if this warp contributed to it.
-    const int e = end - 1 - tid;
+    const int e = end - 1 - lane;
+    bool mine = false;
     if (e >= 0) {
       const uint32_t packed = static_cast<uint32_t>(point_list[range.x + e]);
+      mine = ((packed >> 24) & warp_bit) != 0u;
+      if (mine) {
-      const uint32_t id = packed & 0x00ffffffu;
+        const uint32_t id = packed & 0x00ffffffu;
-      splat_touched[id] = 1u;
+        splat_touched[id] = 1u;
-      const float4 p0 = __ldg(&splats[id].p0);
+        const float4 p0 = __ldg(&splats[id].p0);
-      const float4 p1 = __ldg(&splats[id].p1);
+        const float4 p1 = __ldg(&splats[id].p1);
-      const float4 p2 = __ldg(&splats[id].p2);
+        const float4 p2 = __ldg(&splats[id].p2);
-      const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
+        const float3 c2 = Power2Conic(p0.z, p0.w, p1.x);
-      batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
+        batch_geom0[tid] = {p0.x, p0.y, c2.x, c2.y};
-      batch_geom1[tid] = c2.z;
+        batch_geom1[tid] = c2.z;
-      batch_rgb_o[tid] = {p2.x, p2.y, p2.z, p1.y};
+        batch_rgb_o[tid] = p2;
-      batch_id[tid] = id;
+        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(
+    uint32_t mask = __ballot_sync(warp::kFullMask, mine);
-            warp::kFullMask,
-            c + lane < count && (batch_mask[c + lane] & warp_bit) != 0u);
+    __syncwarp();
+    {
+      {
         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 int j0 = warp_base + bit0, j1 = warp_base + bit1;
             const float4 g0_0 = batch_geom0[j0];
             const float c2_0 = batch_geom1[j0];
             const float4 ro_0 = batch_rgb_o[j0];
             const float4 g0_1 = batch_geom0[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(g0_0, c2_0, ro_0, end - 1 - j0, dx0, ga_sum0, ga_dy0,
+            eval_splat(g0_0, c2_0, ro_0, end - 1 - bit0, dx0, ga_sum0, ga_dy0,
                        ga_dy2_0, drgb0);
             float dx1, ga_sum1, ga_dy1, ga_dy2_1;
             float3 drgb1;
-            eval_splat(g0_1, c2_1, ro_1, end - 1 - j1, dx1, ga_sum1, ga_dy1,
+            eval_splat(g0_1, c2_1, ro_1, end - 1 - bit1, dx1, ga_sum1, ga_dy1,
                        ga_dy2_1, drgb1);
             const bool hi16 = (lane & 16) != 0;
+            const bool row1 = (lane & 8) != 0;
-            const float opacity = hi16 ? ro_1.w : ro_0.w;
+            const float log2_opacity = hi16 ? ro_1.w : ro_0.w;
             const float dx = hi16 ? dx1 : dx0;
+            // This lane's row of the raw conic [[a, b], [b, c]] times -1:
+            // lanes 0-7 hold (a, b), lanes 8-15 (b, c), from the staged
+            // (A, B, C) = kPower2Scale (a, 2b, c). Selecting before scaling
+            // leaves two multiplies.
+            const float4 g0_sel = hi16 ? g0_1 : g0_0;
+            const float c2_sel = hi16 ? c2_1 : c2_0;
-            const float3 cn =
+            const float cn_x =
-                hi16 ? raw_conic(g0_1, c2_1) : raw_conic(g0_0, c2_0);
+                (row1 ? g0_sel.w : g0_sel.z) *
+                (row1 ? -0.5f * kInvPower2Scale : -kInvPower2Scale);
+            const float cn_y =
+                (row1 ? c2_sel : g0_sel.w) *
+                (row1 ? -kInvPower2Scale : -0.5f * kInvPower2Scale);
             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),
+            float b4[4] = {-0.5f * 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_sum = ga_sum, p_dy = ga_dy;
             const float p_dx = dx * p_sum;
-            const float v_dx0 = (lane & 8) == 0 ? -(cn.x * p_dx + cn.y * p_dy)
+            const float v_dx0 = fmaf(cn_x, p_dx, cn_y * p_dy);
-                                                : -(cn.y * p_dx + cn.z * p_dy);
-            const float v_dx1 =
+            const float v_dx1 = dx * (row1 ? -p_dy : -0.5f * p_dx);
-                (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 =
-            const float out_grad = ((lane & 7) == 1) ? ga_sum : v_sum;
+                ((lane & 7) == 1) ? ga_sum * Exp2Approx(-log2_opacity) : 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 j = warp_base + bit;
           const float4 g0 = batch_geom0[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(g0, c2, ro, end - 1 - j, dx, ga_sum, ga_dy, ga_dy2, drgb);
+          eval_splat(g0, c2, ro, end - 1 - bit, dx, ga_sum, ga_dy, ga_dy2,
+                     drgb);
-          const float opacity = ro.w;
+          const float inv_opacity = Exp2Approx(-ro.w);
           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};
+            float b4[4] = {-0.5f * 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_sum = ga_sum, p_dy = 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};
+            float v[3] = {v_dx, b4[0], ga_sum * inv_opacity};
             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_sum = ga_sum, p_dy = 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,
+                                          -0.5f * ga_dy2,
-                                          ga_sum,  // opacity
+                                          ga_sum * inv_opacity,  // 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();
+    __syncwarp();  // The slice is restaged next iteration.
   }
 }
 // ---------------------------------------------------------------------------
 // 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).
+// zeros untouched SH rows with warp-coalesced 16-byte 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

s05_branchless_predication → s06_multirow_and_pipe_math

`s05_branchless_predication` → `s06_multirow_and_pipe_math`