Fix: Add conda CUDA include path for JIT compilation on Linux

[Odysseusailoon]

Fix: Add conda CUDA include path for JIT compilation on Linux

Description

This PR fixes a Just-In-Time (JIT) compilation failure that occurs when trying to use nvdiffrast on Linux systems where the CUDA toolkit was installed via the nvidia channel in a conda environment. This is becoming a common way to manage CUDA dependencies.

Problem

The cuda-toolkit package provided by the nvidia conda channel installs CUDA header files into a non-standard, target-specific directory (e.g., $CONDA_PREFIX/targets/x86_64-linux/include), rather than the standard $CONDA_PREFIX/include.

When nvdiffrast attempts its JIT compilation using torch.utils.cpp_extension.load, it fails to find necessary headers like cuda_runtime_api.h because:

1. The build system doesn't automatically search this non-standard path.
2. The build system appears to ignore standard environment variables (CPPFLAGS, CXXFLAGS, CUDA_HOME) that are typically used to add include directories, preventing a simple workaround.

This results in compilation errors like fatal error: cuda_runtime_api.h: No such file or directory or fatal error: crt/host_defines.h: No such file or directory.

Solution

This PR modifies the _get_plugin function in nvdiffrast/torch/ops.py. It now includes logic to proactively check for this specific conda setup:

1. Check if running within a conda environment (CONDA_PREFIX environment variable is set).
2. Check if the operating system is Linux (platform.system() == 'Linux').
3. Construct the expected non-standard include path ($CONDA_PREFIX/targets/x86_64-linux/include).
4. If this path exists as a directory, add it to the extra_include_paths list argument passed directly to torch.utils.cpp_extension.load.

This directly informs the build system about the correct location of the CUDA headers in this specific environment setup, resolving the compilation failures.

How to Test

1. Set up a Linux conda environment.
2. Install PyTorch with CUDA support (e.g., conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia).
3. Install cuda-toolkit from the nvidia channel (e.g., conda install cuda-toolkit -c nvidia). Note: Ensure the version matches or is compatible with the PyTorch CUDA version.
4. Clone this branch of nvdiffrast.
5. Install nvdiffrast from the local source (e.g., pip install .).
6. Run Python code that triggers the nvdiffrast JIT compilation (e.g., importing nvdiffrast.torch and using a function like rasterize). The compilation should now succeed without errors.

Notes

• This fix currently targets the x86_64-linux target path commonly used by conda on Linux. It might need adaptation for other architectures (e.g., aarch64) or operating systems (Windows, macOS) if they exhibit similar issues with conda packages placing headers in non-standard targets directories.
• Closes # (Link to any relevant issue number if one exists)