diff --git a/pkgs/development/python-modules/cuda-tile/default.nix b/pkgs/development/python-modules/cuda-tile/default.nix index 7e4baa83da35..c63021b61bcf 100644 --- a/pkgs/development/python-modules/cuda-tile/default.nix +++ b/pkgs/development/python-modules/cuda-tile/default.nix @@ -20,6 +20,7 @@ # tests cupy, + llvm, numba-cuda, numpy, pytest-benchmark, @@ -40,7 +41,7 @@ let else "${addDriverRunpath.driverLink}/lib"; - # https://github.com/NVIDIA/cutile-python/blob/v1.4.0/cmake/FetchXLAHeaders.cmake#L5-L6 + # https://github.com/NVIDIA/cutile-python/blob/v1.7.0/cmake/FetchXLAHeaders.cmake#L5-L6 xla = fetchFromGitHub { owner = "openxla"; repo = "xla"; @@ -50,7 +51,7 @@ let in buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: { pname = "cuda-tile"; - version = "1.4.0"; + version = "1.7.0"; pyproject = true; __structuredAttrs = true; @@ -58,7 +59,7 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs owner = "NVIDIA"; repo = "cutile-python"; tag = "v${finalAttrs.version}"; - hash = "sha256-R5V69nJLQ3/1995ezH1/WuueA6cm1vhKZdOECqbwPbU="; + hash = "sha256-mO5VOusQg6wYzrvkD9SB74O6g5fOnyEdCIIKBkjC7s4="; }; patches = [ @@ -67,9 +68,7 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs postPatch = '' substituteInPlace pyproject.toml \ - --replace-fail "setuptools==80.10.2" "setuptools" - substituteInPlace setup.py \ - --replace-fail ' dry_run=self.dry_run)' ' )' + --replace-fail "setuptools==84.0.0" "setuptools" '' # Otherwise fails to find libc # xla_ffi.cpp:(.text+0x308): undefined reference to `__stack_chk_fail' @@ -79,6 +78,13 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs "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 @@ -99,8 +105,8 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs + lib.optionalString cudaPackages.cuda_tileiras.meta.available '' substituteInPlace src/cuda/tile/_compile.py \ --replace-fail \ - "for toolkit_path in _get_default_cuda_toolkit_paths():" \ - "for toolkit_path in [\"${lib.getBin cudaPackages.cuda_tileiras}\"]:" + "res = _find_file(_get_default_cuda_toolkit_paths(), [\"bin\"], [binary_name]," \ + "res = _find_file([\"${lib.getBin cudaPackages.cuda_tileiras}\"], [\"bin\"], [binary_name]," ''; build-system = [ @@ -120,6 +126,7 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs buildInputs = [ cudaPackages.cuda_cudart # cuda.h + cudaPackages.libnvvm # nvvm.h (libnvvm itself is loaded at run time, by the experimental cuda.lang only) ]; dependencies = [ @@ -130,6 +137,7 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs nativeCheckInputs = [ cupy + llvm # FileCheck, for the MLIR tests numba-cuda numpy pytest-benchmark @@ -143,6 +151,29 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs # 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;