Add RTX 50 series GPU compatibility via PyTorch 2.7 and CUDA 12.8

Utils.py

@@ -43,7 +43,7 @@
 except:
   kornia = None
 try:
-  import mycpp.build.mycpp as mycpp
+  import mycpp
 except:
   mycpp = None
 try:

bundlesdf/mycuda/common.cu

@@ -117,7 +117,7 @@ at::Tensor sampleRaysUniformOccupiedVoxels(const at::Tensor z_in_out,  const at:
   const int threadx = 32;
   const int thready = 32;
 
-  AT_DISPATCH_FLOATING_TYPES(z_in_out.type(), "sample_rays_uniform_occupied_voxels_kernel", ([&]
+  AT_DISPATCH_FLOATING_TYPES(z_in_out.scalar_type(), "sample_rays_uniform_occupied_voxels_kernel", ([&]
   {
     sample_rays_uniform_occupied_voxels_kernel<scalar_t><<<{divCeil(N_rays,threadx),divCeil(N_samples,thready)}, {threadx,thready}>>>(z_sampled.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),z_in_out.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),z_vals.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
   }));
@@ -159,7 +159,7 @@ at::Tensor postprocessOctreeRayTracing(const at::Tensor ray_index, const at::Ten
   at::Tensor depths_in_out_padded = at::zeros({N_rays,max_intersections,2}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0).requires_grad(false));
   dim3 threads = {256};
   dim3 blocks = {divCeil(n_unique_ids,threads.x)};
-  AT_DISPATCH_FLOATING_TYPES(depth_in_out.type(), "postprocessOctreeRayTracingKernel", ([&]
+  AT_DISPATCH_FLOATING_TYPES(depth_in_out.scalar_type(), "postprocessOctreeRayTracingKernel", ([&]
   {
     postprocessOctreeRayTracingKernel<scalar_t><<<blocks,threads>>>(ray_index.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depth_in_out.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), unique_intersect_ray_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), start_poss.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depths_in_out_padded.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
   }));
@@ -265,7 +265,7 @@ void rayColorToTextureImageCUDA(const at::Tensor &F, const at::Tensor &V, const
   dim3 threads = {512};
   dim3 blocks = {divCeil(int(hit_locations.sizes()[0]),threads.x)};
 
-  AT_DISPATCH_FLOATING_TYPES(V.type(), "rayColorToTextureImageKernel", ([&]
+  AT_DISPATCH_FLOATING_TYPES(V.scalar_type(), "rayColorToTextureImageKernel", ([&]
   {
     rayColorToTextureImageKernel<scalar_t><<<blocks,threads>>>(F.packed_accessor32<long,2,torch::RestrictPtrTraits>(), V.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_locations.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_face_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), uvs_tex.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), uvs.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
   }));

bundlesdf/mycuda/setup.py

@@ -15,8 +15,8 @@
 code_dir = os.path.dirname(os.path.realpath(__file__))
 
 
-nvcc_flags = ['-Xcompiler', '-O3', '-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__']
-c_flags = ['-O3', '-std=c++14']
+nvcc_flags = ['-Xcompiler', '-O3', '-std=c++17', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__']
+c_flags = ['-O3', '-std=c++17']
 
 setup(
     name='common',

mycpp/src/app/pybind_api.cpp

@@ -23,7 +23,7 @@ namespace py = pybind11;
 //@dist_diff: unit is meter
 vectorMatrix4f cluster_poses(float angle_diff, float dist_diff, const vectorMatrix4f &poses_in, const vectorMatrix4f &symmetry_tfs)
 {
-  printf("num original candidates = %d\n",poses_in.size());
+  printf("num original candidates = %zu\n",poses_in.size());
   vectorMatrix4f poses_out;
   poses_out.push_back(poses_in[0]);
 
@@ -63,7 +63,7 @@ vectorMatrix4f cluster_poses(float angle_diff, float dist_diff, const vectorMatr
     }
   }
 
-  printf("num of pose after clustering: %d\n",poses_out.size());
+  printf("num of pose after clustering: %zu\n",poses_out.size());
   return poses_out;
 }
 

readme.md

@@ -123,6 +123,38 @@ python -m pip install --quiet --no-index --no-cache-dir pytorch3d -f https://dl.
 CMAKE_PREFIX_PATH=$CONDA_PREFIX/lib/python3.9/site-packages/pybind11/share/cmake/pybind11 bash build_all_conda.sh
 ```
 
+## Env setup for RTX 50 series GPU
+Due to the `sm_120` architecture, Pytorch binaries only support Cuda `12.8+` and PyTorch `2.7.0+` (see [here](https://discuss.pytorch.org/t/pytorch-support-for-sm120/216099)). Therefore, the provided Docker images (including the one for RTX 40 series) won't work. Until a Docker image is provided, it is still possible to set up FoundationPose using the Conda instructions on Python `3.9`. The main tricks are:
+
+1. Make sure Cuda 12.8 is installed.
+1. PyTorch3D does not currently have a prebuilt binary for `sm_120`([GitHub issue](https://github.com/facebookresearch/pytorch3d/issues/1970)). Just [install it from source](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#building--installing-from-source) manually.
+1. There may be build failures in the "Build extensions" step. See the troubleshooting section.
+
+### Troubleshooting
+- If the following occurs, there is a mismatch between PyTorch / Cuda /GPU version. Note that it is *possible* to install Cuda 12.0~12.7 on a RTX50 series card and then install some earlier Pytorch binaries, it will just throw the error below upon running. Install Cuda `12.8+` and PyTorch `2.7.0+` to fix this.
+```
+RuntimeError: CUDA error: no kernel image is available for execution on the device
+CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
+For debugging consider passing CUDA_LAUNCH_BLOCKING=1
+```
+
+- An error message like this indicates a problem with PyTorch. See [this issue](https://github.com/facebookresearch/pytorch3d/issues/1315). This can be fixed by installing from source.
+```
+".../python3.9/site-packages/pytorch3d/renderer/blending.py", line 12, in <module>
+    from pytorch3d import _C
+ImportError: .../python3.9/site-packages/pytorch3d/_C.cpython-39-x86_64-linux-gnu.so: undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs
+```
+
+- As suggested by [this issue](https://github.com/NVlabs/FoundationPose/issues/288), the following error from running `python run_demo.py` suggests a build issue in `mycpp`. Run the commands in `build_all_conda.sh` manually and make sure `mycpp` is built correctly.
+```
+    rot_grid = mycpp.cluster_poses(30, 99999, rot_grid, self.symmetry_tfs.data.cpu().numpy())
+AttributeError: 'NoneType' object has no attribute 'cluster_poses'
+```
+
+- This error from building `mycuda` occurs specifically on RTX50 series cards due to using deprecated API. The fix is to modify `bundlesdf/mycuda/common.cu` to use the new API. This is included in commit `25fe3b`.
+```
+.../FoundationPose/bundlesdf/mycuda/common.cu", static_cast<uint32_t>(268), (::c10::detail::torchCheckMsgImpl( "Expected " "false" " to be true, but got false.  " "(Could this error message be improved?  If so, " "please report an enhancement request to PyTorch.)", '"', at_dispatch_name, "\" not implemented for '", toString(_st), "'"))); }; } }()
+```
 
 # Run model-based demo
 The paths have been set in argparse by default. If you need to change the scene, you can pass the args accordingly. By running on the demo data, you should be able to see the robot manipulating the mustard bottle. Pose estimation is conducted on the first frame, then it automatically switches to tracking mode for the rest of the video. The resulting visualizations will be saved to the `debug_dir` specified in the argparse. (Note the first time running could be slower due to online compilation)

requirements.txt

@@ -1,10 +1,8 @@
-# PyTorch 2.1.1
---extra-index-url https://download.pytorch.org/whl/cu118
-# torch==2.1.1+cu118
-# torchvision==0.16.1+cu118
-torch==2.0.0+cu118
-torchvision==0.15.1+cu118
-torchaudio==2.0.1+cu118
+# PyTorch 2.7 for RTX 5080 CUDA 12.8 compatibility
+--extra-index-url https://download.pytorch.org/whl/cu128
+torch==2.7.1+cu128
+torchaudio==2.7.1+cu128
+torchvision==0.22.1+cu128
 
 
 # JupyterLab