From 5a9f3c1e085cf09c6579b97b7439dcfa398f89b5 Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Thu, 11 Jun 2026 19:50:59 +0800 Subject: [PATCH 1/2] python3Packages.tilelang: init at 0.1.11 --- .../python-modules/tilelang/default.nix | 178 ++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 180 insertions(+) create mode 100644 pkgs/development/python-modules/tilelang/default.nix diff --git a/pkgs/development/python-modules/tilelang/default.nix b/pkgs/development/python-modules/tilelang/default.nix new file mode 100644 index 000000000000..5c4e512954bd --- /dev/null +++ b/pkgs/development/python-modules/tilelang/default.nix @@ -0,0 +1,178 @@ +{ + lib, + tilelang, + + cudaSupport ? torch.cudaSupport, + + buildPythonPackage, + cudaPackages, + fetchFromGitHub, + python, + pythonOlder, + stdenv, + + # nativeBuildInputs + autoPatchelfHook, + writableTmpDirAsHomeHook, + + # build-system + cmake, + cython, + ninja, + scikit-build-core, + + # dependencies + apache-tvm-ffi, + cloudpickle, + ml-dtypes, + numpy, + psutil, + torch, + torch-c-dlpack-ext, + tqdm, + typing-extensions, + z3-solver, + + # optional-dependencies, + matplotlib, + + # nativeCheckInputs + einops, + flash-linear-attention, + pytestCheckHook, +}: +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { + pname = "tilelang"; + version = "0.1.11"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "tile-ai"; + repo = "tilelang"; + tag = "v${finalAttrs.version}"; + fetchSubmodules = true; + hash = "sha256-C/c99/26/dBnQJYGrZ+NXl1Rqk3bjM2kpkgP/hWkTGE="; + }; + + postPatch = + # remove binary only packages from build-system.requires (pythonRelaxDeps + # doesn't work) + '' + substituteInPlace pyproject.toml \ + --replace-fail '"z3-solver' '# "z3-solver' \ + --replace-fail '"patchelf' '# "patchelf' + '' + # don't call git + + '' + sed -i '/def get_git_commit_id()/a\\ return None' version_provider.py + ''; + + nativeBuildInputs = [ + writableTmpDirAsHomeHook + ] + ++ lib.optionals stdenv.hostPlatform.isLinux [ + autoPatchelfHook + ] + ++ lib.optionals cudaSupport [ + cudaPackages.cuda_nvcc + ]; + + build-system = [ + cmake + cython + ninja + scikit-build-core + ]; + + env = { + NO_GIT_VERSION = true; + USE_CUDA = cudaSupport; + }; + + dontUseCmakeConfigure = true; + + cmakeFlags = [ + (lib.cmakeFeature "Z3_INCLUDE_DIR" "${z3-solver.dev}/include") + ]; + + buildInputs = lib.optionals cudaSupport [ + cudaPackages.cuda_cudart + cudaPackages.cuda_nvrtc # nvrtc.h + ]; + + dependencies = [ + finalAttrs.passthru.apache-tvm-ffi + cloudpickle + ml-dtypes + numpy + psutil + torch + tqdm + typing-extensions + z3-solver + ] + ++ z3-solver.requiredPythonModules + ++ lib.optionals (pythonOlder "3.14") [ + torch-c-dlpack-ext + ]; + + optional-dependencies = { + vis = [ + matplotlib + ]; + }; + + preFixup = lib.optionalString stdenv.hostPlatform.isLinux '' + # libtvm_ffi.so + addAutoPatchelfSearchPath "${finalAttrs.passthru.apache-tvm-ffi}/${python.sitePackages}/tvm_ffi/lib" + # libz3.so.4.16 + addAutoPatchelfSearchPath "${z3-solver.lib}/lib" + ''; + + pythonImportsCheck = [ "tilelang" ]; + + # requires GPU + doCheck = false; + nativeCheckInputs = [ + einops + flash-linear-attention + pytestCheckHook + ]; + disabledTestPaths = [ + "3rdparty" + # ImportError: cannot import name 'AutoModelForVision2Seq' from 'transformers' + "examples/bitnet-1.58b/vllm_workspace" + # FileNotFoundError: [Errno 2] No such file or directory: './old/bin/python' + "maint/scripts/test_perf_regression.py" + ]; + + passthru.gpuCheck = tilelang.overridePythonAttrs { + requiredSystemFeatures = [ "cuda" ]; + doCheck = true; + }; + + passthru.apache-tvm-ffi = apache-tvm-ffi.overrideAttrs (previousAttrs: { + version = "0.1.10"; + src = previousAttrs.src.override { + tag = null; + rev = "3c35034fd1026011736e19a4e0e1ed0f22058c42"; + hash = "sha256-dqAO6RLLGIRzPk7dNQsQCck+ziyONddhK/t4+S28cn8="; + }; + # fix eval + meta.changelog = ""; + }); + + meta = { + description = "Tile level programming language to generate high performance code"; + homepage = "https://tilelang.com/"; + downloadPage = "https://github.com/tile-ai/tilelang/releases"; + changelog = "https://github.com/tile-ai/tilelang/releases/tag/${finalAttrs.src.tag}"; + license = with lib.licenses; [ + mit + asl20 # 3rdparty/tvm + bsd3 # 3rdparty/cutlass + ]; + maintainers = with lib.maintainers; [ prince213 ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index e82b1fc240a6..00d54ee0cbba 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19731,6 +19731,8 @@ self: super: with self; { tiledb = callPackage ../development/python-modules/tiledb { inherit (pkgs) tiledb; }; + tilelang = callPackage ../development/python-modules/tilelang { }; + tilequant = callPackage ../development/python-modules/tilequant { }; tiler = callPackage ../development/python-modules/tiler { }; From a539b20570beeeec95ff45fd1d8e3a0cba35a8b8 Mon Sep 17 00:00:00 2001 From: Sizhe Zhao Date: Wed, 10 Jun 2026 14:59:01 +0800 Subject: [PATCH 2/2] python3Packages.tilelang: fix build on darwin --- .../python-modules/tilelang/default.nix | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/pkgs/development/python-modules/tilelang/default.nix b/pkgs/development/python-modules/tilelang/default.nix index 5c4e512954bd..b08b91d7bf40 100644 --- a/pkgs/development/python-modules/tilelang/default.nix +++ b/pkgs/development/python-modules/tilelang/default.nix @@ -66,6 +66,15 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { # don't call git + '' sed -i '/def get_git_commit_id()/a\\ return None' version_provider.py + '' + # fix permissions for install_name_tool -id + + lib.optionalString stdenv.hostPlatform.isDarwin '' + substituteInPlace cmake/pypi-z3/FindZ3.cmake \ + --replace-fail COPYONLY 'FILE_PERMISSIONS + OWNER_READ OWNER_WRITE OWNER_EXECUTE + GROUP_READ GROUP_WRITE GROUP_EXECUTE + WORLD_READ WORLD_WRITE WORLD_EXECUTE + COPYONLY' ''; nativeBuildInputs = [ @@ -123,12 +132,18 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { ]; }; - preFixup = lib.optionalString stdenv.hostPlatform.isLinux '' - # libtvm_ffi.so - addAutoPatchelfSearchPath "${finalAttrs.passthru.apache-tvm-ffi}/${python.sitePackages}/tvm_ffi/lib" - # libz3.so.4.16 - addAutoPatchelfSearchPath "${z3-solver.lib}/lib" - ''; + preFixup = + lib.optionalString stdenv.hostPlatform.isLinux '' + # libtvm_ffi.so + addAutoPatchelfSearchPath "${finalAttrs.passthru.apache-tvm-ffi}/${python.sitePackages}/tvm_ffi/lib" + # libz3.so.4.16 + addAutoPatchelfSearchPath "${z3-solver.lib}/lib" + '' + + lib.optionalString stdenv.hostPlatform.isDarwin '' + # @rpath/libz3.dylib + install_name_tool -add_rpath "${z3-solver.lib}/lib" \ + "$out/${python.sitePackages}/tilelang/lib/libtvm_compiler.dylib" + ''; pythonImportsCheck = [ "tilelang" ];