Lazy load cuda (#4)

.readthedocs.yaml

 version: 2
 
+build:
+  os: ubuntu-20.04
+  tools:
+    python: "3.9"
+
 sphinx:
   fail_on_warning: true
+  configuration: docs/source/conf.py
 
 python:
-  version: 3.8
   install:
     # Equivalent to 'pip install .'
     - method: pip

nerfacc/cuda/__init__.py

-from ._backend import _C
+from typing import Callable
+
+def _make_lazy_cuda(name: str) -> Callable:
+    def call_cuda(*args, **kwargs):
+        # pylint: disable=import-outside-toplevel
+        from ._backend import _C
+
+        return getattr(_C, name)(*args, **kwargs)
+
+    return call_cuda
+
+ray_aabb_intersect = _make_lazy_cuda("ray_aabb_intersect")
+volumetric_marching = _make_lazy_cuda("volumetric_marching")
+volumetric_rendering_steps = _make_lazy_cuda("volumetric_rendering_steps")
+volumetric_rendering_weights_forward = _make_lazy_cuda("volumetric_rendering_weights_forward")
+volumetric_rendering_weights_backward = _make_lazy_cuda("volumetric_rendering_weights_backward")
\ No newline at end of file

nerfacc/utils.py

@@ -2,7 +2,7 @@ from typing import Tuple
 
 import torch
 
-from .cuda import _C
+import nerfacc.cuda as nerfacc_cuda
 
 
 @torch.no_grad()
@@ -27,7 +27,7 @@ def ray_aabb_intersect(
         rays_o = rays_o.contiguous()
         rays_d = rays_d.contiguous()
         aabb = aabb.contiguous()
-        t_min, t_max = _C.ray_aabb_intersect(rays_o, rays_d, aabb)
+        t_min, t_max = nerfacc_cuda.ray_aabb_intersect(rays_o, rays_d, aabb)
     else:
         raise NotImplementedError("Only support cuda inputs.")
     return t_min, t_max
@@ -87,7 +87,7 @@ def volumetric_marching(
         frustum_dirs,
         frustum_starts,
         frustum_ends,
-    ) = _C.volumetric_marching(
+    ) = nerfacc_cuda.volumetric_marching(
         # rays
         rays_o.contiguous(),
         rays_d.contiguous(),
@@ -152,7 +152,7 @@ def volumetric_rendering_steps(
         frustum_starts = frustum_starts.contiguous()
         frustum_ends = frustum_ends.contiguous()
         sigmas = sigmas.contiguous()
-        compact_packed_info, compact_selector = _C.volumetric_rendering_steps(
+        compact_packed_info, compact_selector = nerfacc_cuda.volumetric_rendering_steps(
             packed_info, frustum_starts, frustum_ends, sigmas
         )
         compact_frustum_starts = frustum_starts[compact_selector]
@@ -261,7 +261,7 @@ def volumetric_rendering_accumulate(
 class _volumetric_rendering_weights(torch.autograd.Function):
     @staticmethod
     def forward(ctx, packed_info, frustum_starts, frustum_ends, sigmas):
-        weights, ray_indices = _C.volumetric_rendering_weights_forward(
+        weights, ray_indices = nerfacc_cuda.volumetric_rendering_weights_forward(
             packed_info, frustum_starts, frustum_ends, sigmas
         )
         ctx.save_for_backward(
@@ -282,7 +282,7 @@ class _volumetric_rendering_weights(torch.autograd.Function):
             sigmas,
             weights,
         ) = ctx.saved_tensors
-        grad_sigmas = _C.volumetric_rendering_weights_backward(
+        grad_sigmas = nerfacc_cuda.volumetric_rendering_weights_backward(
             weights,
             grad_weights,
             packed_info,