From 4f6eafbb37cb0434f9cf310d96509f04ba4f634b Mon Sep 17 00:00:00 2001
From: Alex James <git@alextjam.es>
Date: Wed, 27 Nov 2024 11:54:29 -0600
Subject: [PATCH] tiny-cuda-nn: fix a couple of build issues

The CUDA architecture detection code in tiny-cuda-nn's Torch bindings
convert CUDA capabilities to `int`s, which breaks with 9.0a. Filter
9.0a from the CUDA capabilities passed to TCNN_CUDA_ARCHITECTURES. Also
patch the Torch bindings to replace the hardcoded `-std=c++14` argument
to nvcc with `-std=c++17` for CUDA 11+ to fix the build.
---
 pkgs/by-name/ti/tiny-cuda-nn/package.nix | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/pkgs/by-name/ti/tiny-cuda-nn/package.nix b/pkgs/by-name/ti/tiny-cuda-nn/package.nix
index ada84ef50767c..74a76d2b9c90f 100644
--- a/pkgs/by-name/ti/tiny-cuda-nn/package.nix
+++ b/pkgs/by-name/ti/tiny-cuda-nn/package.nix
@@ -11,7 +11,7 @@
   which,
 }: let
   inherit (lib) lists strings;
-  inherit (cudaPackages) backendStdenv flags;
+  inherit (cudaPackages) backendStdenv cudaVersion flags;
 
   cuda-common-redist = with cudaPackages; [
     (lib.getDev cuda_cudart) # cuda_runtime.h
@@ -36,6 +36,14 @@
     name = "cuda-redist";
     paths = cuda-common-redist;
   };
+
+  unsupportedCudaCapabilities = [
+    "9.0a"
+  ];
+
+  cudaCapabilities = lists.subtractLists unsupportedCudaCapabilities flags.cudaCapabilities;
+
+  cudaArchitecturesString = strings.concatMapStringsSep ";" flags.dropDot cudaCapabilities;
 in
   stdenv.mkDerivation (finalAttrs: {
     pname = "tiny-cuda-nn";
@@ -52,6 +60,13 @@ in
       hash = "sha256-qW6Fk2GB71fvZSsfu+mykabSxEKvaikZ/pQQZUycOy0=";
     };
 
+    # Remove this once a release is made with
+    # https://github.com/NVlabs/tiny-cuda-nn/commit/78a14fe8c292a69f54e6d0d47a09f52b777127e1
+    postPatch = lib.optionals (strings.versionAtLeast cudaVersion "11.0") ''
+      substituteInPlace bindings/torch/setup.py --replace-fail \
+        "-std=c++14" "-std=c++17"
+    '';
+
     nativeBuildInputs =
       [
         cmake
@@ -89,7 +104,7 @@ in
     doCheck = false;
 
     preConfigure = ''
-      export TCNN_CUDA_ARCHITECTURES="${flags.cmakeCudaArchitecturesString}"
+      export TCNN_CUDA_ARCHITECTURES="${cudaArchitecturesString}"
       export CUDA_HOME="${cuda-native-redist}"
       export LIBRARY_PATH="${cuda-native-redist}/lib/stubs:$LIBRARY_PATH"
       export CC="${backendStdenv.cc}/bin/cc"