From d0adab60d5d20bf0db25bc026fe5c7790d47f5ef Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 26 Dec 2023 20:04:49 +0000
Subject: [PATCH] nix: explicit jetson support

---
 .devops/nix/jetson-support.nix | 32 ++++++++++++++++++++++++++++++++
 flake.nix                      |  1 +
 2 files changed, 33 insertions(+)
 create mode 100644 .devops/nix/jetson-support.nix

diff --git a/.devops/nix/jetson-support.nix b/.devops/nix/jetson-support.nix
new file mode 100644
index 000000000..08426d2ab
--- /dev/null
+++ b/.devops/nix/jetson-support.nix
@@ -0,0 +1,32 @@
+{ inputs, ... }:
+{
+  perSystem =
+    {
+      config,
+      system,
+      lib,
+      pkgsCuda,
+      ...
+    }:
+    lib.optionalAttrs (system == "aarch64-linux") {
+      packages =
+        let
+          caps.jetson-xavier = "7.2";
+          caps.jetson-orin = "8.7";
+          caps.jetson-nano = "5.3";
+
+          pkgsFor =
+            cap:
+            import inputs.nixpkgs {
+              inherit system;
+              config = {
+                cudaSupport = true;
+                cudaCapabilities = [ cap ];
+                cudaEnableForwardCompat = false;
+                inherit (pkgsCuda.config) allowUnfreePredicate;
+              };
+            };
+        in
+        builtins.mapAttrs (name: cap: ((pkgsFor cap).callPackage ./scope.nix { }).llama-cpp) caps;
+    };
+}
diff --git a/flake.nix b/flake.nix
index d240ececa..b0a6abd3c 100644
--- a/flake.nix
+++ b/flake.nix
@@ -34,6 +34,7 @@
           .devops/nix/nixpkgs-instances.nix
           .devops/nix/apps.nix
           .devops/nix/devshells.nix
+          .devops/nix/jetson-support.nix
         ];
 
         # An overlay can be used to have a more granular control over llama-cpp's