From 0d71bb909fd1de7cc43377ea7b0f9fa33e225acf Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 8 Oct 2024 18:02:34 +0200
Subject: [PATCH] python312Packages.tinygrad: patch cuda headers

---
 pkgs/development/python-modules/tinygrad/default.nix | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pkgs/development/python-modules/tinygrad/default.nix b/pkgs/development/python-modules/tinygrad/default.nix
index 8566c0427112..bb0c6065f0a6 100644
--- a/pkgs/development/python-modules/tinygrad/default.nix
+++ b/pkgs/development/python-modules/tinygrad/default.nix
@@ -74,6 +74,13 @@ buildPythonPackage rec {
       substituteInPlace tinygrad/runtime/ops_clang.py \
         --replace-fail "'clang'" "'${lib.getExe clang}'"
     ''
+    # `cuda_fp16.h` and co. are needed at runtime to compile kernels
+    + lib.optionalString cudaSupport ''
+      substituteInPlace tinygrad/runtime/support/compiler_cuda.py \
+        --replace-fail \
+        ', "-I/usr/local/cuda/include", "-I/usr/include", "-I/opt/cuda/include/"' \
+        ', "-I${lib.getDev cudaPackages.cuda_cudart}/include/"'
+    ''
     + lib.optionalString rocmSupport ''
       substituteInPlace tinygrad/runtime/autogen/hip.py \
         --replace-fail "/opt/rocm/lib/libamdhip64.so" "${rocmPackages.clr}/lib/libamdhip64.so" \