tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 buildPythonPackage,
3 python,
4 composable_kernel,
5 lib,
6 setuptools,
7 setuptools-scm,
8 rocm-merged-llvm,
9}:
10buildPythonPackage {
11 pyproject = true;
12 pname = "ck4inductor";
13 build-system = [
14 setuptools
15 setuptools-scm
16 ];
17 version = "6.4.0";
18 inherit (composable_kernel) src;
19 pythonImportsCheck = [
20 "ck4inductor"
21 "ck4inductor.universal_gemm.gen_instances"
22 "ck4inductor.universal_gemm.gen_instances"
23 "ck4inductor.universal_gemm.op"
24 ];
25 propagatedBuildInputs = [
26 # At runtime will fail to compile anything with ck4inductor without this
27 # can't easily use in checks phase because most of the compiler machinery is in torch
28 rocm-merged-llvm
29 ];
30 checkPhase = ''
31 if [ ! -d "$out/${python.sitePackages}/ck4inductor" ]; then
32 echo "ck4inductor isn't at the expected location in $out/${python.sitePackages}/ck4inductor"
33 exit 1
34 fi
35 '';
36 meta = with lib; {
37 description = "Pytorch inductor backend which uses composable_kernel universal GEMM implementations";
38 homepage = "https://github.com/ROCm/composable_kernel";
39 license = with licenses; [ mit ];
40 teams = [ teams.rocm ];
41 platforms = platforms.linux;
42 };
43}