ir/flake.nix at main · dev-nicelee/ir · 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
143
144
145
146
147
148
149
150
151
{
  description = "ir - local markdown semantic search with hybrid BM25+vector retrieval and LLM reranking";

  inputs = {
    flake-parts.url = "github:hercules-ci/flake-parts";
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    rust-overlay = {
      url = "github:oxalica/rust-overlay";
      inputs.nixpkgs.follows = "nixpkgs";
    };
  };

  outputs =
    inputs@{ flake-parts, ... }:
    flake-parts.lib.mkFlake { inherit inputs; } {
      systems = [
        "x86_64-linux"
        "aarch64-linux"
        "aarch64-darwin"
        "x86_64-darwin"
      ];
      perSystem =
        {
          pkgs,
          lib,
          system,
          config,
          ...
        }:
        let
          isDarwin = pkgs.stdenv.isDarwin;

          # GPU backend selection for Linux.
          # Default: CPU + OpenMP only (safe, no external GPU deps).
          # Override: nix build .#cuda / .#rocm / .#vulkan
          linuxGpuVariants = {
            cpu = {
              features = "llama-openmp";
              extraBuildInputs = [ pkgs.llvmPackages.openmp ];
              ldflags = "-lomp";
            };
            cuda = {
              # Requires nixpkgs config.allowUnfree = true and CUDA drivers on host.
              features = "llama-openmp,llama-cuda";
              extraBuildInputs = [
                pkgs.llvmPackages.openmp
                pkgs.cudaPackages.cudatoolkit
                pkgs.cudaPackages.cuda_cudart
                pkgs.cudaPackages.libcublas
              ];
              ldflags = "-lomp";
            };
            rocm = {
              # Requires ROCm stack on host (/opt/rocm or ROCM_PATH set).
              features = "llama-openmp,llama-rocm";
              extraBuildInputs = [
                pkgs.llvmPackages.openmp
                pkgs.rocmPackages.clr
                pkgs.rocmPackages.rocblas
                pkgs.rocmPackages.hipblas
              ];
              ldflags = "-lomp";
            };
            vulkan = {
              features = "llama-openmp,llama-vulkan";
              extraBuildInputs = [
                pkgs.llvmPackages.openmp
                pkgs.vulkan-loader
                pkgs.vulkan-headers
              ];
              ldflags = "-lomp -lvulkan";
            };
          };

          mkPackage =
            { variant ? null }:
            let
              linux = linuxGpuVariants.${if variant != null then variant else "cpu"};
              cargoFlags = lib.optionals (!isDarwin) [
                "--no-default-features"
                "--features"
                linux.features
              ];
            in
            pkgs.rustPlatform.buildRustPackage {
              pname = "ir";
              version = (builtins.fromTOML (builtins.readFile ./Cargo.toml)).package.version;
              src = ./.;
              cargoLock.lockFile = ./Cargo.lock;

              nativeBuildInputs = [
                pkgs.cmake
                pkgs.python3
              ];

              buildInputs =
                lib.optionals isDarwin [
                  pkgs.darwin.apple_sdk.frameworks.Accelerate
                  pkgs.darwin.apple_sdk.frameworks.Foundation
                  pkgs.darwin.apple_sdk.frameworks.Metal
                ]
                ++ lib.optionals (!isDarwin) linux.extraBuildInputs;

              cargoBuildFlags = cargoFlags;
              cargoTestFlags = cargoFlags;

              # ggml is statically linked; cargo's final link step needs explicit
              # GPU library flags that cmake's detection only set in the static archive.
              env.NIX_LDFLAGS = lib.optionalString (!isDarwin) linux.ldflags;

              CMAKE_GENERATOR = "Unix Makefiles";

              meta = {
                description = (builtins.fromTOML (builtins.readFile ./Cargo.toml)).package.description;
                homepage = "https://github.com/vlwkaos/ir";
                license = lib.licenses.mit;
                mainProgram = "ir";
                platforms = lib.platforms.unix;
              };
            };
        in
        {
          _module.args.pkgs = import inputs.nixpkgs {
            inherit system;
            overlays = [ (import inputs.rust-overlay) ];
            config = { };
          };

          packages = {
            default = mkPackage { };
            cuda = mkPackage { variant = "cuda"; };
            rocm = mkPackage { variant = "rocm"; };
            vulkan = mkPackage { variant = "vulkan"; };
          };

          devShells.default = pkgs.mkShell {
            inputsFrom = [ config.packages.default ];
            nativeBuildInputs = [
              (pkgs.rust-bin.stable."1.95.0".default.override {
                extensions = [
                  "rust-src"
                  "rust-analyzer"
                ];
              })
            ];
          };

          formatter = pkgs.nixfmt-rfc-style;
        };
    };
}