Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198{ lib, addDriverRunpath, config, buildPythonPackage, fetchFromGitHub,
# build-system setuptools,
# dependencies typing-extensions,
# nativeBuildInputs cmake, cudaPackages,
# buildInputs dlpack,
# tests cupy, llvm, numba-cuda, numpy, pytest-benchmark, pytest-env, pytestCheckHook, torch, writableTmpDirAsHomeHook,
# passthru cuda-tile,}:let libCudaDir = # `cuda_compat` provides a newer `libcuda.so` in its `compat` directory if (cudaPackages.cuda_compat.meta.available or false) then "${cudaPackages.cuda_compat}/compat" # Else: the host driver library else "${addDriverRunpath.driverLink}/lib";
# https://github.com/NVIDIA/cutile-python/blob/v1.7.0/cmake/FetchXLAHeaders.cmake#L5-L6 xla = fetchFromGitHub { owner = "openxla"; repo = "xla"; rev = "b6f37ab7767f428fd6f993de5e211643d47d4deb"; hash = "sha256-U4e3k4nm9gB1x5hahXwycWSryBQuxIPmOzVf6kuahY0="; };inbuildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: { pname = "cuda-tile"; version = "1.7.0"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "NVIDIA"; repo = "cutile-python"; tag = "v${finalAttrs.version}"; hash = "sha256-mO5VOusQg6wYzrvkD9SB74O6g5fOnyEdCIIKBkjC7s4="; };
patches = [ ./xla_ffi-include-python-first.patch ];
postPatch = '' substituteInPlace pyproject.toml \ --replace-fail "setuptools==84.0.0" "setuptools" '' # Otherwise fails to find libc # xla_ffi.cpp:(.text+0x308): undefined reference to `__stack_chk_fail' + '' substituteInPlace cext/CMakeLists.txt \ --replace-fail \ "target_link_libraries(_cext_shared _cext_static \''${Python_LIBRARIES} \''${asan_library})" \ "target_link_libraries(_cext_shared _cext_static \''${Python_LIBRARIES} \''${asan_library} c)" '' # Don't build with -Werror: GCC 15 flags variables in cext/tile_kernel.cpp as set but not used + '' substituteInPlace CMakeLists.txt \ --replace-fail \ 'treat warnings as errors" OFF)' \ 'treat warnings as errors" ON)' '' # Get rid of the vendored broken logic for finding the CUDA toolkit + '' rm cmake/FindCUDAToolkit.cmake '' # Manually inject the library version + '' echo "${finalAttrs.version}" >src/cuda/tile/VERSION '' # `libcuda.so.1` is loaded by name, which the dynamic loader doesn't find outside of FHS dirs + '' substituteInPlace src/cuda/tile/_load_libcuda.py \ --replace-fail \ '_dll = CDLL("libcuda.so.1")' \ '_dll = CDLL("${libCudaDir}/libcuda.so.1")' '' # Use tileiras from Nixpkgs instead of searching for it in /usr/local/cuda*, which is the last place cuTile looks # (after the nvidia-cuda-tileiras Python package, PATH, and $CUDA_HOME/bin). + lib.optionalString cudaPackages.cuda_tileiras.meta.available '' substituteInPlace src/cuda/tile/_compile.py \ --replace-fail \ "res = _find_file(_get_default_cuda_toolkit_paths(), [\"bin\"], [binary_name]," \ "res = _find_file([\"${lib.getBin cudaPackages.cuda_tileiras}\"], [\"bin\"], [binary_name]," '';
build-system = [ setuptools ];
env = { CUDA_TILE_CMAKE_DLPACK_PATH = dlpack; CUDA_TILE_CMAKE_XLA_PATH = xla; };
nativeBuildInputs = [ cmake cudaPackages.cuda_nvcc ]; dontUseCmakeConfigure = true;
buildInputs = [ cudaPackages.cuda_cudart # cuda.h cudaPackages.libnvvm # nvvm.h (libnvvm itself is loaded at run time, by the experimental cuda.lang only) ];
dependencies = [ typing-extensions ];
pythonImportsCheck = [ "cuda.tile" ];
nativeCheckInputs = [ cupy llvm # FileCheck, for the MLIR tests numba-cuda numpy pytest-benchmark pytest-env pytestCheckHook torch # cuTile fails to compile kernels if it cannot create its cache directory in $HOME writableTmpDirAsHomeHook ];
# Run the benchmarks once as tests, without timing them pytestFlags = [ "--benchmark-disable" ];
# The benchmarks look up their tuned configurations by GPU name, which they otherwise get from nvidia-smi preCheck = '' export CUDA_TILE_BENCHMARK_GPU_NAME="$(python -c 'import torch; print(torch.cuda.get_device_name(0))')" '';
disabledTestPaths = [ # Separate packages (task_scheduling, tile_preview, cuda-lang), not installed here; upstream's CI ignores them too "experimental" ];
disabledTests = [ # Deselected by upstream's CI too "int64_index" # These need more memory than a 24 GiB GPU has: a 65536x16384 layer norm, and float32 vector adds of 2^30 elements # (in each layout). Failed tests keep their tensors alive (through their tracebacks), so they also made later # benchmarks run out of memory. "64K-16K" "1G-1G-f32" "1G-1Mx1K-f32" "1G-8Mx128-f32" "1G-32Mx32-f32" ];
# Tests require access to a physical GPU doCheck = false;
passthru.gpuCheck = cuda-tile.overridePythonAttrs { requiredSystemFeatures = [ "cuda" ]; doCheck = true; };
meta = { description = "Programming model for writing parallel kernels for NVIDIA GPUs"; homepage = "https://docs.nvidia.com/cuda/cutile-python/"; downloadPage = "https://github.com/NVIDIA/cutile-python"; changelog = "https://docs.nvidia.com/cuda/cutile-python/generated/release_notes.html"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ GaetanLepage prince213 ]; broken = !config.cudaSupport; };})