tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 fetchFromGitHub,
4 fetchpatch,
5 rocmUpdateScript,
6 buildPythonPackage,
7 pytestCheckHook,
8 setuptools,
9 distro,
10 pyyaml,
11 msgpack,
12 pandas,
13 joblib,
14 filelock,
15 clr,
16 rich,
17}:
18
19buildPythonPackage rec {
20 pname = "tensile";
21 version = "7.1.1";
22 pyproject = true;
23
24 src = fetchFromGitHub {
25 owner = "ROCm";
26 repo = "Tensile";
27 rev = "rocm-${version}";
28 hash = "sha256-uYLlbtXrc8m8/RpMbP7WaS7VYGIUm2LHYzv00DZAzvE=";
29 };
30
31 # TODO: It should be possible to run asm caps test ONCE for all supported arches
32 # We currently disable the test because it's slow and runs each time tensile launches
33 postPatch = ''
34 substituteInPlace Tensile/Common.py \
35 --replace-fail 'if globalParameters["AssemblerPath"] is not None:' "if False:"
36 # Add an assert that the fallback 9,0,0 is supported before setting the kernel to it
37 # If it's not detected as supported we have an issue with compiler paths or the compiler is broken
38 # and it's better to stop immediately
39 substituteInPlace Tensile/KernelWriter.py \
40 --replace-fail '= (9,0,0)' '= (9,0,0);assert(globalParameters["AsmCaps"][(9,0,0)]["SupportedISA"])'
41 find . -type f -iname "*.sh" -exec chmod +x {} \;
42 patchShebangs Tensile
43 '';
44
45 buildInputs = [ setuptools ];
46
47 propagatedBuildInputs = [
48 pyyaml
49 msgpack
50 pandas
51 joblib
52 distro
53 rich
54 ];
55
56 patches = [
57 ./tensile-solutionstructs-perf-fix.diff
58 ./tensile-create-library-dont-copy-twice.diff
59 (fetchpatch {
60 # [PATCH] Extend Tensile HIP ISA compatibility
61 sha256 = "sha256-d+fVf/vz+sxGqJ96vuxe0jRMgbC5K6j5FQ5SJ1e3Sl8=";
62 url = "https://github.com/GZGavinZhao/Tensile/commit/855cb15839849addb0816a6dde45772034a3e41f.patch";
63 })
64 ];
65
66 doCheck = false; # Too many errors, not sure how to set this up properly
67
68 nativeCheckInputs = [
69 pytestCheckHook
70 filelock
71 clr
72 ];
73
74 env.ROCM_PATH = "${clr}";
75
76 pythonImportsCheck = [ "Tensile" ];
77
78 passthru.updateScript = rocmUpdateScript {
79 name = pname;
80 inherit (src) owner repo;
81 };
82
83 meta = {
84 description = "GEMMs and tensor contractions";
85 homepage = "https://github.com/ROCm/Tensile";
86 license = with lib.licenses; [ mit ];
87 teams = [ lib.teams.rocm ];
88 platforms = lib.platforms.linux;
89 };
90}