Skip to content

Add experimental GEMM forward rasterization backends for 3DGS and 2DGS#935

Open
zhubowen1 wants to merge 1 commit into
nerfstudio-project:mainfrom
zhubowen1:feat/gemm-raster-fwd
Open

Add experimental GEMM forward rasterization backends for 3DGS and 2DGS#935
zhubowen1 wants to merge 1 commit into
nerfstudio-project:mainfrom
zhubowen1:feat/gemm-raster-fwd

Conversation

@zhubowen1

Copy link
Copy Markdown

This PR adds experimental GEMM-based forward rasterization backends for both 3DGS and 2DGS.

Background

The existing rasterization pipeline relies on the original per-pixel forward rasterization kernels. This PR introduces an alternative forward backend designed to accelerate rasterization by reformulating key parts of the computation into GEMM-style, tensor-core-friendly workloads.

Changes

  • Added a GEMM-based forward rasterization kernel for 3DGS in RasterizeToPixels3DGSFwd.cu.
  • Added a GEMM-based forward rasterization kernel for 2DGS in RasterizeToPixels2DGSFwd.cu.
  • Added shared low-level GEMM raster helpers in GemmRasterUtils.cuh.
  • Added explicit backend selection plumbing through Python, C++, and CUDA: rasterize_fwd_impl="gemm"

Evaluation Setup

I evaluated the GEMM backend on seven real scenes: garden, bicycle, stump, bonsai, counter, kitchen, and room.
These scenes cover both outdoor and indoor settings and vary in scene complexity. I report rendering quality measured by PSNR, and performance measured by the latency of the forward rasterization kernel only.

Performance Summary

Across these seven scenes, the GEMM backend consistently reduces the latency of the forward rasterization kernel on both RTX 3090 and Jetson AGX Orin, while maintaining nearly identical PSNR.
Note that the reported latency refers only to the single forward rasterization kernel, rather than the full end-to-end rendering pipeline.

Platform Backend Avg legacy latency Avg GEMM latency Speedup Latency reduction
RTX 3090 3DGS 2.61 ms 2.00 ms 1.33x 23.30%
RTX 3090 2DGS 4.80 ms 3.29 ms 1.47x 31.30%
Jetson AGX Orin 3DGS 51.79 ms 29.18 ms 1.87x 43.70%
Jetson AGX Orin 2DGS 87.89 ms 51.84 ms 1.70x 41.00%

Quality

The GEMM backend preserves rendering quality closely compared with the legacy backend.

Backend Avg PSNR change Max absolute PSNR change
3DGS -0.019 dB 0.041 dB
2DGS -0.004 dB 0.022 dB

Known Limitations

  • Forward-only support
    The GEMM rasterization backend currently supports forward inference only. Backward/gradient computation is not implemented yet.
  • Half-precision computation
    The implementation relies on Tensor Core (MMA) operations in FP16. While this enables significant speedup, it may introduce minor numerical differences compared to the FP32 legacy backend.
  • Fixed tile size
    The current implementation only supports a tile size of 16×16.
  • Limited channel configuration
    Only CDIM = 3 is supported in the current version.
  • Hardware requirement
    Requires GPUs with Tensor Core support (SM 7.5 and above).

@zhubowen1 zhubowen1 closed this Apr 27, 2026
@zhubowen1 zhubowen1 reopened this Apr 27, 2026
@zhubowen1

Copy link
Copy Markdown
Author

This GEMM-based rasterization backend is related to our recent work:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant