kernel-builder/lib/deps.nix at 6e60da9ea56ac22bf7843f295c6c2af3abc6ddbc · githubnemo/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
{
  lib,
}:

{
  pkgs,
  torch,
  deps,
}:

let
  knownDeps = with pkgs.cudaPackages; {
    "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
      torch.cxxdev
    ];
    "cutlass_sycl" = [ torch.xpuPackages.cutlass-sycl ];
  };
in
let
  depToPkg =
    dep:
    assert lib.assertMsg (builtins.hasAttr dep knownDeps) "Unknown dependency: ${dep}";
    knownDeps.${dep};
in
lib.flatten (map depToPkg deps)