Add experimental GEMM forward rasterization backends for 3DGS and 2DGS#935
Open
zhubowen1 wants to merge 1 commit into
Open
Add experimental GEMM forward rasterization backends for 3DGS and 2DGS#935zhubowen1 wants to merge 1 commit into
zhubowen1 wants to merge 1 commit into
Conversation
Author
|
This GEMM-based rasterization backend is related to our recent work:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
Evaluation Setup
I evaluated the GEMM backend on seven real scenes:
garden,bicycle,stump,bonsai,counter,kitchen, androom.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.
Quality
The GEMM backend preserves rendering quality closely compared with the legacy backend.
Known Limitations
The GEMM rasterization backend currently supports forward inference only. Backward/gradient computation is not implemented yet.
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.
The current implementation only supports a tile size of 16×16.
Only CDIM = 3 is supported in the current version.
Requires GPUs with Tensor Core support (SM 7.5 and above).