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