Something went wrong. Try again.
lol
Something went wrong. Try again.
4.7 kB · 180 lines
Nix
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181{ lib, amd-aiter, buildPythonPackage, einops, fetchFromGitHub, ninja, numpy, packaging, pandas, psutil, pybind11, python, rocmPackages, runCommand, setuptools, setuptools-scm, torch, writableTmpDirAsHomeHook,}:let # Provide a default set of include paths needed by aiter at runtime for JIT modules defaultRocmIncl = lib.makeIncludePath ( (with rocmPackages; [ clr hipblas hipblas-common hipblaslt hipcub hipfft hipsolver hipsparse rocblas rocprim rocsolver rocsparse rocthrust ]) ++ [ pybind11 ] );inbuildPythonPackage (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
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}"'
# NIX_AITER_ROCM_INCL is a colon-separated list of include dirs # defaults to packages known to be needed by aiter substituteInPlace aiter/jit/utils/cpp_extension.py \ --replace-fail \ 'paths.append(_join_rocm_home("include"))' \ 'paths.append(_join_rocm_home("include")); paths.extend(os.environ.get("NIX_AITER_ROCM_INCL", "${defaultRocmIncl}").split(":"))'
# 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 = { BUILD_TARGET = "rocm"; PREBUILD_KERNELS = "0"; ROCM_PATH = "${rocmPackages.clr}"; SETUPTOOLS_SCM_PRETEND_VERSION = finalAttrs.version; };
build-system = [ ninja packaging pandas psutil pybind11 setuptools setuptools-scm ];
buildInputs = [ rocmPackages.clr ];
nativeBuildInputs = [ rocmPackages.hipcc writableTmpDirAsHomeHook ];
dependencies = [ einops ninja numpy packaging pandas psutil pybind11 torch ];
# Most tests and imports require a GPU and writable $HOME for JIT cache doCheck = false;
# Test JIT module builds for CDNA3 iff rocm enabled for torch passthru.tests = lib.optionalAttrs torch.rocmSupport ( let mkJitTest = name: moduleName: runCommand "amd-aiter-jit-${name}" { nativeBuildInputs = [ (python.withPackages (_: [ amd-aiter ])) rocmPackages.clr writableTmpDirAsHomeHook ]; env = { CXX = "amdclang++"; GPU_ARCHS = "gfx942"; PYTORCH_ROCM_ARCH = "gfx942"; }; } '' export AITER_JIT_DIR=$(mktemp -d) python3 -c " from aiter.jit.core import get_args_of_build, build_module args = get_args_of_build('${moduleName}') build_module( '${moduleName}', args['srcs'], args['flags_extra_cc'], args['flags_extra_hip'], args['blob_gen_cmd'], args['extra_include'], args['extra_ldflags'], args['verbose'], args['is_python_module'], args['is_standalone'], args['torch_exclude'], ) print('JIT compile of ${moduleName} finished') " touch $out ''; in { jit-module-opus-sort = mkJitTest "opus-sort" "module_moe_sorting_opus"; jit-module-mhc = mkJitTest "mhc" "module_mhc"; } );
meta = { description = "AI Tensor Engine for ROCm"; homepage = "https://github.com/ROCm/aiter"; license = lib.licenses.mit; sourceProvenance = with lib.sourceTypes; [ fromSource binaryNativeCode ]; maintainers = with lib.maintainers; [ lach ]; teams = [ lib.teams.rocm ]; platforms = lib.platforms.linux; };})