Eman Resu
43 lines
1.2 KiB
Nix
43 lines
1.2 KiB
Nix
{
|
|
buildPythonPackage,
|
|
python,
|
|
composable_kernel,
|
|
lib,
|
|
rocm-toolchain,
|
|
setuptools,
|
|
setuptools-scm,
|
|
}:
|
|
buildPythonPackage {
|
|
pyproject = true;
|
|
pname = "ck4inductor";
|
|
build-system = [
|
|
setuptools
|
|
setuptools-scm
|
|
];
|
|
inherit (composable_kernel) version src sourceRoot;
|
|
pythonImportsCheck = [
|
|
"ck4inductor"
|
|
"ck4inductor.universal_gemm.gen_instances"
|
|
"ck4inductor.universal_gemm.gen_instances"
|
|
"ck4inductor.universal_gemm.op"
|
|
];
|
|
propagatedBuildInputs = [
|
|
# At runtime will fail to compile anything with ck4inductor without this
|
|
# can't easily use in checks phase because most of the compiler machinery is in torch
|
|
rocm-toolchain
|
|
];
|
|
checkPhase = ''
|
|
if [ ! -d "$out/${python.sitePackages}/ck4inductor" ]; then
|
|
echo "ck4inductor isn't at the expected location in $out/${python.sitePackages}/ck4inductor"
|
|
exit 1
|
|
fi
|
|
'';
|
|
meta = {
|
|
description = "Pytorch inductor backend which uses composable_kernel universal GEMM implementations";
|
|
homepage = "https://github.com/ROCm/rocm-libraries/tree/develop/projects/composablekernel";
|
|
license = lib.licenses.mit;
|
|
teams = [ lib.teams.rocm ];
|
|
platforms = lib.platforms.linux;
|
|
};
|
|
}
|