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 pyyaml,
10 msgpack,
11 simplejson,
12 ujson,
13 orjson,
14 pandas,
15 joblib,
16 filelock,
17 clr,
18 rich,
19 isTensileLite ? false,
20}:
21
22buildPythonPackage rec {
23 pname = if isTensileLite then "tensilelite" else "tensile";
24 # Using a specific commit which has code object compression support from after the 6.3 release
25 # Without compression packages are too large for hydra
26 version = "6.3-unstable-2024-12-10";
27 format = "pyproject";
28
29 src = fetchFromGitHub {
30 owner = "ROCm";
31 repo = "Tensile";
32 rev = "1752af518190500891a865379a4569b8abf6ba01";
33 hash = "sha256-Wvz4PVs//3Ox7ykZHpjPzOVwlyATyc+MmVVenfTzWK4=";
34 };
35
36 # TODO: It should be possible to run asm caps test ONCE for all supported arches
37 # We currently disable the test because it's slow and runs each time tensile launches
38
39 postPatch =
40 lib.optionalString (!isTensileLite) ''
41 if grep -F .SafeLoader Tensile/LibraryIO.py; then
42 substituteInPlace Tensile/LibraryIO.py \
43 --replace-fail "yaml.SafeLoader" "yaml.CSafeLoader"
44 fi
45 # See TODO above about asm caps test
46 substituteInPlace Tensile/Common.py \
47 --replace-fail 'if globalParameters["AssemblerPath"] is not None:' "if False:"
48 ''
49 + ''
50 # Add an assert that the fallback 9,0,0 is supported before setting the kernel to it
51 # If it's not detected as supported we have an issue with compiler paths or the compiler is broken
52 # and it's better to stop immediately
53 substituteInPlace Tensile/KernelWriter.py \
54 --replace-fail '= (9,0,0)' '= (9,0,0);assert(globalParameters["AsmCaps"][(9,0,0)]["SupportedISA"])'
55 find . -type f -iname "*.sh" -exec chmod +x {} \;
56 patchShebangs Tensile
57 '';
58
59 buildInputs = [ setuptools ];
60
61 propagatedBuildInputs = [
62 pyyaml
63 msgpack
64 pandas
65 joblib
66 ]
67 ++ lib.optionals (!isTensileLite) [
68 rich
69 ]
70 ++ lib.optionals isTensileLite [
71 simplejson
72 ujson
73 orjson
74 ];
75
76 patches =
77 lib.optional (!isTensileLite) ./tensile-solutionstructs-perf-fix.diff
78 ++ lib.optional (!isTensileLite) ./tensile-create-library-dont-copy-twice.diff
79 ++ lib.optional (!isTensileLite) (fetchpatch {
80 # [PATCH] Extend Tensile HIP ISA compatibility
81 sha256 = "sha256-d+fVf/vz+sxGqJ96vuxe0jRMgbC5K6j5FQ5SJ1e3Sl8=";
82 url = "https://github.com/GZGavinZhao/Tensile/commit/855cb15839849addb0816a6dde45772034a3e41f.patch";
83 })
84 ++ lib.optional isTensileLite ./tensilelite-create-library-dont-copy-twice.diff
85 ++ lib.optional isTensileLite ./tensilelite-gen_assembly-venv-err-handling.diff
86 ++ lib.optional isTensileLite ./tensilelite-compression.diff;
87
88 doCheck = false; # Too many errors, not sure how to set this up properly
89
90 nativeCheckInputs = [
91 pytestCheckHook
92 filelock
93 clr
94 ];
95
96 env.ROCM_PATH = "${clr}";
97
98 pythonImportsCheck = [ "Tensile" ];
99
100 passthru.updateScript = rocmUpdateScript {
101 name = pname;
102 inherit (src) owner repo;
103 };
104
105 meta = with lib; {
106 description = "GEMMs and tensor contractions";
107 homepage = "https://github.com/ROCm/Tensile";
108 license = with licenses; [ mit ];
109 teams = [ teams.rocm ];
110 platforms = platforms.linux;
111 };
112}