cu-embed 0.1.0

Compile CUDA kernels with nvcc, embed cubin/PTX artifacts, and load the best module at runtime.
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

{
  description = "Rust + CUDA kernels embedded with rust-embed";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    crane.url = "github:ipetkov/crane";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, crane, flake-utils }:
    flake-utils.lib.eachSystem [ "x86_64-linux" ] (system:
      let
        pkgs = import nixpkgs {
          inherit system;
          config.allowUnfree = true;
        };
        lib = pkgs.lib;
        craneLib = crane.mkLib pkgs;

        cudaToolkit = pkgs.buildEnv {
          name = "cuda-toolkit-with-nvrtc-static";
          paths = [
            pkgs.cudaPackages.cudatoolkit
            pkgs.cudaPackages.cuda_nvrtc.static
          ];
          pathsToLink = [
            "/bin"
            "/include"
            "/lib"
            "/nix-support"
          ];
          ignoreCollisions = true;
        };
        cudaDriverLibPath = "/run/opengl-driver/lib";
        nvccCcbin = lib.getExe' pkgs.cudaPackages.backendStdenv.cc.cc "g++";
        gccStaticLibDir = "${pkgs.cudaPackages.backendStdenv.cc.cc}/lib";

        cuFilter = path: type:
          (builtins.baseNameOf path == "README.md") ||
          (builtins.baseNameOf path == "LICENSE-MIT") ||
          (builtins.baseNameOf path == "LICENSE-APACHE") ||
          (lib.hasSuffix ".cu" path) ||
          (craneLib.filterCargoSources path type);

        src = lib.cleanSourceWith {
          src = ./.;
          filter = cuFilter;
        };

        cudaDeps = [ cudaToolkit ];

        commonArgs = {
          inherit src;
          strictDeps = true;
          nativeBuildInputs = cudaDeps;
          NVCC = "${cudaToolkit}/bin/nvcc";
          NVCC_CCBIN = nvccCcbin;
          CUDA_PATH = cudaToolkit;
          RUSTFLAGS = "-L native=${gccStaticLibDir}";
        };

        cargoArtifacts = craneLib.buildDepsOnly commonArgs;
      in
      {
        packages.default = craneLib.buildPackage (commonArgs // {
          inherit cargoArtifacts;
        });

        packages.example = craneLib.buildPackage (commonArgs // {
          inherit cargoArtifacts;
          cargoExtraArgs = "--example add_scalar";
          nativeBuildInputs = commonArgs.nativeBuildInputs ++ [ pkgs.makeWrapper ];
          postFixup = ''
            wrapProgram $out/bin/add_scalar \
              --prefix LD_LIBRARY_PATH : ${cudaDriverLibPath}
          '';
        });

        devShells.default = pkgs.mkShell {
          packages = [
            pkgs.cargo
            pkgs.rustc
            pkgs.rustfmt
          ] ++ cudaDeps;

          NVCC = "${cudaToolkit}/bin/nvcc";
          NVCC_CCBIN = nvccCcbin;
          CUDA_PATH = cudaToolkit;
          RUSTFLAGS = "-L native=${gccStaticLibDir}";

          shellHook = ''
            export LD_LIBRARY_PATH=${cudaDriverLibPath}''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
          '';
        };
      });
}