Featured image of post watch: 3DGS | AI葵 Cuda Code Walkthrough

watch: 3DGS | AI葵 Cuda Code Walkthrough

forward.cu

Trade-off:

  1. The 2D-projections of 3D ellipsoids are circles rather than ellipses to reduce shading tiles.

  2. Determine visibility of tiles instead of pixels (shaded by circle shadows) for fast rasterization.

Main steps:

  1. Determine the radii of circles shadows (preprocessCUDA)

    1. Project a 3D ellipsoid will yield a 2x2 covariance matrix

    2. Solve the 2 Eigenvalues for the covariance matrix, and the bigger one is the length of the major axis.

    3. Use the major axis as the radius of the circle.

  2. Determine pixels covered by the projected circles.

    If the distance betwen a pixel to the circle center is smaller than the circle radius, the pixel is visible to the circle (a disc corresponding to a 3D Gaussian in the ray space).

    1. GetRect
  3. Sort 3D Gaussians (discs) by depths

    1. Each tile is shaded by multiple ellipsoids, i.e., visible to multiple Gaussians.

    2. Pair each pixel with each contributing ellipsoid, and form a 64-bit identifier for each pair.

      i.e., stitching the index (32-bit) of a pixel and the depth (32-bit) of a ellipsoid.

    3. Sort all the identifier, and obtain a sequence

    For example:

    0 2 c a b 1 3
    • tile-0 has 2 pairs: 0-a, 0-c
    • tile-1 has 3 pairs: 1-a, 1-b, 1-c

    Suppose the depths of 3 Gaussians are b > a > c, the sequence of pairs is as follows:

    1
    2
    3
    4
    5
    
    0-c
    0-a
    1-c
    1-a
    1-b
    
  4. Alpha compositing for each pixel’s color (renderCUDA)

    Implementation tricks:

    • A tile is a Block, in which each pixel is a worker.
    • Every pixel in a tile uses the same Gaussian distributions, so those data are stored in __share__ memory. Docs
    1. Calculate alpha, which is proportional to the probability in a Gaussian distribution.

    2. The probability is calculated according to the expression of 2D Gaussian.

    3. Blend alpha * color (SH) of each ellipsoid front-to-back.


backward.cu

output input
pixels’ color (p,3) Gaussians’ color (g,3)
Gaussians’ alpha
Gaussians’ position (g,3)
Gaussians’ rotation (g,4)
Gaussians’ length of axis (g,3)
  • p is number of pixels; g is number of 3D Gaussians

Loss = color + SSIM

  • A full image is produced at once, so image metrics, like SSIM, can be added.

Write parital derivative for each input tensor.