kernel-builder/lib/build-variants.nix at 1fdc859a69b2b5760dfc897e6b1ef952e1ee75d2 · huggingface/kernel-builder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
{ lib }:
let
  inherit (import ./torch-version-utils.nix { inherit lib; })
    backend
    flattenSystems
    ;
in
rec {

  buildName =
    let
      inherit (import ./version-utils.nix { inherit lib; }) flattenVersion;
      computeString =
        version:
        if backend version == "cpu" then
          "cpu"
        else if backend version == "cuda" then
          "cu${flattenVersion (lib.versions.majorMinor version.cudaVersion)}"
        else if backend version == "rocm" then
          "rocm${flattenVersion (lib.versions.majorMinor version.rocmVersion)}"
        else if backend version == "metal" then
          "metal"
        else if backend version == "xpu" then
          "xpu${flattenVersion (lib.versions.majorMinor version.xpuVersion)}"
        else
          throw "No compute framework set in Torch version";
    in
    version:
    if version.system == "aarch64-darwin" then
      "torch${flattenVersion version.torchVersion}-${computeString version}-${version.system}"
    else
      "torch${flattenVersion version.torchVersion}-cxx11-${computeString version}-${version.system}";

  # Build variants included in bundle builds.
  buildVariants =
    torchVersions:
    let
      bundleBuildVersions = lib.filter (version: version.bundleBuild or false);
    in
    lib.foldl' (
      acc: version:
      let
        path = [
          version.system
          (backend version)
        ];
        pathVersions = lib.attrByPath path [ ] acc ++ [ (buildName version) ];
      in
      lib.recursiveUpdate acc (lib.setAttrByPath path pathVersions)
    ) { } (flattenSystems (bundleBuildVersions torchVersions));
}