Commit 0d71bb90 authored by Gaetan Lepage's avatar Gaetan Lepage Committed by Gaétan Lepage
Browse files

python312Packages.tinygrad: patch cuda headers

parent ba70efee

pkgs/development/python-modules/tinygrad/default.nix

+7 −0
Original line number Diff line number Diff line
@@ -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" \