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

inria → s00_baseline

Inria's Reference diff-gaussian-rasterization  ·  +510 / −196 lines in 1 harness wrapper file (all 9 core CUDA kernel files unchanged)

Changed files: csrc/ffi.cu.cc +510/−196
csrc/ffi.cu.cc +510 −196
inria/rasterize_points.cu.cc
s00_baseline/csrc/ffi.cu.cc
1/*
2 * Copyright (C) 2023, Inria
3 * GRAPHDECO research group, https://team.inria.fr/graphdeco
4 * All rights reserved.
5 *
6 * This software is free for non-commercial, research and evaluation use
7 * under the terms of the LICENSE.md file.
8 *
9 * For inquiries contact george.drettakis@inria.fr
10 */
1// XLA FFI entry points driving the verbatim Inria `diff-gaussian-rasterization`
2// CUDA kernels (`forward.cu.cc`, `backward.cu.cc`, `rasterizer_impl.cu.cc`).
3#include "xla/ffi/api/ffi.h"
114
12#include <math.h>5#include <cuda_runtime.h>
13#include <torch/extension.h>
14#include <cstdio>
6
15#include <sstream>7#include <array>
16#include <iostream>8#include <cfloat>
17#include <tuple>9#include <cstddef>
18#include <stdio.h>10#include <cstdint>
19#include <cuda_runtime_api.h>11#include <cub/cub.cuh>
20#include <memory>
21#include "cuda_rasterizer/config.h"
22#include "cuda_rasterizer/rasterizer.h"
23#include <fstream>12#include <optional>
24#include <string>13#include <string>
25#include <functional>
2614
27std::function<char*(size_t N)> resizeFunctional(torch::Tensor& t) {
28 auto lambda = [&t](size_t N) {
29 t.resize_({(long long)N});
30 return reinterpret_cast<char*>(t.contiguous().data_ptr());
31 };
32 return lambda;
15#include "csrc/backward.h"
16#include "csrc/config.h"
17#include "csrc/forward.h"
18#include "csrc/rasterizer_impl.h"
19#include "xla/ffi/api/c_api.h"
20
21// Non-static helpers defined in `rasterizer_impl.cu.cc`.
22uint32_t getHigherMsb(uint32_t n);
23__global__ void duplicateWithKeys(int P, const float2* points_xy,
24 const float* depths, const uint32_t* offsets,
25 uint64_t* gaussian_keys_unsorted,
26 uint32_t* gaussian_values_unsorted,
27 int* radii, dim3 grid, int2* rects);
28__global__ void identifyTileRanges(int L, uint64_t* point_list_keys,
29 uint2* ranges);
30
31namespace gs {
32namespace {
33
34namespace ffi = ::xla::ffi;
35
36constexpr int kThreads = 256;
37constexpr int kMaxTimings = 8;
38constexpr int kCameraFloats = 20;
39constexpr int kShDegree = 3;
40constexpr int kShCoeffs = 16;
41constexpr float kScaleModifier = 1.0f;
42constexpr float kLowPassFilter = 0.3f;
43
44int Blocks(int64_t n) {
45 return static_cast<int>((n + kThreads - 1) / kThreads);
33}46}
3447
35std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
36RasterizeGaussiansCUDA(
37 const torch::Tensor& background,
38 const torch::Tensor& means3D,
39 const torch::Tensor& colors,
40 const torch::Tensor& opacity,
41 const torch::Tensor& scales,
42 const torch::Tensor& rotations,
43 const float scale_modifier,
44 const torch::Tensor& cov3D_precomp,
45 const torch::Tensor& viewmatrix,
46 const torch::Tensor& projmatrix,
47 const float tan_fovx,
48 const float tan_fovy,
49 const float low_pass_2d_kernel_size,
50 const bool mip_splatting_enabled,
51 const int image_height,
52 const int image_width,
53 const torch::Tensor& sh,
54 const int degree,
55 const torch::Tensor& campos,
56 const bool prefiltered)
57{
58 if (means3D.ndimension() != 2 || means3D.size(1) != 3) {
59 AT_ERROR("means3D must have dimensions (num_points, 3)");
48ffi::Error CudaError(cudaError_t e, const char* what) {
49 return ffi::Error(ffi::ErrorCode::kInternal,
50 std::string(what) + ": " + cudaGetErrorString(e));
60 }51}
61 52
62 const int P = means3D.size(0);
63 const int H = image_height;
64 const int W = image_width;
53#define GS_RETURN_IF_CUDA_ERROR(expr, what) \
54 do { \
55 cudaError_t gs_err_ = (expr); \
56 if (gs_err_ != cudaSuccess) return CudaError(gs_err_, what); \
57 } while (0)
6558
66 auto int_opts = means3D.options().dtype(torch::kInt32);
67 auto float_opts = means3D.options().dtype(torch::kFloat32);
68
69 torch::Tensor out_color = torch::full({NUM_CHANNELS, H, W}, 0.0, float_opts);
70 torch::Tensor radii = torch::full({P}, 0, means3D.options().dtype(torch::kInt32));
71
72 torch::Device device(torch::kCUDA);
73 torch::TensorOptions options(torch::kByte);
74 torch::Tensor geomBuffer = torch::empty({0}, options.device(device));
75 torch::Tensor binningBuffer = torch::empty({0}, options.device(device));
76 torch::Tensor imgBuffer = torch::empty({0}, options.device(device));
77 std::function<char*(size_t)> geomFunc = resizeFunctional(geomBuffer);
78 std::function<char*(size_t)> binningFunc = resizeFunctional(binningBuffer);
79 std::function<char*(size_t)> imgFunc = resizeFunctional(imgBuffer);
80
81 int rendered = 0;
82 if(P != 0)
83 {
84 int M = 0;
85 if(sh.size(0) != 0)
86 {
87 M = sh.size(1);
59template <typename T>
60std::optional<T*> Alloc(ffi::ScratchAllocator& scratch, size_t count) {
61 auto p = scratch.Allocate(sizeof(T) * (count > 0 ? count : 1), 16);
62 if (!p.has_value()) return std::nullopt;
63 return static_cast<T*>(*p);
88 }64}
8965
90 rendered = CudaRasterizer::Rasterizer::forward(
91 geomFunc,
92 binningFunc,
93 imgFunc,
94 P, degree, M,
95 background.contiguous().data<float>(),
96 W, H,
97 means3D.contiguous().data<float>(),
98 sh.contiguous().data_ptr<float>(),
99 colors.contiguous().data<float>(),
100 opacity.contiguous().data<float>(),
101 scales.contiguous().data_ptr<float>(),
102 scale_modifier,
103 rotations.contiguous().data_ptr<float>(),
104 cov3D_precomp.contiguous().data<float>(),
105 viewmatrix.contiguous().data<float>(),
106 projmatrix.contiguous().data<float>(),
107 campos.contiguous().data<float>(),
108 tan_fovx,
109 tan_fovy,
110 low_pass_2d_kernel_size,
111 mip_splatting_enabled,
112 prefiltered,
113 out_color.contiguous().data<float>(),
114 radii.contiguous().data<int>());
66// Bridges XLA's non-blocking `stream` to the default stream (`nullptr`) used by
67// Inria's unmodified `FORWARD::*` and `BACKWARD::*` kernel launches.
68class DefaultStreamBridge {
69 public:
70 explicit DefaultStreamBridge(cudaStream_t xla_stream)
71 : xla_stream_(xla_stream) {
72 cudaEventCreateWithFlags(&ev_, cudaEventDisableTiming);
73 cudaEventRecord(ev_, xla_stream_);
74 cudaStreamWaitEvent(nullptr, ev_, 0);
115 }75 }
116 return std::make_tuple(rendered, out_color, radii, geomBuffer, binningBuffer, imgBuffer);
76 ~DefaultStreamBridge() {
77 cudaEventRecord(ev_, nullptr);
78 cudaStreamWaitEvent(xla_stream_, ev_, 0);
79 cudaEventDestroy(ev_);
117}80 }
11881
119std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
120 RasterizeGaussiansBackwardCUDA(
121 const torch::Tensor& background,
122 const torch::Tensor& means3D,
123 const torch::Tensor& radii,
124 const torch::Tensor& colors,
125 const torch::Tensor& scales,
126 const torch::Tensor& rotations,
127 const float scale_modifier,
128 const torch::Tensor& cov3D_precomp,
129 const torch::Tensor& viewmatrix,
130 const torch::Tensor& projmatrix,
131 const float tan_fovx,
132 const float tan_fovy,
133 const float low_pass_2d_kernel_size,
134 const bool mip_splatting_enabled,
135 const torch::Tensor& dL_dout_color,
136 const torch::Tensor& sh,
137 const int degree,
138 const torch::Tensor& campos,
139 const torch::Tensor& geomBuffer,
140 const int R,
141 const torch::Tensor& binningBuffer,
142 const torch::Tensor& imageBuffer)
143{
144 const int P = means3D.size(0);
145 const int H = dL_dout_color.size(1);
146 const int W = dL_dout_color.size(2);
82 private:
83 cudaStream_t xla_stream_;
84 cudaEvent_t ev_{};
85};
147 86
148 int M = 0;
149 if(sh.size(0) != 0)
150 {
151 M = sh.size(1);
87class StageTimer {
88 public:
89 StageTimer(cudaStream_t stream, bool enabled)
90 : stream_(stream), enabled_(enabled) {
91 if (enabled_) Mark();
92 }
93 ~StageTimer() {
94 for (int i = 0; i < num_; ++i) cudaEventDestroy(events_[i]);
95 }
96 void Mark() {
97 if (!enabled_ || num_ > kMaxTimings) return;
98 cudaEventCreate(&events_[num_]);
99 cudaEventRecord(events_[num_], stream_);
100 ++num_;
101 }
102 void Write(float* out) {
103 std::array<float, kMaxTimings> ms{};
104 if (enabled_ && num_ > 1) {
105 cudaEventSynchronize(events_[num_ - 1]);
106 for (int i = 1; i < num_; ++i) {
107 cudaEventElapsedTime(&ms[i - 1], events_[i - 1], events_[i]);
108 }
109 }
110 cudaMemcpyAsync(out, ms.data(), sizeof(ms), cudaMemcpyHostToDevice,
111 stream_);
112 if (enabled_) cudaStreamSynchronize(stream_);
152 }113 }
153114
154 torch::Tensor dL_dmeans3D = torch::zeros({P, 3}, means3D.options());
155 torch::Tensor dL_dmeans2D = torch::zeros({P, 3}, means3D.options());
156 torch::Tensor dL_dcolors = torch::zeros({P, NUM_CHANNELS}, means3D.options());
157 torch::Tensor dL_dconic = torch::zeros({P, 2, 2}, means3D.options());
158 torch::Tensor dL_dopacity = torch::zeros({P, 1}, means3D.options());
159 torch::Tensor dL_dcov3D = torch::zeros({P, 6}, means3D.options());
160 torch::Tensor dL_dsh = torch::zeros({P, M, 3}, means3D.options());
161 torch::Tensor dL_dscales = torch::zeros({P, 3}, means3D.options());
162 torch::Tensor dL_drotations = torch::zeros({P, 4}, means3D.options());
115 private:
116 cudaStream_t stream_;
117 bool enabled_;
118 int num_ = 0;
119 std::array<cudaEvent_t, kMaxTimings + 1> events_{};
120};
163 121
164 if(P != 0 && R > 0)
165 {
166 CudaRasterizer::Rasterizer::backward(P, degree, M, R,
167 background.contiguous().data<float>(),
168 W, H,
169 means3D.contiguous().data<float>(),
170 sh.contiguous().data<float>(),
171 colors.contiguous().data<float>(),
172 scales.contiguous().data<float>(),
173 scale_modifier,
174 rotations.contiguous().data<float>(),
175 cov3D_precomp.contiguous().data<float>(),
176 viewmatrix.contiguous().data<float>(),
177 projmatrix.contiguous().data<float>(),
178 campos.contiguous().data<float>(),
179 tan_fovx,
180 tan_fovy,
181 low_pass_2d_kernel_size,
182 mip_splatting_enabled,
183 radii.contiguous().data<int>(),
184 reinterpret_cast<char*>(geomBuffer.contiguous().data_ptr()),
185 reinterpret_cast<char*>(binningBuffer.contiguous().data_ptr()),
186 reinterpret_cast<char*>(imageBuffer.contiguous().data_ptr()),
187 dL_dout_color.contiguous().data<float>(),
188 dL_dmeans2D.contiguous().data<float>(),
189 dL_dconic.contiguous().data<float>(),
190 dL_dopacity.contiguous().data<float>(),
191 dL_dcolors.contiguous().data<float>(),
192 dL_dmeans3D.contiguous().data<float>(),
193 dL_dcov3D.contiguous().data<float>(),
194 dL_dsh.contiguous().data<float>(),
195 dL_dscales.contiguous().data<float>(),
196 dL_drotations.contiguous().data<float>());
122__global__ void WriteStatsKernel(int32_t* stats, int num_pairs, int overflow) {
123 stats[0] = num_pairs;
124 stats[1] = overflow;
197 }125}
198126
199 return std::make_tuple(dL_dmeans2D, dL_dcolors, dL_dopacity, dL_dmeans3D, dL_dcov3D, dL_dsh, dL_dscales, dL_drotations);
127// Unpacks `Camera.packed()` (`f32[20]`) into Inria's column-major `viewmatrix`
128// (16 floats), column-major `projmatrix` (16 floats), `cam_pos` (3 floats), and
129// `bg_color` (3 floats).
130__global__ void PrepareCameraKernel(const float* __restrict__ cam, int width,
131 int height, float* __restrict__ out) {
132 float* view = out;
133 float* proj = out + 16;
134 float* cam_pos = out + 32;
135 float* bg = out + 35;
136
137 view[0] = cam[0];
138 view[1] = cam[3];
139 view[2] = cam[6];
140 view[3] = 0.0f;
141 view[4] = cam[1];
142 view[5] = cam[4];
143 view[6] = cam[7];
144 view[7] = 0.0f;
145 view[8] = cam[2];
146 view[9] = cam[5];
147 view[10] = cam[8];
148 view[11] = 0.0f;
149 view[12] = cam[9];
150 view[13] = cam[10];
151 view[14] = cam[11];
152 view[15] = 1.0f;
153
154 const float fx = cam[12];
155 const float fy = cam[13];
156 const float sx = 2.0f * fx / static_cast<float>(width);
157 const float sy = 2.0f * fy / static_cast<float>(height);
158 constexpr float kZNear = 0.01f;
159 constexpr float kZFar = 100.0f;
160 constexpr float kZScale = kZFar / (kZFar - kZNear);
161 constexpr float kZBias = -(kZFar * kZNear) / (kZFar - kZNear);
162 for (int c = 0; c < 4; ++c) {
163 proj[4 * c + 0] = sx * view[4 * c + 0];
164 proj[4 * c + 1] = sy * view[4 * c + 1];
165 proj[4 * c + 2] = kZScale * view[4 * c + 2] + kZBias * view[4 * c + 3];
166 proj[4 * c + 3] = view[4 * c + 2];
200}167 }
201168
202torch::Tensor markVisible(
203 torch::Tensor& means3D,
204 torch::Tensor& viewmatrix,
205 torch::Tensor& projmatrix)
206{
207 const int P = means3D.size(0);
169 cam_pos[0] = cam[16];
170 cam_pos[1] = cam[17];
171 cam_pos[2] = cam[18];
208 172
209 torch::Tensor present = torch::full({P}, false, means3D.options().dtype(at::kBool));
210
211 if(P != 0)
212 {
213 CudaRasterizer::Rasterizer::markVisible(P,
214 means3D.contiguous().data<float>(),
215 viewmatrix.contiguous().data<float>(),
216 projmatrix.contiguous().data<float>(),
217 present.contiguous().data<bool>());
173 bg[0] = 0.0f;
174 bg[1] = 0.0f;
175 bg[2] = 0.0f;
218 }176}
219 177
178// Slices the `(N, 18)` `splats` tensor into Inria's per-Gaussian geometry
179// arrays (`conic_opacity`, `means2D`, `cov3Ds`, `rgb`, `depths`, `radii`,
180// `clamped`).
181struct SplatsView {
182 float4* conic_opacity;
183 float2* means2D;
184 float* cov3Ds;
185 float* rgb;
186 float* depths;
187 int* radii;
188 bool* clamped;
189
190 static SplatsView From(float* base, int n) {
220 return present;191 return {
192 reinterpret_cast<float4*>(base + 0 * n),
193 reinterpret_cast<float2*>(base + 4 * n),
194 base + 6 * n,
195 base + 12 * n,
196 base + 15 * n,
197 reinterpret_cast<int*>(base + 16 * n),
198 reinterpret_cast<bool*>(base + 17 * n),
199 };
221}200 }
201};
202
203void ResolveFocal(const float* d_camera, float attr_fx, float attr_fy,
204 float* fx_out, float* fy_out) {
205 if (attr_fx > 0.0f && attr_fy > 0.0f) {
206 *fx_out = attr_fx;
207 *fy_out = attr_fy;
208 return;
209 }
210 float fxy[2] = {0.0f, 0.0f};
211 cudaMemcpyAsync(fxy, d_camera + 12, sizeof(fxy), cudaMemcpyDeviceToHost,
212 nullptr);
213 cudaStreamSynchronize(nullptr);
214 *fx_out = fxy[0];
215 *fy_out = fxy[1];
216}
217
218// ---------------------------------------------------------------------------
219
220ffi::Error Forward(
221 cudaStream_t stream, ffi::ScratchAllocator scratch,
222 ffi::Buffer<ffi::F32> means, ffi::Buffer<ffi::F32> scales,
223 ffi::Buffer<ffi::F32> quats, ffi::Buffer<ffi::F32> opacities,
224 ffi::Buffer<ffi::F32> sh, ffi::Buffer<ffi::F32> camera,
225 ffi::ResultBuffer<ffi::F32> image, ffi::ResultBuffer<ffi::F32> splats,
226 ffi::ResultBuffer<ffi::S32> point_list, ffi::ResultBuffer<ffi::S32> ranges,
227 ffi::ResultBuffer<ffi::F32> final_t, ffi::ResultBuffer<ffi::S32> n_contrib,
228 ffi::ResultBuffer<ffi::S32> stats, ffi::ResultBuffer<ffi::F32> timings,
229 int32_t width, int32_t height, int32_t culling, int32_t profile,
230 float fx_attr, float fy_attr) {
231 (void)culling;
232 const int n = static_cast<int>(means.element_count() / 3);
233 if (camera.element_count() != kCameraFloats) {
234 return ffi::Error(ffi::ErrorCode::kInvalidArgument,
235 "camera must have 20 floats");
236 }
237 if (n <= 0) {
238 return ffi::Error(ffi::ErrorCode::kInvalidArgument, "no Gaussians");
239 }
240
241 DefaultStreamBridge bridge(stream);
242 auto cam_buf = Alloc<float>(scratch, 40);
243 auto tiles_touched = Alloc<uint32_t>(scratch, n);
244 auto point_offsets = Alloc<uint32_t>(scratch, n);
245 if (!cam_buf || !tiles_touched || !point_offsets) {
246 return ffi::Error(ffi::ErrorCode::kResourceExhausted, "scratch");
247 }
248 PrepareCameraKernel<<<1, 1>>>(camera.typed_data(), width, height, *cam_buf);
249 float focal_x = 0.0f, focal_y = 0.0f;
250 ResolveFocal(camera.typed_data(), fx_attr, fy_attr, &focal_x, &focal_y);
251 const float tan_fovx = static_cast<float>(width) / (2.0f * focal_x);
252 const float tan_fovy = static_cast<float>(height) / (2.0f * focal_y);
253
254 const float* viewmatrix = *cam_buf;
255 const float* projmatrix = *cam_buf + 16;
256 const auto* cam_pos = reinterpret_cast<const glm::vec3*>(*cam_buf + 32);
257 const float* bg_color = *cam_buf + 35;
258
259 const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
260 (height + BLOCK_Y - 1) / BLOCK_Y, 1);
261 const dim3 block(BLOCK_X, BLOCK_Y, 1);
262 const int num_tiles = static_cast<int>(tile_grid.x * tile_grid.y);
263 const SplatsView geom = SplatsView::From(splats->typed_data(), n);
264 const float3 minn = {-FLT_MAX, -FLT_MAX, -FLT_MAX};
265 const float3 maxx = {FLT_MAX, FLT_MAX, FLT_MAX};
266
267 StageTimer timer(nullptr, profile != 0);
268
269 // 1. Preprocess.
270 FORWARD::preprocess(
271 n, kShDegree, kShCoeffs, means.typed_data(),
272 reinterpret_cast<const glm::vec3*>(scales.typed_data()), kScaleModifier,
273 reinterpret_cast<const glm::vec4*>(quats.typed_data()),
274 opacities.typed_data(), sh.typed_data(), geom.clamped, nullptr, nullptr,
275 viewmatrix, projmatrix, cam_pos, width, height, focal_x, focal_y,
276 tan_fovx, tan_fovy, kLowPassFilter, false, geom.radii, geom.means2D,
277 geom.depths, geom.cov3Ds, geom.rgb, geom.conic_opacity, tile_grid,
278 *tiles_touched, false, nullptr, minn, maxx);
279 timer.Mark();
280
281 // 2. Inclusive scan over touched tile counts.
282 size_t scan_bytes = 0;
283 cub::DeviceScan::InclusiveSum(nullptr, scan_bytes, *tiles_touched,
284 *point_offsets, n);
285 auto scan_tmp = Alloc<uint8_t>(scratch, scan_bytes);
286 if (!scan_tmp) return ffi::Error(ffi::ErrorCode::kResourceExhausted, "scan");
287 GS_RETURN_IF_CUDA_ERROR(
288 cub::DeviceScan::InclusiveSum(*scan_tmp, scan_bytes, *tiles_touched,
289 *point_offsets, n),
290 "scan");
291 int32_t num_pairs = 0;
292 GS_RETURN_IF_CUDA_ERROR(
293 cudaMemcpyAsync(&num_pairs, *point_offsets + n - 1, sizeof(int32_t),
294 cudaMemcpyDeviceToHost, nullptr),
295 "copy num_pairs");
296 GS_RETURN_IF_CUDA_ERROR(cudaStreamSynchronize(nullptr), "sync");
297 timer.Mark();
298
299 const int64_t pixels = static_cast<int64_t>(width) * height;
300 const int64_t capacity = point_list->element_count();
301 auto* range_ptr = reinterpret_cast<uint2*>(ranges->typed_data());
302 auto* point_list_ptr = reinterpret_cast<uint32_t*>(point_list->typed_data());
303 auto* n_contrib_ptr = reinterpret_cast<uint32_t*>(n_contrib->typed_data());
304
305 GS_RETURN_IF_CUDA_ERROR(
306 cudaMemsetAsync(range_ptr, 0, sizeof(uint2) * num_tiles, nullptr),
307 "memset ranges");
308 if (num_pairs > capacity) {
309 cudaMemsetAsync(image->typed_data(), 0, sizeof(float) * 3 * pixels,
310 nullptr);
311 cudaMemsetAsync(final_t->typed_data(), 0, sizeof(float) * pixels, nullptr);
312 cudaMemsetAsync(n_contrib_ptr, 0, sizeof(uint32_t) * pixels, nullptr);
313 WriteStatsKernel<<<1, 1>>>(stats->typed_data(), num_pairs, 1);
314 timer.Write(timings->typed_data());
315 return ffi::Error::Success();
316 }
317
318 // 3. Emit and sort 64-bit (tile | depth) keys.
319 if (num_pairs > 0) {
320 auto keys = Alloc<uint64_t>(scratch, num_pairs);
321 auto keys_sorted = Alloc<uint64_t>(scratch, num_pairs);
322 auto values = Alloc<uint32_t>(scratch, num_pairs);
323 if (!keys || !keys_sorted || !values) {
324 return ffi::Error(ffi::ErrorCode::kResourceExhausted, "pairs");
325 }
326 duplicateWithKeys<<<Blocks(n), kThreads>>>(n, geom.means2D, geom.depths,
327 *point_offsets, *keys, *values,
328 geom.radii, tile_grid, nullptr);
329 timer.Mark();
330 const int bit = static_cast<int>(getHigherMsb(num_tiles));
331 size_t sort_bytes = 0;
332 cub::DeviceRadixSort::SortPairs(nullptr, sort_bytes, *keys, *keys_sorted,
333 *values, point_list_ptr, num_pairs, 0,
334 32 + bit);
335 auto sort_tmp = Alloc<uint8_t>(scratch, sort_bytes);
336 if (!sort_tmp) {
337 return ffi::Error(ffi::ErrorCode::kResourceExhausted, "sort");
338 }
339 GS_RETURN_IF_CUDA_ERROR(
340 cub::DeviceRadixSort::SortPairs(*sort_tmp, sort_bytes, *keys,
341 *keys_sorted, *values, point_list_ptr,
342 num_pairs, 0, 32 + bit),
343 "sort");
344 timer.Mark();
345 identifyTileRanges<<<Blocks(num_pairs), kThreads>>>(num_pairs, *keys_sorted,
346 range_ptr);
347 } else {
348 timer.Mark();
349 timer.Mark();
350 }
351 timer.Mark();
352
353 // 4. Rasterize.
354 FORWARD::render(tile_grid, block, range_ptr, point_list_ptr, width, height,
355 geom.means2D, geom.rgb, geom.conic_opacity,
356 final_t->typed_data(), n_contrib_ptr, bg_color,
357 image->typed_data());
358 timer.Mark();
359 WriteStatsKernel<<<1, 1>>>(stats->typed_data(), num_pairs, 0);
360 timer.Write(timings->typed_data());
361 GS_RETURN_IF_CUDA_ERROR(cudaGetLastError(), "forward launch");
362 return ffi::Error::Success();
363}
364
365XLA_FFI_DEFINE_HANDLER(kForward, Forward,
366 ffi::Ffi::Bind()
367 .Ctx<ffi::PlatformStream<cudaStream_t>>()
368 .Ctx<ffi::ScratchAllocator>()
369 .Arg<ffi::Buffer<ffi::F32>>() // means
370 .Arg<ffi::Buffer<ffi::F32>>() // scales
371 .Arg<ffi::Buffer<ffi::F32>>() // quats
372 .Arg<ffi::Buffer<ffi::F32>>() // opacities
373 .Arg<ffi::Buffer<ffi::F32>>() // sh
374 .Arg<ffi::Buffer<ffi::F32>>() // camera
375 .Ret<ffi::Buffer<ffi::F32>>() // image
376 .Ret<ffi::Buffer<ffi::F32>>() // splats
377 .Ret<ffi::Buffer<ffi::S32>>() // point_list
378 .Ret<ffi::Buffer<ffi::S32>>() // ranges
379 .Ret<ffi::Buffer<ffi::F32>>() // final_t
380 .Ret<ffi::Buffer<ffi::S32>>() // n_contrib
381 .Ret<ffi::Buffer<ffi::S32>>() // stats
382 .Ret<ffi::Buffer<ffi::F32>>() // timings
383 .Attr<int32_t>("width")
384 .Attr<int32_t>("height")
385 .Attr<int32_t>("culling")
386 .Attr<int32_t>("profile")
387 .Attr<float>("fx")
388 .Attr<float>("fy"));
389
390// ---------------------------------------------------------------------------
391
392ffi::Error Backward(
393 cudaStream_t stream, ffi::ScratchAllocator scratch,
394 ffi::Buffer<ffi::F32> means, ffi::Buffer<ffi::F32> scales,
395 ffi::Buffer<ffi::F32> quats, ffi::Buffer<ffi::F32> opacities,
396 ffi::Buffer<ffi::F32> sh, ffi::Buffer<ffi::F32> camera,
397 ffi::Buffer<ffi::F32> splats, ffi::Buffer<ffi::S32> point_list,
398 ffi::Buffer<ffi::S32> ranges, ffi::Buffer<ffi::F32> final_t,
399 ffi::Buffer<ffi::S32> n_contrib, ffi::Buffer<ffi::F32> dl_dimage,
400 ffi::ResultBuffer<ffi::F32> dmeans, ffi::ResultBuffer<ffi::F32> dscales,
401 ffi::ResultBuffer<ffi::F32> dquats, ffi::ResultBuffer<ffi::F32> dopacities,
402 ffi::ResultBuffer<ffi::F32> dsh, ffi::ResultBuffer<ffi::F32> timings,
403 int32_t width, int32_t height, int32_t profile, float fx_attr,
404 float fy_attr) {
405 (void)opacities;
406 const int n = static_cast<int>(means.element_count() / 3);
407 if (camera.element_count() != kCameraFloats) {
408 return ffi::Error(ffi::ErrorCode::kInvalidArgument,
409 "camera must have 20 floats");
410 }
411 if (n <= 0) {
412 return ffi::Error(ffi::ErrorCode::kInvalidArgument, "no Gaussians");
413 }
414
415 DefaultStreamBridge bridge(stream);
416 auto cam_buf = Alloc<float>(scratch, 40);
417 auto dl_dmean2d = Alloc<float3>(scratch, n);
418 auto dl_dconic2d = Alloc<float4>(scratch, n);
419 auto dl_dcolor = Alloc<float>(scratch, 3 * n);
420 auto dl_dcov3d = Alloc<float>(scratch, 6 * n);
421 if (!cam_buf || !dl_dmean2d || !dl_dconic2d || !dl_dcolor || !dl_dcov3d) {
422 return ffi::Error(ffi::ErrorCode::kResourceExhausted, "backward scratch");
423 }
424
425 PrepareCameraKernel<<<1, 1>>>(camera.typed_data(), width, height, *cam_buf);
426 float focal_x = 0.0f, focal_y = 0.0f;
427 ResolveFocal(camera.typed_data(), fx_attr, fy_attr, &focal_x, &focal_y);
428 const float tan_fovx = static_cast<float>(width) / (2.0f * focal_x);
429 const float tan_fovy = static_cast<float>(height) / (2.0f * focal_y);
430
431 const float* viewmatrix = *cam_buf;
432 const float* projmatrix = *cam_buf + 16;
433 const auto* cam_pos = reinterpret_cast<const glm::vec3*>(*cam_buf + 32);
434 const float* bg_color = *cam_buf + 35;
435
436 const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
437 (height + BLOCK_Y - 1) / BLOCK_Y, 1);
438 const dim3 block(BLOCK_X, BLOCK_Y, 1);
439 const SplatsView geom =
440 SplatsView::From(const_cast<float*>(splats.typed_data()), n);
441
442 StageTimer timer(nullptr, profile != 0);
443
444 GS_RETURN_IF_CUDA_ERROR(
445 cudaMemsetAsync(*dl_dmean2d, 0, sizeof(float3) * n, nullptr), "memset");
446 GS_RETURN_IF_CUDA_ERROR(
447 cudaMemsetAsync(*dl_dconic2d, 0, sizeof(float4) * n, nullptr), "memset");
448 GS_RETURN_IF_CUDA_ERROR(
449 cudaMemsetAsync(*dl_dcolor, 0, sizeof(float) * 3 * n, nullptr), "memset");
450 GS_RETURN_IF_CUDA_ERROR(
451 cudaMemsetAsync(*dl_dcov3d, 0, sizeof(float) * 6 * n, nullptr), "memset");
452 GS_RETURN_IF_CUDA_ERROR(
453 cudaMemsetAsync(dopacities->typed_data(), 0, sizeof(float) * n, nullptr),
454 "memset");
455
456 // 1. Backward rasterize.
457 BACKWARD::render(tile_grid, block,
458 reinterpret_cast<const uint2*>(ranges.typed_data()),
459 reinterpret_cast<const uint32_t*>(point_list.typed_data()),
460 width, height, bg_color, geom.means2D, geom.conic_opacity,
461 geom.rgb, final_t.typed_data(),
462 reinterpret_cast<const uint32_t*>(n_contrib.typed_data()),
463 dl_dimage.typed_data(), *dl_dmean2d, *dl_dconic2d,
464 dopacities->typed_data(), *dl_dcolor);
465 timer.Mark();
466
467 // 2. Backward preprocess (as Inria's PyTorch wrapper `torch::zeros`: bulk
468 // zero parameter gradient tensors before `BACKWARD::preprocess`).
469 GS_RETURN_IF_CUDA_ERROR(
470 cudaMemsetAsync(dmeans->typed_data(), 0, sizeof(float) * 3 * n, nullptr),
471 "memset");
472 GS_RETURN_IF_CUDA_ERROR(
473 cudaMemsetAsync(dscales->typed_data(), 0, sizeof(float) * 3 * n, nullptr),
474 "memset");
475 GS_RETURN_IF_CUDA_ERROR(
476 cudaMemsetAsync(dquats->typed_data(), 0, sizeof(float) * 4 * n, nullptr),
477 "memset");
478 GS_RETURN_IF_CUDA_ERROR(
479 cudaMemsetAsync(dsh->typed_data(), 0, sizeof(float) * 3 * kShCoeffs * n,
480 nullptr),
481 "memset");
482
483 BACKWARD::preprocess(n, kShDegree, kShCoeffs,
484 reinterpret_cast<const float3*>(means.typed_data()),
485 geom.radii, sh.typed_data(), geom.clamped,
486 reinterpret_cast<const glm::vec3*>(scales.typed_data()),
487 reinterpret_cast<const glm::vec4*>(quats.typed_data()),
488 kScaleModifier, geom.cov3Ds, geom.conic_opacity,
489 viewmatrix, projmatrix, focal_x, focal_y, tan_fovx,
490 tan_fovy, kLowPassFilter, false, cam_pos, *dl_dmean2d,
491 reinterpret_cast<const float*>(*dl_dconic2d),
492 reinterpret_cast<glm::vec3*>(dmeans->typed_data()),
493 *dl_dcolor, *dl_dcov3d, dsh->typed_data(),
494 reinterpret_cast<glm::vec3*>(dscales->typed_data()),
495 reinterpret_cast<glm::vec4*>(dquats->typed_data()),
496 dopacities->typed_data());
497 timer.Mark();
498 timer.Write(timings->typed_data());
499 GS_RETURN_IF_CUDA_ERROR(cudaGetLastError(), "backward launch");
500 return ffi::Error::Success();
501}
502
503XLA_FFI_DEFINE_HANDLER(kBackward, Backward,
504 ffi::Ffi::Bind()
505 .Ctx<ffi::PlatformStream<cudaStream_t>>()
506 .Ctx<ffi::ScratchAllocator>()
507 .Arg<ffi::Buffer<ffi::F32>>() // means
508 .Arg<ffi::Buffer<ffi::F32>>() // scales
509 .Arg<ffi::Buffer<ffi::F32>>() // quats
510 .Arg<ffi::Buffer<ffi::F32>>() // opacities
511 .Arg<ffi::Buffer<ffi::F32>>() // sh
512 .Arg<ffi::Buffer<ffi::F32>>() // camera
513 .Arg<ffi::Buffer<ffi::F32>>() // splats
514 .Arg<ffi::Buffer<ffi::S32>>() // point_list
515 .Arg<ffi::Buffer<ffi::S32>>() // ranges
516 .Arg<ffi::Buffer<ffi::F32>>() // final_t
517 .Arg<ffi::Buffer<ffi::S32>>() // n_contrib
518 .Arg<ffi::Buffer<ffi::F32>>() // dl_dimage
519 .Ret<ffi::Buffer<ffi::F32>>() // dmeans
520 .Ret<ffi::Buffer<ffi::F32>>() // dscales
521 .Ret<ffi::Buffer<ffi::F32>>() // dquats
522 .Ret<ffi::Buffer<ffi::F32>>() // dopacities
523 .Ret<ffi::Buffer<ffi::F32>>() // dsh
524 .Ret<ffi::Buffer<ffi::F32>>() // timings
525 .Attr<int32_t>("width")
526 .Attr<int32_t>("height")
527 .Attr<int32_t>("profile")
528 .Attr<float>("fx")
529 .Attr<float>("fy"));
530
531XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "gs_forward", "CUDA", kForward);
532XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "gs_backward", "CUDA", kBackward);
533
534} // namespace
535} // namespace gs
1−/*
2− * Copyright (C) 2023, Inria
3− * GRAPHDECO research group, https://team.inria.fr/graphdeco
4− * All rights reserved.
5− *
6− * This software is free for non-commercial, research and evaluation use
7− * under the terms of the LICENSE.md file.
8− *
9− * For inquiries contact george.drettakis@inria.fr
10− */
1+// XLA FFI entry points driving the verbatim Inria `diff-gaussian-rasterization`
2+// CUDA kernels (`forward.cu.cc`, `backward.cu.cc`, `rasterizer_impl.cu.cc`).
3+#include "xla/ffi/api/ffi.h"
114
12−#include <math.h>
5+#include <cuda_runtime.h>
13−#include <torch/extension.h>
14−#include <cstdio>
6+
15−#include <sstream>
7+#include <array>
16−#include <iostream>
8+#include <cfloat>
17−#include <tuple>
9+#include <cstddef>
18−#include <stdio.h>
10+#include <cstdint>
19−#include <cuda_runtime_api.h>
11+#include <cub/cub.cuh>
20−#include <memory>
21−#include "cuda_rasterizer/config.h"
22−#include "cuda_rasterizer/rasterizer.h"
23−#include <fstream>
12+#include <optional>
2413 #include <string>
25−#include <functional>
2614
27−std::function<char*(size_t N)> resizeFunctional(torch::Tensor& t) {
28− auto lambda = [&t](size_t N) {
29− t.resize_({(long long)N});
30− return reinterpret_cast<char*>(t.contiguous().data_ptr());
31− };
32− return lambda;
15+#include "csrc/backward.h"
16+#include "csrc/config.h"
17+#include "csrc/forward.h"
18+#include "csrc/rasterizer_impl.h"
19+#include "xla/ffi/api/c_api.h"
20+
21+// Non-static helpers defined in `rasterizer_impl.cu.cc`.
22+uint32_t getHigherMsb(uint32_t n);
23+__global__ void duplicateWithKeys(int P, const float2* points_xy,
24+ const float* depths, const uint32_t* offsets,
25+ uint64_t* gaussian_keys_unsorted,
26+ uint32_t* gaussian_values_unsorted,
27+ int* radii, dim3 grid, int2* rects);
28+__global__ void identifyTileRanges(int L, uint64_t* point_list_keys,
29+ uint2* ranges);
30+
31+namespace gs {
32+namespace {
33+
34+namespace ffi = ::xla::ffi;
35+
36+constexpr int kThreads = 256;
37+constexpr int kMaxTimings = 8;
38+constexpr int kCameraFloats = 20;
39+constexpr int kShDegree = 3;
40+constexpr int kShCoeffs = 16;
41+constexpr float kScaleModifier = 1.0f;
42+constexpr float kLowPassFilter = 0.3f;
43+
44+int Blocks(int64_t n) {
45+ return static_cast<int>((n + kThreads - 1) / kThreads);
3346 }
3447
35−std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
36−RasterizeGaussiansCUDA(
37− const torch::Tensor& background,
38− const torch::Tensor& means3D,
39− const torch::Tensor& colors,
40− const torch::Tensor& opacity,
41− const torch::Tensor& scales,
42− const torch::Tensor& rotations,
43− const float scale_modifier,
44− const torch::Tensor& cov3D_precomp,
45− const torch::Tensor& viewmatrix,
46− const torch::Tensor& projmatrix,
47− const float tan_fovx,
48− const float tan_fovy,
49− const float low_pass_2d_kernel_size,
50− const bool mip_splatting_enabled,
51− const int image_height,
52− const int image_width,
53− const torch::Tensor& sh,
54− const int degree,
55− const torch::Tensor& campos,
56− const bool prefiltered)
57−{
58− if (means3D.ndimension() != 2 || means3D.size(1) != 3) {
59− AT_ERROR("means3D must have dimensions (num_points, 3)");
48+ffi::Error CudaError(cudaError_t e, const char* what) {
49+ return ffi::Error(ffi::ErrorCode::kInternal,
50+ std::string(what) + ": " + cudaGetErrorString(e));
6051⇥}
6152
62− const int P = means3D.size(0);
63− const int H = image_height;
64− const int W = image_width;
53+#define GS_RETURN_IF_CUDA_ERROR(expr, what) \
54+ do { \
55+ cudaError_t gs_err_ = (expr); \
56+ if (gs_err_ != cudaSuccess) return CudaError(gs_err_, what); \
57+ } while (0)
6558
66− auto int_opts = means3D.options().dtype(torch::kInt32);
67− auto float_opts = means3D.options().dtype(torch::kFloat32);
68−
69− torch::Tensor out_color = torch::full({NUM_CHANNELS, H, W}, 0.0, float_opts);
70− torch::Tensor radii = torch::full({P}, 0, means3D.options().dtype(torch::kInt32));
71−
72− torch::Device device(torch::kCUDA);
73− torch::TensorOptions options(torch::kByte);
74− torch::Tensor geomBuffer = torch::empty({0}, options.device(device));
75− torch::Tensor binningBuffer = torch::empty({0}, options.device(device));
76− torch::Tensor imgBuffer = torch::empty({0}, options.device(device));
77− std::function<char*(size_t)> geomFunc = resizeFunctional(geomBuffer);
78− std::function<char*(size_t)> binningFunc = resizeFunctional(binningBuffer);
79− std::function<char*(size_t)> imgFunc = resizeFunctional(imgBuffer);
80−
81− int rendered = 0;
82− if(P != 0)
83− {
84− int M = 0;
85− if(sh.size(0) != 0)
86− {
87− M = sh.size(1);
59+template <typename T>
60+std::optional<T*> Alloc(ffi::ScratchAllocator& scratch, size_t count) {
61+ auto p = scratch.Allocate(sizeof(T) * (count > 0 ? count : 1), 16);
62+ if (!p.has_value()) return std::nullopt;
63+ return static_cast<T*>(*p);
8864⇥}
8965
90− rendered = CudaRasterizer::Rasterizer::forward(
91− geomFunc,
92− binningFunc,
93− imgFunc,
94− P, degree, M,
95− background.contiguous().data<float>(),
96− W, H,
97− means3D.contiguous().data<float>(),
98− sh.contiguous().data_ptr<float>(),
99− colors.contiguous().data<float>(),
100− opacity.contiguous().data<float>(),
101− scales.contiguous().data_ptr<float>(),
102− scale_modifier,
103− rotations.contiguous().data_ptr<float>(),
104− cov3D_precomp.contiguous().data<float>(),
105− viewmatrix.contiguous().data<float>(),
106− projmatrix.contiguous().data<float>(),
107− campos.contiguous().data<float>(),
108− tan_fovx,
109− tan_fovy,
110− low_pass_2d_kernel_size,
111− mip_splatting_enabled,
112− prefiltered,
113− out_color.contiguous().data<float>(),
114− radii.contiguous().data<int>());
66+// Bridges XLA's non-blocking `stream` to the default stream (`nullptr`) used by
67+// Inria's unmodified `FORWARD::*` and `BACKWARD::*` kernel launches.
68+class DefaultStreamBridge {
69+ public:
70+ explicit DefaultStreamBridge(cudaStream_t xla_stream)
71+ : xla_stream_(xla_stream) {
72+ cudaEventCreateWithFlags(&ev_, cudaEventDisableTiming);
73+ cudaEventRecord(ev_, xla_stream_);
74+ cudaStreamWaitEvent(nullptr, ev_, 0);
11575 }
116− return std::make_tuple(rendered, out_color, radii, geomBuffer, binningBuffer, imgBuffer);
76+ ~DefaultStreamBridge() {
77+ cudaEventRecord(ev_, nullptr);
78+ cudaStreamWaitEvent(xla_stream_, ev_, 0);
79+ cudaEventDestroy(ev_);
11780⇥ }
11881
119−std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
120− RasterizeGaussiansBackwardCUDA(
121− const torch::Tensor& background,
122− const torch::Tensor& means3D,
123− const torch::Tensor& radii,
124− const torch::Tensor& colors,
125− const torch::Tensor& scales,
126− const torch::Tensor& rotations,
127− const float scale_modifier,
128− const torch::Tensor& cov3D_precomp,
129− const torch::Tensor& viewmatrix,
130− const torch::Tensor& projmatrix,
131− const float tan_fovx,
132− const float tan_fovy,
133− const float low_pass_2d_kernel_size,
134− const bool mip_splatting_enabled,
135− const torch::Tensor& dL_dout_color,
136− const torch::Tensor& sh,
137− const int degree,
138− const torch::Tensor& campos,
139− const torch::Tensor& geomBuffer,
140− const int R,
141− const torch::Tensor& binningBuffer,
142− const torch::Tensor& imageBuffer)
143−{
144− const int P = means3D.size(0);
145− const int H = dL_dout_color.size(1);
146− const int W = dL_dout_color.size(2);
82+ private:
83+ cudaStream_t xla_stream_;
84+ cudaEvent_t ev_{};
85+};
14786
148− int M = 0;
149− if(sh.size(0) != 0)
150− {
151− M = sh.size(1);
87+class StageTimer {
88+ public:
89+ StageTimer(cudaStream_t stream, bool enabled)
90+ : stream_(stream), enabled_(enabled) {
91+ if (enabled_) Mark();
92+ }
93+ ~StageTimer() {
94+ for (int i = 0; i < num_; ++i) cudaEventDestroy(events_[i]);
95+ }
96+ void Mark() {
97+ if (!enabled_ || num_ > kMaxTimings) return;
98+ cudaEventCreate(&events_[num_]);
99+ cudaEventRecord(events_[num_], stream_);
100+ ++num_;
101+ }
102+ void Write(float* out) {
103+ std::array<float, kMaxTimings> ms{};
104+ if (enabled_ && num_ > 1) {
105+ cudaEventSynchronize(events_[num_ - 1]);
106+ for (int i = 1; i < num_; ++i) {
107+ cudaEventElapsedTime(&ms[i - 1], events_[i - 1], events_[i]);
108+ }
109+ }
110+ cudaMemcpyAsync(out, ms.data(), sizeof(ms), cudaMemcpyHostToDevice,
111+ stream_);
112+ if (enabled_) cudaStreamSynchronize(stream_);
152113 }
153114
154− torch::Tensor dL_dmeans3D = torch::zeros({P, 3}, means3D.options());
155− torch::Tensor dL_dmeans2D = torch::zeros({P, 3}, means3D.options());
156− torch::Tensor dL_dcolors = torch::zeros({P, NUM_CHANNELS}, means3D.options());
157− torch::Tensor dL_dconic = torch::zeros({P, 2, 2}, means3D.options());
158− torch::Tensor dL_dopacity = torch::zeros({P, 1}, means3D.options());
159− torch::Tensor dL_dcov3D = torch::zeros({P, 6}, means3D.options());
160− torch::Tensor dL_dsh = torch::zeros({P, M, 3}, means3D.options());
161− torch::Tensor dL_dscales = torch::zeros({P, 3}, means3D.options());
162− torch::Tensor dL_drotations = torch::zeros({P, 4}, means3D.options());
115+ private:
116+ cudaStream_t stream_;
117+ bool enabled_;
118+ int num_ = 0;
119+ std::array<cudaEvent_t, kMaxTimings + 1> events_{};
120+};
163121
164− if(P != 0 && R > 0)
165− {
166− CudaRasterizer::Rasterizer::backward(P, degree, M, R,
167− background.contiguous().data<float>(),
168− W, H,
169− means3D.contiguous().data<float>(),
170− sh.contiguous().data<float>(),
171− colors.contiguous().data<float>(),
172− scales.contiguous().data<float>(),
173− scale_modifier,
174− rotations.contiguous().data<float>(),
175− cov3D_precomp.contiguous().data<float>(),
176− viewmatrix.contiguous().data<float>(),
177− projmatrix.contiguous().data<float>(),
178− campos.contiguous().data<float>(),
179− tan_fovx,
180− tan_fovy,
181− low_pass_2d_kernel_size,
182− mip_splatting_enabled,
183− radii.contiguous().data<int>(),
184− reinterpret_cast<char*>(geomBuffer.contiguous().data_ptr()),
185− reinterpret_cast<char*>(binningBuffer.contiguous().data_ptr()),
186− reinterpret_cast<char*>(imageBuffer.contiguous().data_ptr()),
187− dL_dout_color.contiguous().data<float>(),
188− dL_dmeans2D.contiguous().data<float>(),
189− dL_dconic.contiguous().data<float>(),
190− dL_dopacity.contiguous().data<float>(),
191− dL_dcolors.contiguous().data<float>(),
192− dL_dmeans3D.contiguous().data<float>(),
193− dL_dcov3D.contiguous().data<float>(),
194− dL_dsh.contiguous().data<float>(),
195− dL_dscales.contiguous().data<float>(),
196− dL_drotations.contiguous().data<float>());
122+__global__ void WriteStatsKernel(int32_t* stats, int num_pairs, int overflow) {
123+ stats[0] = num_pairs;
124+ stats[1] = overflow;
197125⇥}
198126
199− return std::make_tuple(dL_dmeans2D, dL_dcolors, dL_dopacity, dL_dmeans3D, dL_dcov3D, dL_dsh, dL_dscales, dL_drotations);
127+// Unpacks `Camera.packed()` (`f32[20]`) into Inria's column-major `viewmatrix`
128+// (16 floats), column-major `projmatrix` (16 floats), `cam_pos` (3 floats), and
129+// `bg_color` (3 floats).
130+__global__ void PrepareCameraKernel(const float* __restrict__ cam, int width,
131+ int height, float* __restrict__ out) {
132+ float* view = out;
133+ float* proj = out + 16;
134+ float* cam_pos = out + 32;
135+ float* bg = out + 35;
136+
137+ view[0] = cam[0];
138+ view[1] = cam[3];
139+ view[2] = cam[6];
140+ view[3] = 0.0f;
141+ view[4] = cam[1];
142+ view[5] = cam[4];
143+ view[6] = cam[7];
144+ view[7] = 0.0f;
145+ view[8] = cam[2];
146+ view[9] = cam[5];
147+ view[10] = cam[8];
148+ view[11] = 0.0f;
149+ view[12] = cam[9];
150+ view[13] = cam[10];
151+ view[14] = cam[11];
152+ view[15] = 1.0f;
153+
154+ const float fx = cam[12];
155+ const float fy = cam[13];
156+ const float sx = 2.0f * fx / static_cast<float>(width);
157+ const float sy = 2.0f * fy / static_cast<float>(height);
158+ constexpr float kZNear = 0.01f;
159+ constexpr float kZFar = 100.0f;
160+ constexpr float kZScale = kZFar / (kZFar - kZNear);
161+ constexpr float kZBias = -(kZFar * kZNear) / (kZFar - kZNear);
162+ for (int c = 0; c < 4; ++c) {
163+ proj[4 * c + 0] = sx * view[4 * c + 0];
164+ proj[4 * c + 1] = sy * view[4 * c + 1];
165+ proj[4 * c + 2] = kZScale * view[4 * c + 2] + kZBias * view[4 * c + 3];
166+ proj[4 * c + 3] = view[4 * c + 2];
200167⇥ }
201168
202−torch::Tensor markVisible(
203− torch::Tensor& means3D,
204− torch::Tensor& viewmatrix,
205− torch::Tensor& projmatrix)
206−{
207− const int P = means3D.size(0);
169+ cam_pos[0] = cam[16];
170+ cam_pos[1] = cam[17];
171+ cam_pos[2] = cam[18];
208172
209− torch::Tensor present = torch::full({P}, false, means3D.options().dtype(at::kBool));
210−
211− if(P != 0)
212− {
213− CudaRasterizer::Rasterizer::markVisible(P,
214− means3D.contiguous().data<float>(),
215− viewmatrix.contiguous().data<float>(),
216− projmatrix.contiguous().data<float>(),
217− present.contiguous().data<bool>());
173+ bg[0] = 0.0f;
174+ bg[1] = 0.0f;
175+ bg[2] = 0.0f;
218176⇥}
219177
178+// Slices the `(N, 18)` `splats` tensor into Inria's per-Gaussian geometry
179+// arrays (`conic_opacity`, `means2D`, `cov3Ds`, `rgb`, `depths`, `radii`,
180+// `clamped`).
181+struct SplatsView {
182+ float4* conic_opacity;
183+ float2* means2D;
184+ float* cov3Ds;
185+ float* rgb;
186+ float* depths;
187+ int* radii;
188+ bool* clamped;
189+
190+ static SplatsView From(float* base, int n) {
220− return present;
191+ return {
192+ reinterpret_cast<float4*>(base + 0 * n),
193+ reinterpret_cast<float2*>(base + 4 * n),
194+ base + 6 * n,
195+ base + 12 * n,
196+ base + 15 * n,
197+ reinterpret_cast<int*>(base + 16 * n),
198+ reinterpret_cast<bool*>(base + 17 * n),
199+ };
221200⇥ }
201+};
202+
203+void ResolveFocal(const float* d_camera, float attr_fx, float attr_fy,
204+ float* fx_out, float* fy_out) {
205+ if (attr_fx > 0.0f && attr_fy > 0.0f) {
206+ *fx_out = attr_fx;
207+ *fy_out = attr_fy;
208+ return;
209+ }
210+ float fxy[2] = {0.0f, 0.0f};
211+ cudaMemcpyAsync(fxy, d_camera + 12, sizeof(fxy), cudaMemcpyDeviceToHost,
212+ nullptr);
213+ cudaStreamSynchronize(nullptr);
214+ *fx_out = fxy[0];
215+ *fy_out = fxy[1];
216+}
217+
218+// ---------------------------------------------------------------------------
219+
220+ffi::Error Forward(
221+ cudaStream_t stream, ffi::ScratchAllocator scratch,
222+ ffi::Buffer<ffi::F32> means, ffi::Buffer<ffi::F32> scales,
223+ ffi::Buffer<ffi::F32> quats, ffi::Buffer<ffi::F32> opacities,
224+ ffi::Buffer<ffi::F32> sh, ffi::Buffer<ffi::F32> camera,
225+ ffi::ResultBuffer<ffi::F32> image, ffi::ResultBuffer<ffi::F32> splats,
226+ ffi::ResultBuffer<ffi::S32> point_list, ffi::ResultBuffer<ffi::S32> ranges,
227+ ffi::ResultBuffer<ffi::F32> final_t, ffi::ResultBuffer<ffi::S32> n_contrib,
228+ ffi::ResultBuffer<ffi::S32> stats, ffi::ResultBuffer<ffi::F32> timings,
229+ int32_t width, int32_t height, int32_t culling, int32_t profile,
230+ float fx_attr, float fy_attr) {
231+ (void)culling;
232+ const int n = static_cast<int>(means.element_count() / 3);
233+ if (camera.element_count() != kCameraFloats) {
234+ return ffi::Error(ffi::ErrorCode::kInvalidArgument,
235+ "camera must have 20 floats");
236+ }
237+ if (n <= 0) {
238+ return ffi::Error(ffi::ErrorCode::kInvalidArgument, "no Gaussians");
239+ }
240+
241+ DefaultStreamBridge bridge(stream);
242+ auto cam_buf = Alloc<float>(scratch, 40);
243+ auto tiles_touched = Alloc<uint32_t>(scratch, n);
244+ auto point_offsets = Alloc<uint32_t>(scratch, n);
245+ if (!cam_buf || !tiles_touched || !point_offsets) {
246+ return ffi::Error(ffi::ErrorCode::kResourceExhausted, "scratch");
247+ }
248+ PrepareCameraKernel<<<1, 1>>>(camera.typed_data(), width, height, *cam_buf);
249+ float focal_x = 0.0f, focal_y = 0.0f;
250+ ResolveFocal(camera.typed_data(), fx_attr, fy_attr, &focal_x, &focal_y);
251+ const float tan_fovx = static_cast<float>(width) / (2.0f * focal_x);
252+ const float tan_fovy = static_cast<float>(height) / (2.0f * focal_y);
253+
254+ const float* viewmatrix = *cam_buf;
255+ const float* projmatrix = *cam_buf + 16;
256+ const auto* cam_pos = reinterpret_cast<const glm::vec3*>(*cam_buf + 32);
257+ const float* bg_color = *cam_buf + 35;
258+
259+ const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
260+ (height + BLOCK_Y - 1) / BLOCK_Y, 1);
261+ const dim3 block(BLOCK_X, BLOCK_Y, 1);
262+ const int num_tiles = static_cast<int>(tile_grid.x * tile_grid.y);
263+ const SplatsView geom = SplatsView::From(splats->typed_data(), n);
264+ const float3 minn = {-FLT_MAX, -FLT_MAX, -FLT_MAX};
265+ const float3 maxx = {FLT_MAX, FLT_MAX, FLT_MAX};
266+
267+ StageTimer timer(nullptr, profile != 0);
268+
269+ // 1. Preprocess.
270+ FORWARD::preprocess(
271+ n, kShDegree, kShCoeffs, means.typed_data(),
272+ reinterpret_cast<const glm::vec3*>(scales.typed_data()), kScaleModifier,
273+ reinterpret_cast<const glm::vec4*>(quats.typed_data()),
274+ opacities.typed_data(), sh.typed_data(), geom.clamped, nullptr, nullptr,
275+ viewmatrix, projmatrix, cam_pos, width, height, focal_x, focal_y,
276+ tan_fovx, tan_fovy, kLowPassFilter, false, geom.radii, geom.means2D,
277+ geom.depths, geom.cov3Ds, geom.rgb, geom.conic_opacity, tile_grid,
278+ *tiles_touched, false, nullptr, minn, maxx);
279+ timer.Mark();
280+
281+ // 2. Inclusive scan over touched tile counts.
282+ size_t scan_bytes = 0;
283+ cub::DeviceScan::InclusiveSum(nullptr, scan_bytes, *tiles_touched,
284+ *point_offsets, n);
285+ auto scan_tmp = Alloc<uint8_t>(scratch, scan_bytes);
286+ if (!scan_tmp) return ffi::Error(ffi::ErrorCode::kResourceExhausted, "scan");
287+ GS_RETURN_IF_CUDA_ERROR(
288+ cub::DeviceScan::InclusiveSum(*scan_tmp, scan_bytes, *tiles_touched,
289+ *point_offsets, n),
290+ "scan");
291+ int32_t num_pairs = 0;
292+ GS_RETURN_IF_CUDA_ERROR(
293+ cudaMemcpyAsync(&num_pairs, *point_offsets + n - 1, sizeof(int32_t),
294+ cudaMemcpyDeviceToHost, nullptr),
295+ "copy num_pairs");
296+ GS_RETURN_IF_CUDA_ERROR(cudaStreamSynchronize(nullptr), "sync");
297+ timer.Mark();
298+
299+ const int64_t pixels = static_cast<int64_t>(width) * height;
300+ const int64_t capacity = point_list->element_count();
301+ auto* range_ptr = reinterpret_cast<uint2*>(ranges->typed_data());
302+ auto* point_list_ptr = reinterpret_cast<uint32_t*>(point_list->typed_data());
303+ auto* n_contrib_ptr = reinterpret_cast<uint32_t*>(n_contrib->typed_data());
304+
305+ GS_RETURN_IF_CUDA_ERROR(
306+ cudaMemsetAsync(range_ptr, 0, sizeof(uint2) * num_tiles, nullptr),
307+ "memset ranges");
308+ if (num_pairs > capacity) {
309+ cudaMemsetAsync(image->typed_data(), 0, sizeof(float) * 3 * pixels,
310+ nullptr);
311+ cudaMemsetAsync(final_t->typed_data(), 0, sizeof(float) * pixels, nullptr);
312+ cudaMemsetAsync(n_contrib_ptr, 0, sizeof(uint32_t) * pixels, nullptr);
313+ WriteStatsKernel<<<1, 1>>>(stats->typed_data(), num_pairs, 1);
314+ timer.Write(timings->typed_data());
315+ return ffi::Error::Success();
316+ }
317+
318+ // 3. Emit and sort 64-bit (tile | depth) keys.
319+ if (num_pairs > 0) {
320+ auto keys = Alloc<uint64_t>(scratch, num_pairs);
321+ auto keys_sorted = Alloc<uint64_t>(scratch, num_pairs);
322+ auto values = Alloc<uint32_t>(scratch, num_pairs);
323+ if (!keys || !keys_sorted || !values) {
324+ return ffi::Error(ffi::ErrorCode::kResourceExhausted, "pairs");
325+ }
326+ duplicateWithKeys<<<Blocks(n), kThreads>>>(n, geom.means2D, geom.depths,
327+ *point_offsets, *keys, *values,
328+ geom.radii, tile_grid, nullptr);
329+ timer.Mark();
330+ const int bit = static_cast<int>(getHigherMsb(num_tiles));
331+ size_t sort_bytes = 0;
332+ cub::DeviceRadixSort::SortPairs(nullptr, sort_bytes, *keys, *keys_sorted,
333+ *values, point_list_ptr, num_pairs, 0,
334+ 32 + bit);
335+ auto sort_tmp = Alloc<uint8_t>(scratch, sort_bytes);
336+ if (!sort_tmp) {
337+ return ffi::Error(ffi::ErrorCode::kResourceExhausted, "sort");
338+ }
339+ GS_RETURN_IF_CUDA_ERROR(
340+ cub::DeviceRadixSort::SortPairs(*sort_tmp, sort_bytes, *keys,
341+ *keys_sorted, *values, point_list_ptr,
342+ num_pairs, 0, 32 + bit),
343+ "sort");
344+ timer.Mark();
345+ identifyTileRanges<<<Blocks(num_pairs), kThreads>>>(num_pairs, *keys_sorted,
346+ range_ptr);
347+ } else {
348+ timer.Mark();
349+ timer.Mark();
350+ }
351+ timer.Mark();
352+
353+ // 4. Rasterize.
354+ FORWARD::render(tile_grid, block, range_ptr, point_list_ptr, width, height,
355+ geom.means2D, geom.rgb, geom.conic_opacity,
356+ final_t->typed_data(), n_contrib_ptr, bg_color,
357+ image->typed_data());
358+ timer.Mark();
359+ WriteStatsKernel<<<1, 1>>>(stats->typed_data(), num_pairs, 0);
360+ timer.Write(timings->typed_data());
361+ GS_RETURN_IF_CUDA_ERROR(cudaGetLastError(), "forward launch");
362+ return ffi::Error::Success();
363+}
364+
365+XLA_FFI_DEFINE_HANDLER(kForward, Forward,
366+ ffi::Ffi::Bind()
367+ .Ctx<ffi::PlatformStream<cudaStream_t>>()
368+ .Ctx<ffi::ScratchAllocator>()
369+ .Arg<ffi::Buffer<ffi::F32>>() // means
370+ .Arg<ffi::Buffer<ffi::F32>>() // scales
371+ .Arg<ffi::Buffer<ffi::F32>>() // quats
372+ .Arg<ffi::Buffer<ffi::F32>>() // opacities
373+ .Arg<ffi::Buffer<ffi::F32>>() // sh
374+ .Arg<ffi::Buffer<ffi::F32>>() // camera
375+ .Ret<ffi::Buffer<ffi::F32>>() // image
376+ .Ret<ffi::Buffer<ffi::F32>>() // splats
377+ .Ret<ffi::Buffer<ffi::S32>>() // point_list
378+ .Ret<ffi::Buffer<ffi::S32>>() // ranges
379+ .Ret<ffi::Buffer<ffi::F32>>() // final_t
380+ .Ret<ffi::Buffer<ffi::S32>>() // n_contrib
381+ .Ret<ffi::Buffer<ffi::S32>>() // stats
382+ .Ret<ffi::Buffer<ffi::F32>>() // timings
383+ .Attr<int32_t>("width")
384+ .Attr<int32_t>("height")
385+ .Attr<int32_t>("culling")
386+ .Attr<int32_t>("profile")
387+ .Attr<float>("fx")
388+ .Attr<float>("fy"));
389+
390+// ---------------------------------------------------------------------------
391+
392+ffi::Error Backward(
393+ cudaStream_t stream, ffi::ScratchAllocator scratch,
394+ ffi::Buffer<ffi::F32> means, ffi::Buffer<ffi::F32> scales,
395+ ffi::Buffer<ffi::F32> quats, ffi::Buffer<ffi::F32> opacities,
396+ ffi::Buffer<ffi::F32> sh, ffi::Buffer<ffi::F32> camera,
397+ ffi::Buffer<ffi::F32> splats, ffi::Buffer<ffi::S32> point_list,
398+ ffi::Buffer<ffi::S32> ranges, ffi::Buffer<ffi::F32> final_t,
399+ ffi::Buffer<ffi::S32> n_contrib, ffi::Buffer<ffi::F32> dl_dimage,
400+ ffi::ResultBuffer<ffi::F32> dmeans, ffi::ResultBuffer<ffi::F32> dscales,
401+ ffi::ResultBuffer<ffi::F32> dquats, ffi::ResultBuffer<ffi::F32> dopacities,
402+ ffi::ResultBuffer<ffi::F32> dsh, ffi::ResultBuffer<ffi::F32> timings,
403+ int32_t width, int32_t height, int32_t profile, float fx_attr,
404+ float fy_attr) {
405+ (void)opacities;
406+ const int n = static_cast<int>(means.element_count() / 3);
407+ if (camera.element_count() != kCameraFloats) {
408+ return ffi::Error(ffi::ErrorCode::kInvalidArgument,
409+ "camera must have 20 floats");
410+ }
411+ if (n <= 0) {
412+ return ffi::Error(ffi::ErrorCode::kInvalidArgument, "no Gaussians");
413+ }
414+
415+ DefaultStreamBridge bridge(stream);
416+ auto cam_buf = Alloc<float>(scratch, 40);
417+ auto dl_dmean2d = Alloc<float3>(scratch, n);
418+ auto dl_dconic2d = Alloc<float4>(scratch, n);
419+ auto dl_dcolor = Alloc<float>(scratch, 3 * n);
420+ auto dl_dcov3d = Alloc<float>(scratch, 6 * n);
421+ if (!cam_buf || !dl_dmean2d || !dl_dconic2d || !dl_dcolor || !dl_dcov3d) {
422+ return ffi::Error(ffi::ErrorCode::kResourceExhausted, "backward scratch");
423+ }
424+
425+ PrepareCameraKernel<<<1, 1>>>(camera.typed_data(), width, height, *cam_buf);
426+ float focal_x = 0.0f, focal_y = 0.0f;
427+ ResolveFocal(camera.typed_data(), fx_attr, fy_attr, &focal_x, &focal_y);
428+ const float tan_fovx = static_cast<float>(width) / (2.0f * focal_x);
429+ const float tan_fovy = static_cast<float>(height) / (2.0f * focal_y);
430+
431+ const float* viewmatrix = *cam_buf;
432+ const float* projmatrix = *cam_buf + 16;
433+ const auto* cam_pos = reinterpret_cast<const glm::vec3*>(*cam_buf + 32);
434+ const float* bg_color = *cam_buf + 35;
435+
436+ const dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X,
437+ (height + BLOCK_Y - 1) / BLOCK_Y, 1);
438+ const dim3 block(BLOCK_X, BLOCK_Y, 1);
439+ const SplatsView geom =
440+ SplatsView::From(const_cast<float*>(splats.typed_data()), n);
441+
442+ StageTimer timer(nullptr, profile != 0);
443+
444+ GS_RETURN_IF_CUDA_ERROR(
445+ cudaMemsetAsync(*dl_dmean2d, 0, sizeof(float3) * n, nullptr), "memset");
446+ GS_RETURN_IF_CUDA_ERROR(
447+ cudaMemsetAsync(*dl_dconic2d, 0, sizeof(float4) * n, nullptr), "memset");
448+ GS_RETURN_IF_CUDA_ERROR(
449+ cudaMemsetAsync(*dl_dcolor, 0, sizeof(float) * 3 * n, nullptr), "memset");
450+ GS_RETURN_IF_CUDA_ERROR(
451+ cudaMemsetAsync(*dl_dcov3d, 0, sizeof(float) * 6 * n, nullptr), "memset");
452+ GS_RETURN_IF_CUDA_ERROR(
453+ cudaMemsetAsync(dopacities->typed_data(), 0, sizeof(float) * n, nullptr),
454+ "memset");
455+
456+ // 1. Backward rasterize.
457+ BACKWARD::render(tile_grid, block,
458+ reinterpret_cast<const uint2*>(ranges.typed_data()),
459+ reinterpret_cast<const uint32_t*>(point_list.typed_data()),
460+ width, height, bg_color, geom.means2D, geom.conic_opacity,
461+ geom.rgb, final_t.typed_data(),
462+ reinterpret_cast<const uint32_t*>(n_contrib.typed_data()),
463+ dl_dimage.typed_data(), *dl_dmean2d, *dl_dconic2d,
464+ dopacities->typed_data(), *dl_dcolor);
465+ timer.Mark();
466+
467+ // 2. Backward preprocess (as Inria's PyTorch wrapper `torch::zeros`: bulk
468+ // zero parameter gradient tensors before `BACKWARD::preprocess`).
469+ GS_RETURN_IF_CUDA_ERROR(
470+ cudaMemsetAsync(dmeans->typed_data(), 0, sizeof(float) * 3 * n, nullptr),
471+ "memset");
472+ GS_RETURN_IF_CUDA_ERROR(
473+ cudaMemsetAsync(dscales->typed_data(), 0, sizeof(float) * 3 * n, nullptr),
474+ "memset");
475+ GS_RETURN_IF_CUDA_ERROR(
476+ cudaMemsetAsync(dquats->typed_data(), 0, sizeof(float) * 4 * n, nullptr),
477+ "memset");
478+ GS_RETURN_IF_CUDA_ERROR(
479+ cudaMemsetAsync(dsh->typed_data(), 0, sizeof(float) * 3 * kShCoeffs * n,
480+ nullptr),
481+ "memset");
482+
483+ BACKWARD::preprocess(n, kShDegree, kShCoeffs,
484+ reinterpret_cast<const float3*>(means.typed_data()),
485+ geom.radii, sh.typed_data(), geom.clamped,
486+ reinterpret_cast<const glm::vec3*>(scales.typed_data()),
487+ reinterpret_cast<const glm::vec4*>(quats.typed_data()),
488+ kScaleModifier, geom.cov3Ds, geom.conic_opacity,
489+ viewmatrix, projmatrix, focal_x, focal_y, tan_fovx,
490+ tan_fovy, kLowPassFilter, false, cam_pos, *dl_dmean2d,
491+ reinterpret_cast<const float*>(*dl_dconic2d),
492+ reinterpret_cast<glm::vec3*>(dmeans->typed_data()),
493+ *dl_dcolor, *dl_dcov3d, dsh->typed_data(),
494+ reinterpret_cast<glm::vec3*>(dscales->typed_data()),
495+ reinterpret_cast<glm::vec4*>(dquats->typed_data()),
496+ dopacities->typed_data());
497+ timer.Mark();
498+ timer.Write(timings->typed_data());
499+ GS_RETURN_IF_CUDA_ERROR(cudaGetLastError(), "backward launch");
500+ return ffi::Error::Success();
501+}
502+
503+XLA_FFI_DEFINE_HANDLER(kBackward, Backward,
504+ ffi::Ffi::Bind()
505+ .Ctx<ffi::PlatformStream<cudaStream_t>>()
506+ .Ctx<ffi::ScratchAllocator>()
507+ .Arg<ffi::Buffer<ffi::F32>>() // means
508+ .Arg<ffi::Buffer<ffi::F32>>() // scales
509+ .Arg<ffi::Buffer<ffi::F32>>() // quats
510+ .Arg<ffi::Buffer<ffi::F32>>() // opacities
511+ .Arg<ffi::Buffer<ffi::F32>>() // sh
512+ .Arg<ffi::Buffer<ffi::F32>>() // camera
513+ .Arg<ffi::Buffer<ffi::F32>>() // splats
514+ .Arg<ffi::Buffer<ffi::S32>>() // point_list
515+ .Arg<ffi::Buffer<ffi::S32>>() // ranges
516+ .Arg<ffi::Buffer<ffi::F32>>() // final_t
517+ .Arg<ffi::Buffer<ffi::S32>>() // n_contrib
518+ .Arg<ffi::Buffer<ffi::F32>>() // dl_dimage
519+ .Ret<ffi::Buffer<ffi::F32>>() // dmeans
520+ .Ret<ffi::Buffer<ffi::F32>>() // dscales
521+ .Ret<ffi::Buffer<ffi::F32>>() // dquats
522+ .Ret<ffi::Buffer<ffi::F32>>() // dopacities
523+ .Ret<ffi::Buffer<ffi::F32>>() // dsh
524+ .Ret<ffi::Buffer<ffi::F32>>() // timings
525+ .Attr<int32_t>("width")
526+ .Attr<int32_t>("height")
527+ .Attr<int32_t>("profile")
528+ .Attr<float>("fx")
529+ .Attr<float>("fy"));
530+
531+XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "gs_forward", "CUDA", kForward);
532+XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "gs_backward", "CUDA", kBackward);
533+
534+} // namespace
535+} // namespace gs
Unchanged Inria CUDA files in s00_baseline (100% identical to diff-gaussian-rasterization/cuda_rasterizer/): csrc/config.h, csrc/auxiliary.h, csrc/forward.h, csrc/forward.cu.cc, csrc/backward.h, csrc/backward.cu.cc, csrc/rasterizer.h, csrc/rasterizer_impl.h, csrc/rasterizer_impl.cu.cc