Fast SH implementation

[jb-ye]

I recently read this paper: https://jcgt.org/published/0002/02/06/ (it has sample codes) there is a faster way to evaluate spherical harmonics. I made an attempt to implement the idea in gsplat and was able to observe 40% time saving in forward calls of degree 3 and 50% time saving of degree 4.

[liruilong940607]

This is very cool! And I also verified the diff is basically zero comparing with the original one.

basis_dim = 25
dirs = torch.rand((100000, 3), device="cuda")
dirs = dirs / dirs.norm(dim=-1, keepdim=True)
result1 = eval_sh_bases(basis_dim, dirs)
result2 = eval_sh_bases_fast(basis_dim, dirs)
diff = (result1 - result2).abs().mean(dim=0)
>> tensor([0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 7.1335e-09,
        2.3338e-08, 7.0950e-09, 0.0000e+00, 1.2602e-08, 8.1815e-09, 1.8609e-08,
        3.7991e-08, 1.8580e-08, 0.0000e+00, 1.2594e-08, 1.0718e-08, 1.6333e-08,
        1.4538e-08, 2.3519e-08, 3.1723e-08, 2.3399e-08, 1.2196e-08, 1.6297e-08,
        1.3505e-08], device='cuda:0')

[liruilong940607]

@jb-ye I add some profile code in the test_sh.py, and this is what I get by running python tests/test_sh.py on a Tesla V100-SXM2:

[Fwd] Method: poly, ellipsed: 2.19 ms
[Bwd] Method: poly, ellipsed: 35.19 ms
[Fwd] Method: fast, ellipsed: 2.46 ms
[Bwd] Method: fast, ellipsed: 35.16 ms

Somehow I couldn't get speedup on my end. Could you run the profiler on your end and post the results?

[jb-ye]

@liruilong940607 Your previous profile code has a typo ... after fixing it, this is what I got:

[Fwd] Method: poly, ellipsed: 5.58 ms
[Bwd] Method: poly, ellipsed: 31.67 ms
[Fwd] Method: fast, ellipsed: 2.02 ms
[Bwd] Method: fast, ellipsed: 31.47 ms

[liruilong940607]

Oh yeah you are right thanks for the fix! It now makes sense.

[liruilong940607]

LGTM! A few minor comments

[liruilong940607]

is it intentional that this 0.4886025119029199 does not equal to -fTmpA?

[jb-ye]

This is just for matching the sample code from the original authors. They should be equal to -fTmpA.

[ichsan2895]

How to activate this feature with splatfacto? @jb-ye ?

I did not find any speed improvement, maybe I did it wrong because it was not activated?

black line = newest main gsplat (included this merged PR)
red line = nerfstudio v1.0.3 and gsplat 0.1.10

[jb-ye]

@ichsan2895 It doesn't bring speed-ups for training since the bwd calls take rough the same time. It may increase ns-viewer fps rate.