Add aarch64

hasktorch · May 28, 2023 · b534310 · b534310
1 parent 3ba8bc2
commit b534310
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 12 deletions.
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -21,6 +21,7 @@
       # System types to support.
       supportedSystems = [
         "x86_64-linux"
+        "aarch64-darwin"
       ];
 
       # Helper function to generate an attrset '{ x86_64-linux = f "x86_64-linux"; ... }'.
@@ -30,7 +31,8 @@
       nixpkgsFor =
         forAllSystems (system: import nixpkgs { inherit system; overlays = [ stacklock2nix.overlay self.overlay ]; });
 
-      deviceConfig = {cudaSupport=true;device="cuda-11";};
+#      deviceConfig = {cudaSupport=true;device="cuda-11";};
+      deviceConfig = {cudaSupport=false;device="cpu";};
     in
     {
       # A Nixpkgs overlay.
@@ -47,7 +49,7 @@
 
           # The Haskell package set to use as a base.  You should change this
           # based on the compiler version from the resolver in your stack.yaml.
-          baseHaskellPkgSet = final.haskell.packages.ghc924;
+          baseHaskellPkgSet = final.haskell.packages.ghc927;
 
           # Any additional Haskell package overrides you may want to add.
           additionalHaskellPkgSetOverrides = hfinal: hprev: {
@@ -65,6 +67,26 @@
                { doCheck = false;
                }
                hprev.tar;
+            half =
+              final.haskell.lib.compose.overrideCabal
+               { doCheck = false;
+               }
+               hprev.half;
+            sysinfo =
+              final.haskell.lib.compose.overrideCabal
+               { doCheck = false;
+               }
+               hprev.sysinfo;
+            inline-c-cpp =
+              final.haskell.lib.compose.overrideCabal
+               { doCheck = false;
+               }
+               hprev.inline-c-cpp;
+            happy =
+              final.haskell.lib.compose.overrideCabal
+               { doCheck = false;
+               }
+               hprev.happy;
             streaming-cassava =
               final.haskell.lib.compose.overrideCabal
                { preConfigure = ''
@@ -83,6 +105,10 @@
             torch = null;
           };
 
+          cabal2nixArgsOverrides = args: args // {
+            "splitmix" = verion: {};
+          };          
+
           # Additional packages that should be available for development.
           additionalDevShellNativeBuildInputs = stacklockHaskellPkgSet: [
             # Some Haskell tools (like cabal-install and ghcid) can be taken from the
@@ -137,10 +163,11 @@
       });
 
       packages = forAllSystems (system: {
+        hasktorch = nixpkgsFor.${system}.hasktorch-stacklock.pkgSet.hasktorch;
         examples = nixpkgsFor.${system}.hasktorch-stacklock.pkgSet.examples;
       });
 
-      defaultPackage = forAllSystems (system: self.packages.${system}.examples);
+      defaultPackage = forAllSystems (system: self.packages.${system}.hasktorch);
 
       devShells = forAllSystems (system: {
         hasktorch-dev-shell = nixpkgsFor.${system}.hasktorch-dev-shell;

diff --git a/nix/libtorch-binary-hashes.nix b/nix/libtorch-binary-hashes.nix
@@ -5,6 +5,11 @@ version : builtins.getAttr version {
       url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-1.11.0.zip";
       hash = "sha256-oTbvPjrREXPQanSjxzHbgJOtY5Yzb9FFgQsUG78o6eQ=";
     };
+    aarch64-darwin-cpu = {
+      name = "libtorch-macos-1.11.0.zip";
+      url = "https://download.pytorch.org/whl/cpu/torch-1.11.0-cp310-none-macosx_11_0_arm64.whl";
+      hash = "sha256-taOGgnabVEyHXsw0vLgfutXJIhObYTGarP/P2KMvUow=";
+    };
     x86_64-linux-cpu = {
       name = "libtorch-cxx11-abi-shared-with-deps-1.11.0-cpu.zip";
       url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.11.0%2Bcpu.zip";

diff --git a/nix/libtorch.nix b/nix/libtorch.nix
@@ -1,6 +1,8 @@
 { callPackage
 , stdenv
 , fetchzip
+, fetchurl
+, unzip
 , lib
 , libcxx
 
@@ -26,10 +28,15 @@ in stdenv.mkDerivation {
   inherit version;
   pname = "libtorch";
 
-  src = fetchzip srcs."${stdenv.targetPlatform.system}-${device}" or unavailable;
+  src = fetchurl srcs."${stdenv.targetPlatform.system}-${device}" or unavailable;
+
+  unpackPhase = ''
+    unzip $src
+    cp -r torch/* .
+  '';
 
   nativeBuildInputs =
-    if stdenv.isDarwin then [ fixDarwinDylibNames ]
+    if stdenv.isDarwin then [ fixDarwinDylibNames unzip ]
     else [ patchelf ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
 
   dontBuild = true;

diff --git a/stack.yaml b/stack.yaml
@@ -1,4 +1,4 @@
-resolver: nightly-2022-10-11
+resolver: lts-20.22
 
 #compiler: ghc-8.10.7