flakeify

exe/Main.hs

@@ -0,0 +1,6 @@
+module Main where
+
+import HasktorchSkeleton
+
+main :: IO ()
+main = putStrLn "Hello Hasktorch!"

flake.nix

@@ -0,0 +1,112 @@
+{
+  inputs = {
+    haskell-nix.url = "github:input-output-hk/haskell.nix";
+    stackageSrc = {
+      url = "github:input-output-hk/stackage.nix";
+      flake = false;
+    };
+    hackageSrc = {
+      url = "github:input-output-hk/hackage.nix";
+      flake = false;
+    };
+    utils.url = "github:numtide/flake-utils";
+    libtorch-nix = {
+      url = "github:hasktorch/libtorch-nix/d14b0fd10b96d192b92b8ccc7254ade4b3489331";
+      flake = false;
+    };
+  };
+
+  outputs = inputs@{ self, nixpkgs, haskell-nix, utils, ... }:
+    let
+      name = "hasktorchSkeleton";
+      compiler = "ghc8104"; # Not used for `stack.yaml` based projects.
+      cudaSupport = false;
+      cudaMajorVersion = null;
+      project-name = "${name}HaskellPackages";
+
+
+      # This overlay adds our project to pkgs
+      project-overlay = final: prev: {
+        ${project-name} =
+            #assert compiler == supported-compilers;
+            final.haskell-nix.project' {
+              # 'cleanGit' cleans a source directory based on the files known by git
+              src = prev.haskell-nix.haskellLib.cleanGit {
+                inherit name;
+                src = ./.;
+              };
+
+              compiler-nix-name = compiler;
+              projectFileName = "cabal.project"; # Not used for `stack.yaml` based projects.
+              modules = [
+                # Fixes for libtorch-ffi
+                {
+                  packages.libtorch-ffi = {
+                    configureFlags = with final; [
+                      "--extra-lib-dirs=${torch}/lib"
+                      "--extra-include-dirs=${torch}/include"
+                      "--extra-include-dirs=${torch}/include/torch/csrc/api/include"
+                    ];
+                    flags = {
+                      cuda = cudaSupport;
+                      gcc = !cudaSupport && final.stdenv.hostPlatform.isDarwin;
+                    };
+                  };
+                }
+              ];
+
+            };
+
+      };
+    in
+      { overlay = final: prev: {
+          "${name}" = ("${project-name}-overlay" final prev)."${project-name}".flake {};
+        };
+      } // (utils.lib.eachSystem [ "x86_64-linux" ] (system:
+        let
+          pkgs = import nixpkgs {
+            inherit system;
+            overlays = [
+              haskell-nix.overlay
+              (final: prev: {
+                haskell-nix = prev.haskell-nix // {
+                  sources = prev.haskell-nix.sources // {
+                    hackage = inputs.hackageSrc;
+                    stackage = inputs.stackageSrc;
+                  };
+                  modules = [
+                    # Fixes for libtorch-ffi
+                    {
+                      packages.libtorch-ffi = {
+                        configureFlags = with final; [
+                          "--extra-lib-dirs=${torch}/lib"
+                          "--extra-include-dirs=${torch}/include"
+                          "--extra-include-dirs=${torch}/include/torch/csrc/api/include"
+                        ];
+                        flags = {
+                          cuda = cudaSupport;
+                          gcc = !cudaSupport && final.stdenv.hostPlatform.isDarwin;
+                        };
+                      };
+                    }
+                  ];
+                };
+              })
+              (import ./nix/overlays/libtorch.nix { inherit inputs cudaSupport cudaMajorVersion; })
+              project-overlay
+            ];
+          };
+          flake = pkgs."${project-name}".flake {};
+        in flake // rec {
+
+          packages.example = flake.packages."${name}:exe:example";
+
+          defaultPackage = packages.example;
+
+          devShell = (import ./shell.nix {
+              inherit cudaSupport cudaMajorVersion pkgs;
+              withHoogle = false;
+            });
+        }
+        ));
+}

hasktorch-skeleton.cabal

@@ -1,4 +1,4 @@
-cabal-version:       1.24
+cabal-version:       2.2
 name:                hasktorch-skeleton
 version:             0.0.0.0
 synopsis:            See README for more info
@@ -14,13 +14,25 @@ extra-doc-files:     README.md
                    , CHANGELOG.md
 tested-with:         GHC == 8.8.3
 
-library
-  hs-source-dirs:    src
-  exposed-modules:   HasktorchSkeleton
+common base
+  ghc-options: -Wall -Wextra -Wno-unrecognised-pragmas -Wno-orphans
+  default-language: Haskell2010
+  build-depends:
+      base >= 4.12 && < 5
+    , hasktorch >= 0.2 && < 0.3
+
+common binary-base
+  import: base
+  ghc-options: -threaded -rtsopts -with-rtsopts=-N
+  build-depends: hasktorch-skeleton
 
-  build-depends:     base >= 4.7 && < 5
-                   , hasktorch >= 0.2 && < 0.3
-
-  ghc-options:       -Wall
+library
+  import: base
+  exposed-modules:
+    HasktorchSkeleton
+  hs-source-dirs: src
 
-  default-language:  Haskell2010
+executable example
+  import: binary-base
+  main-is: Main.hs
+  hs-source-dirs: exe

nix/default.nix

@@ -42,43 +42,11 @@ let
     ++ iohkNix.overlays.haskell-nix-extra
     # the iohkNix overlay contains nix utilities and niv
     ++ iohkNix.overlays.iohkNix
-    # libtorch overlays from pytorch-world
-    # TODO: pull in libGL_driver and cudatoolkit as done in https://github.com/NixOS/nixpkgs/blob/master/pkgs/games/katago/default.nix
     ++ [
-      (pkgs: _: with pkgs;
-        let libtorchSrc = callPackage "${sources.pytorch-world}/libtorch/release.nix" { }; in
-        if cudaSupport && cudaMajorVersion == "9" then
-          let libtorch = libtorchSrc.libtorch_cudatoolkit_9_2; in
-          {
-            c10 = libtorch;
-            torch = libtorch;
-            torch_cpu = libtorch;
-            torch_cuda = libtorch;
-          }
-        else if cudaSupport && cudaMajorVersion == "10" then
-          let libtorch = libtorchSrc.libtorch_cudatoolkit_10_2; in
-          {
-            c10 = libtorch;
-            torch = libtorch;
-            torch_cpu = libtorch;
-            torch_cuda = libtorch;
-          }
-        else if cudaSupport && cudaMajorVersion == "11" then
-          let libtorch = libtorchSrc.libtorch_cudatoolkit_11_0; in
-          {
-            c10 = libtorch;
-            torch = libtorch;
-            torch_cpu = libtorch;
-            torch_cuda = libtorch;
-          }
-        else
-          let libtorch = libtorchSrc.libtorch_cpu; in
-          {
-            c10 = libtorch;
-            torch = libtorch;
-            torch_cpu = libtorch;
-          }
-      )
+      (import ./overlays/libtorch.nix {
+        inherit cudaSupport cudaMajorVersion;
+        inputs = sources;
+      })
     ]
     # our own overlays:
     ++ [