fix: use JIT return module for CUDA 12.8.1 (Blackwell) to avoid import failures (#166)

* Fix JIT extension loading for CUDA 12.8.1/Blackwell: use returned module instead of re-import

---------

Co-authored-by: Qi Wu <wilson.over.cloud@gmail.com>

Author: Niloux

threedgrt_tracer/setup_3dgrt.py

@@ -71,7 +71,7 @@ def to_cpp_bool(value):
             f"{os.path.join(slang_build_file_dir,'models/gaussianParticles.slang')}",
             f"{os.path.join(slang_build_file_dir,'models/shRadiativeParticles.slang')}",
             "-o",
-            f"{os.path.join(slang_build_file_dir,'gaussianParticles.cuh')}",
+            f"{os.path.join(slang_build_file_dir, 'gaussianParticles.cuh')}",
         ],
         env=slang_build_env,
     )
@@ -84,4 +84,4 @@ def to_cpp_bool(value):
         extra_cflags=cflags,
         extra_cuda_cflags=cuda_flags,
         extra_include_paths=include_paths,
-    )
+    )

threedgrut_playground/setup_playground.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import os
-import sys
 from pathlib import Path
 
 from threedgrut.utils import jit
@@ -72,11 +71,11 @@ def to_cpp_bool(value):
             f"-DGAUSSIAN_PARTICLE_MIN_ALPHA={conf.render.particle_kernel_min_alpha}",
             f"-DGAUSSIAN_PARTICLE_MAX_ALPHA={conf.render.particle_kernel_max_alpha}",
             f"-DGAUSSIAN_PARTICLE_ENABLE_NORMAL={to_cpp_bool(conf.render.enable_normals)}",
-            f"-DGAUSSIAN_PARTICLE_SURFEL={to_cpp_bool(conf.render.primitive_type=='trisurfel')}",
-            f"{os.path.join(slang_build_file_dir,'models/gaussianParticles.slang')}",
-            f"{os.path.join(slang_build_file_dir,'models/shRadiativeParticles.slang')}",
+            f"-DGAUSSIAN_PARTICLE_SURFEL={to_cpp_bool(conf.render.primitive_type == 'trisurfel')}",
+            f"{os.path.join(slang_build_file_dir, 'models/gaussianParticles.slang')}",
+            f"{os.path.join(slang_build_file_dir, 'models/shRadiativeParticles.slang')}",
             "-o",
-            f"{os.path.join(slang_build_file_dir,'gaussianParticles.cuh')}",
+            f"{os.path.join(slang_build_file_dir, 'gaussianParticles.cuh')}",
         ],
         env=slang_build_env,
     )
@@ -87,4 +86,4 @@ def to_cpp_bool(value):
         name="libplayground_cc",
         sources=source_paths,
         extra_include_paths=include_paths,
-    )
+    )

threedgrut_playground/tracer.py

@@ -15,7 +15,6 @@
 
 import logging
 import os
-from enum import IntEnum
 
 import torch
 import torch.utils.cpp_extension
@@ -46,9 +45,7 @@ def load_playground_plugin(conf):
 
 
 class Tracer:
-
     def __init__(self, conf):
-
         self.device = "cuda"
         self.conf = conf
         self.num_update_bvh = 0

threedgut_tracer/setup_3dgut.py

@@ -24,7 +24,6 @@
 # ----------------------------------------------------------------------------
 #
 def setup_3dgut(conf):
-
     build_dir = torch.utils.cpp_extension._get_build_directory("lib3dgut_cc", verbose=True)
 
     include_paths = []
@@ -126,9 +125,9 @@ def to_cpp_bool(value):
             "-Wno-41018",
             "-O2",
             *defines,
-            f"{os.path.join(slang_build_inc_dir,'threedgut.slang')}",
+            f"{os.path.join(slang_build_inc_dir, 'threedgut.slang')}",
             "-o",
-            f"{os.path.join(build_dir,'threedgutSlang.cuh')}",
+            f"{os.path.join(build_dir, 'threedgutSlang.cuh')}",
         ],
         env=slang_build_env,
     )
@@ -142,4 +141,4 @@ def to_cpp_bool(value):
         extra_cuda_cflags=cuda_cflags,
         extra_include_paths=include_paths,
         build_directory=build_dir,
-    )
+    )

threedgut_tracer/tracer.py

@@ -19,7 +19,6 @@
 
 import numpy as np
 import torch
-import torch.utils.cpp_extension
 from omegaconf import OmegaConf
 
 from threedgrut.datasets.protocols import Batch