Last active
March 19, 2026 15:00
-
-
Save LunNova/f30810259d4977554c60586321fca2dd to your computer and use it in GitHub Desktop.
CDNA4 (MI355x) tested only, initial aiter packaging work.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| lib, | |
| buildPythonPackage, | |
| fetchFromGitHub, | |
| setuptools, | |
| setuptools-scm, | |
| packaging, | |
| psutil, | |
| pybind11, | |
| einops, | |
| ninja, | |
| numpy, | |
| pandas, | |
| rocmPackages, | |
| torch, | |
| writableTmpDirAsHomeHook, | |
| }: | |
| # FIXME FIXME will fail at runtime without env vars like this | |
| # aiter JIT writes compiled .so here; also added to sys.path so | |
| # importlib.import_module(bare_name) finds them. if unset aiter | |
| # uses home dir *but does not add to import path* and fails | |
| # with confusing double import error | |
| # AITER_JIT_DIR = "/tmp/aiter-jit"; | |
| # hipcc (clang++) needs these for JIT kernel compilation | |
| #CPLUS_INCLUDE_PATH = builtins.concatStringsSep ":" (map | |
| # (p: "${lib.getInclude p}/include") | |
| # ((with pkgs.rocmPackages; [ | |
| # rocthrust rocprim clr | |
| # hipsparse hipblas hipblas-common hipblaslt hipsolver rocsparse rocblas rocsolver hipfft | |
| # ]) ++ [ | |
| # pkgs.python3.pkgs.pybind11 | |
| # ]) | |
| #); | |
| # FIXME: needs hipcc at runtime, inexplicably, despite that being deprecated before aiter existing. | |
| buildPythonPackage (finalAttrs: { | |
| pname = "amd-aiter"; | |
| version = "0.1.11.post1"; | |
| pyproject = true; | |
| src = fetchFromGitHub { | |
| owner = "ROCm"; | |
| repo = "aiter"; | |
| tag = "v${finalAttrs.version}"; | |
| hash = "sha256-9CCSmEw0kIoxERhtkKhBkaAGx42kCssH7IPjTgbg9LA="; | |
| }; | |
| postPatch = '' | |
| rmdir 3rdparty/composable_kernel | |
| ln -sf ${rocmPackages.composable_kernel.src} 3rdparty/composable_kernel | |
| # python relax deps hook? | |
| substituteInPlace pyproject.toml \ | |
| --replace-fail '"flydsl==0.0.1.dev95158637"' "" | |
| # TODO: upstream fix. get_git_commit_id_short() assumes a git clone; | |
| # fails in hermetic builds, should fall back gracefully. | |
| substituteInPlace csrc/cpp_itfs/utils.py \ | |
| --replace-fail \ | |
| 'commit_id = get_git_commit_id_short()' \ | |
| 'commit_id = "${finalAttrs.version}"' | |
| # setuptools runs setup.py twice (metadata + wheel). prepare_packaging() | |
| # copies 3rdparty/ (with nix store read-only files) into aiter_meta/, then | |
| # the second run can't rmtree or overwrite them. | |
| substituteInPlace setup.py \ | |
| --replace-fail \ | |
| $'prepare_packaging()\n\n\nclass' \ | |
| $'if not os.path.exists("aiter_meta"): prepare_packaging()\n\n\nclass' \ | |
| --replace-fail \ | |
| 'if os.path.exists("aiter_meta") and os.path.isdir("aiter_meta"):' \ | |
| 'if False:' | |
| ''; | |
| env = { | |
| SETUPTOOLS_SCM_PRETEND_VERSION = finalAttrs.version; | |
| PREBUILD_KERNELS = "0"; | |
| BUILD_TARGET = "rocm"; | |
| ROCM_PATH = "${rocmPackages.clr}"; | |
| GPU_ARCHS = "gfx950"; # FIXME clr.localGpuTargets or clr.gpuTargets ? or does it even matter when PREBUILD_KERNELS is off | |
| }; | |
| build-system = [ | |
| setuptools | |
| setuptools-scm | |
| packaging | |
| psutil | |
| pybind11 | |
| ninja | |
| pandas | |
| ]; | |
| buildInputs = [ rocmPackages.clr ]; | |
| nativeBuildInputs = [ | |
| rocmPackages.hipcc | |
| writableTmpDirAsHomeHook | |
| ]; | |
| pythonImportsCheck = [ "aiter" ]; | |
| dependencies = [ | |
| einops | |
| ninja | |
| numpy | |
| packaging | |
| pandas | |
| psutil | |
| pybind11 | |
| torch | |
| ]; | |
| doCheck = false; | |
| meta = { | |
| description = "AI Tensor Engine for ROCm"; | |
| homepage = "https://github.com/ROCm/aiter"; | |
| license = lib.licenses.mit; | |
| platforms = lib.platforms.linux; | |
| }; | |
| }) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment