kernel-builder/lib/build-sets.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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
{
  nixpkgs,
  system,
  torchVersions,
}:

let
  inherit (nixpkgs) lib;

  overlay = import ../overlay.nix;

  inherit (import ./torch-version-utils.nix { inherit lib; })
    backend
    flattenSystems
    ;

  # All build configurations supported by Torch.
  buildConfigs =
    system:
    let
      filterMap = f: xs: builtins.filter (x: x != null) (builtins.map f xs);
      systemBuildConfigs = filterMap (version: if version.system == system then version else null) (
        flattenSystems torchVersions
      );
    in
    builtins.map (buildConfig: buildConfig // { backend = backend buildConfig; }) systemBuildConfigs;

  flattenVersion = version: lib.replaceStrings [ "." ] [ "_" ] (lib.versions.pad 2 version);

  overlayForTorchVersion = torchVersion: sourceBuild: self: super: {
    pythonPackagesExtensions = super.pythonPackagesExtensions ++ [
      (
        python-self: python-super: with python-self; {
          torch =
            if sourceBuild then
              python-self."torch_${flattenVersion torchVersion}"
            else
              python-self."torch-bin_${flattenVersion torchVersion}";
        }
      )
    ];
  };

  # An overlay that overides CUDA to the given version.
  overlayForCudaVersion = cudaVersion: self: super: {
    cudaPackages = super."cudaPackages_${flattenVersion cudaVersion}";
  };

  overlayForRocmVersion = rocmVersion: self: super: {
    rocmPackages = super."rocmPackages_${flattenVersion rocmVersion}";
  };

  overlayForXpuVersion = xpuVersion: self: super: {
    xpuPackages = super."xpuPackages_${flattenVersion xpuVersion}";
  };

  backendConfig = {
    cpu = {
      allowUnfree = true;
    };

    cuda = {
      allowUnfree = true;
      cudaSupport = true;
    };

    metal = {
      allowUnfree = true;
      metalSupport = true;
    };

    rocm = {
      allowUnfree = true;
      rocmSupport = true;
    };

    xpu = {
      allowUnfree = true;
      xpuSupport = true;
    };
  };

  xpuConfig = {
    allowUnfree = true;
    xpuSupport = true;
  };

  # Construct the nixpkgs package set for the given versions.
  mkBuildSet =
    buildConfig@{
      backend,
      cpu ? false,
      cudaVersion ? null,
      metal ? false,
      rocmVersion ? null,
      xpuVersion ? null,
      torchVersion,
      system,
      bundleBuild ? false,
      sourceBuild ? false,
    }:
    let
      backendOverlay =
        if buildConfig.backend == "cpu" then
          [ ]
        else if buildConfig.backend == "cuda" then
          [ (overlayForCudaVersion buildConfig.cudaVersion) ]
        else if buildConfig.backend == "rocm" then
          [ (overlayForRocmVersion buildConfig.rocmVersion) ]
        else if buildConfig.backend == "metal" then
          [ ]
        else if buildConfig.backend == "xpu" then
          [ (overlayForXpuVersion buildConfig.xpuVersion) ]
        else
          throw "No compute framework set in Torch version";
      config =
        backendConfig.${buildConfig.backend} or (throw "No backend config for ${buildConfig.backend}");

      pkgs = import nixpkgs {
        inherit config system;
        overlays = [
          overlay
        ]
        ++ backendOverlay
        ++ [ (overlayForTorchVersion torchVersion sourceBuild) ];
      };

      torch = pkgs.python3.pkgs.torch;

      extension = pkgs.callPackage ./torch-extension { inherit torch; };
    in
    {
      inherit
        buildConfig
        extension
        pkgs
        torch
        bundleBuild
        ;
    };
in
map mkBuildSet (buildConfigs system)