kernels/builder/lib/deps.nix at c81046058f00220ee146e42195bc13ba1262ce30 · robtaylor/kernels · 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
{
  lib,
  pkgs,
  torch,
}:

let
  cppDeps = {
    "cutlass_2_10" = [
      pkgs.cutlass_2_10
    ];
    "cutlass_3_5" = [
      pkgs.cutlass_3_5
    ];
    "cutlass_3_6" = [
      pkgs.cutlass_3_6
    ];
    "cutlass_3_8" = [
      pkgs.cutlass_3_8
    ];
    "cutlass_3_9" = [
      pkgs.cutlass_3_9
    ];
    "cutlass_4_0" = [
      pkgs.cutlass_4_0
    ];
    "torch" = [
      torch
    ];
    "sycl_tla" = [ torch.xpuPackages.sycl-tla ];
    "metal-cpp" = [
      pkgs.metal-cpp.dev
    ];
  };

  pythonDeps =
    let
      depsJson = builtins.fromJSON (builtins.readFile ../../build2cmake/src/python_dependencies.json);
      # Map the Nix package names to actual Nix packages.
      updatePackage = _name: dep: dep // { nix = map (pkg: pkgs.python3.pkgs.${pkg}) dep.nix; };
      updateBackend = _backend: backendDeps: lib.mapAttrs updatePackage backendDeps;
    in
    depsJson
    // {
      general = lib.mapAttrs updatePackage depsJson.general;
      backends = lib.mapAttrs updateBackend depsJson.backends;
    };

  getCppDep = dep: cppDeps.${dep} or (throw "Unknown dependency: ${dep}");
  getPythonDep =
    dep: lib.attrByPath [ "general" dep "nix" ] (throw "Unknown Python dependency: ${dep}") pythonDeps;
  getBackendPythonDep =
    backend: dep:
    let
      backendDeps = lib.attrByPath [
        "backends"
        backend
      ] (throw "Unknown backend: ${backend}") pythonDeps;
    in
    lib.attrByPath [
      dep
      "nix"
    ] (throw "Unknown Python dependency for backend `${backend}`: ${dep}") backendDeps;
in
{
  resolveCppDeps = deps: lib.flatten (map getCppDep deps);
  resolvePythonDeps = deps: lib.flatten (map getPythonDep deps);
  resolveBackendPythonDeps = backend: deps: lib.flatten (map (getBackendPythonDep backend) deps);
}