[ROCm] Add AMD GPU support for TRELLIS.2

o-voxel/setup.py

@@ -19,7 +19,7 @@
 if not IS_HIP:
     cc_flag = []
 else:
-    archs = os.getenv("GPU_ARCHS", "native").split(";")
+    archs = os.getenv("GPU_ARCHS", "gfx942").split(";")
     cc_flag = [f"--offload-arch={arch}" for arch in archs]
 
 setup(

o-voxel/src/hash/hash.cu

@@ -1,6 +1,10 @@
 #include <torch/extension.h>
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
+#endif
 
 #include "api.h"
 #include "hash.cuh"

o-voxel/src/rasterize/rasterize.cu

@@ -1,6 +1,10 @@
 #include <cstdint>
 
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
+#endif
 #include "cuda_runtime.h"
 
 #include <cooperative_groups.h>

o-voxel/src/serialize/hilbert.cu

@@ -1,5 +1,9 @@
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
+#endif
 
 #include <cooperative_groups.h>
 namespace cg = cooperative_groups;

o-voxel/src/serialize/z_order.cu

@@ -1,5 +1,9 @@
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
+#endif
 
 #include <cooperative_groups.h>
 namespace cg = cooperative_groups;