Enable profile by default

.github/workflows/cabal-linux.yml

@@ -15,7 +15,7 @@ jobs:
         sudo add-apt-repository -y ppa:hvr/ghc
         sudo apt update -qq
         sudo apt -y purge ghc* cabal-install* || true
-        sudo apt -y --allow-downgrades --allow-remove-essential --allow-change-held-packages install build-essential zlib1g-dev liblapack-dev libblas-dev ghc-8.6.5 cabal-install-3.0 devscripts debhelper python3-pip cmake curl wget unzip git libtinfo-dev python3 python3-yaml
+        sudo apt -y --allow-downgrades --allow-remove-essential --allow-change-held-packages install build-essential zlib1g-dev liblapack-dev libblas-dev ghc-8.6.5-prof cabal-install-3.0 devscripts debhelper python3-pip cmake curl wget unzip git libtinfo-dev python3 python3-yaml
     - name: Setup repos
       run: |
         git submodule init && git submodule update
@@ -33,5 +33,5 @@ jobs:
         export PATH=/opt/ghc/bin:$PATH
         source setenv
         cabal v2-test all --jobs=2 --write-ghc-environment-files=always
-        cabal v2-exec codegen-exe
-        cabal exec xor-mlp
+        cabal v2-run codegen-exe
+        cabal v2-run xor-mlp

.github/workflows/nix-linux.yml

@@ -16,3 +16,19 @@ jobs:
       with:
         name: hasktorch
         signingKey: '${{ secrets.CACHIX_SIGNING_KEY }}'
+    - name: Run benchmark
+      run: |
+        nix-env -i mnist
+        ln -s ~/.nix-profile data
+        nix-build -A hasktorch-examples_cpu ./default.nix
+        export DEVICE=cpu
+        export NUM_ITERS=3000
+        ./result/bin/static-mnist-mlp +RTS -hc -p -A128m
+        cat static-mnist-mlp.prof
+        mkdir artifact
+        mv static-mnist-mlp.* artifact
+        mv loss.html artifact
+    - uses: actions/upload-artifact@master
+      with:
+        name: benchmark
+        path: artifact

.github/workflows/stack-linux.yml

@@ -29,5 +29,9 @@ jobs:
         stack test codegen
         stack test libtorch-ffi
         stack test hasktorch
-        stack exec codegen-exe
-        stack exec xor-mlp
+        stack run codegen-exe
+        stack run regression
+        stack run gaussian-process
+        stack run vae
+        stack run optimizers
+        stack run xor-mlp

.github/workflows/stack-macos.yaml

@@ -29,9 +29,9 @@ jobs:
         stack test codegen
         stack test libtorch-ffi
         stack test hasktorch
-        stack exec codegen-exe
-        stack exec xor-mlp
-        stack exec regression
-        stack exec gaussian-process
-        stack exec vae
-        stack exec optimizers
+        stack run codegen-exe
+        stack run regression
+        stack run gaussian-process
+        stack run vae
+        stack run optimizers
+        stack run xor-mlp

examples/examples.cabal

@@ -13,7 +13,7 @@ build-type:          Simple
 
 common config
   default-language:    Haskell2010
-  ghc-options:         -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver -fplugin GHC.TypeLits.Extra.Solver -fconstraint-solver-iterations=0
+  ghc-options:         -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver -fplugin GHC.TypeLits.Extra.Solver -fconstraint-solver-iterations=0 -rtsopts
   build-depends:       base >= 4.7 && < 5
                      , hasktorch
                     --  , ghc-typelits-extra >= 0.3.1

examples/static-mnist-cnn/Main.hs

@@ -34,6 +34,7 @@ import           System.Random
 
 import qualified ATen.Cast                     as ATen
 import qualified ATen.Class                    as ATen
+import qualified ATen.GC                       as ATen
 import qualified ATen.Type                     as ATen
 import qualified ATen.Managed.Type.Tensor      as ATen
 import qualified ATen.Managed.Type.Context     as ATen
@@ -115,9 +116,13 @@ instance ( KnownDType dtype
 type BatchSize = 512
 type TestBatchSize = 8192
 
-train :: forall (device :: (D.DeviceType, Nat)) . _ => IO ()
-train = do
-  let (numIters, printEvery) = (1000000, 250)
+train
+  :: forall (device :: (D.DeviceType, Nat))
+  . _
+  => Int
+  -> IO ()
+train numIters = do
+  let printEvery = 250
   (trainingData, testData) <- I.initMnist
   ATen.manual_seed_L 123
   init            <- A.sample (CNNSpec @ 'D.Float @device)
@@ -178,9 +183,13 @@ train = do
     return (crossEntropyLoss prediction target, errorRate prediction target)
 
 main :: IO ()
-main = do
+main = ATen.monitorMemory $ do
   deviceStr <- try (getEnv "DEVICE") :: IO (Either SomeException String)
+  numItersStr <- try (getEnv "NUM_ITERS") :: IO (Either SomeException String)
+  numIters <- case numItersStr of
+    Right iters -> return $ read iters
+    _ -> return 1000000
   case deviceStr of
-    Right "cpu"    -> train @'( 'D.CPU, 0)
-    Right "cuda:0" -> train @'( 'D.CUDA, 0)
+    Right "cpu"    -> train @'( 'D.CPU, 0) numIters
+    Right "cuda:0" -> train @'( 'D.CUDA, 0) numIters
     _              -> error "Don't know what to do or how."

examples/static-mnist-mlp/Main.hs

@@ -32,6 +32,7 @@ import           System.Random
 
 import qualified ATen.Cast                     as ATen
 import qualified ATen.Class                    as ATen
+import qualified ATen.GC                       as ATen
 import qualified ATen.Type                     as ATen
 import qualified ATen.Managed.Type.Tensor      as ATen
 import qualified ATen.Managed.Type.Context     as ATen
@@ -137,9 +138,10 @@ type HiddenFeatures1 = 256
 train
   :: forall (device :: (D.DeviceType, Nat))
    . _
-  => IO ()
-train = do
-  let (numIters, printEvery) = (1000000, 250)
+  => Int
+  -> IO ()
+train numIters = do
+  let printEvery = 250
       dropoutProb            = 0.5
   (trainingData, testData) <- I.initMnist
   ATen.manual_seed_L 123
@@ -210,9 +212,13 @@ train = do
     return (crossEntropyLoss prediction target, errorRate prediction target)
 
 main :: IO ()
-main = do
+main = ATen.monitorMemory $ do
   deviceStr <- try (getEnv "DEVICE") :: IO (Either SomeException String)
+  numItersStr <- try (getEnv "NUM_ITERS") :: IO (Either SomeException String)
+  numIters <- case numItersStr of
+    Right iters -> return $ read iters
+    _ -> return 1000000
   case deviceStr of
-    Right "cpu"    -> train @'( 'D.CPU, 0)
-    Right "cuda:0" -> train @'( 'D.CUDA, 0)
+    Right "cpu"    -> train @'( 'D.CPU, 0) numIters
+    Right "cuda:0" -> train @'( 'D.CUDA, 0) numIters
     _              -> error "Don't know what to do or how."

libtorch-ffi/libtorch-ffi.cabal

@@ -83,9 +83,10 @@ library
  if os(darwin)
   ld-options: -Wl,-keep_dwarf_unwind
   ghc-options:       -optc-std=c++11 -optc-xc++
+  cc-options:        -std=c++11 -xc++
  else
   ghc-options:       -optc-std=c++11
- cc-options:        -std=c++11
+  cc-options:        -std=c++11
  cxx-options:       -std=c++11
  default-extensions:          Strict
                             , StrictData
@@ -114,9 +115,10 @@ test-suite spec
   if os(darwin)
     ld-options: -Wl,-keep_dwarf_unwind
     ghc-options:       -optc-std=c++11 -optc-xc++
+    cc-options:        -std=c++11 -xc++
   else
     ghc-options:       -optc-std=c++11
-  cc-options:        -std=c++11
+    cc-options:        -std=c++11
   cxx-options:       -std=c++11
   default-extensions:          Strict
                              , StrictData

nix/shared.nix

@@ -28,11 +28,15 @@ let
                 overrideCabal = pkgsNew.haskell.lib.overrideCabal;
                 optionalString = pkgsNew.stdenv.lib.optionalString;
                 isDarwin = pkgsNew.stdenv.isDarwin;
+                enableLibraryProfiling = pkgsNew.haskell.lib.enableLibraryProfiling;
+                disableLibraryProfiling = pkgsNew.haskell.lib.disableLibraryProfiling;
+                enableExecutableProfiling = pkgsNew.haskell.lib.enableExecutableProfiling;
 
                 mkHasktorchExtension = postfix:
                   haskellPackagesNew: haskellPackagesOld: {
                     "libtorch-ffi_${postfix}" =
-                        appendConfigureFlag
+                      enableLibraryProfiling
+                        (appendConfigureFlag
                           (overrideCabal
                             (haskellPackagesOld.callCabal2nix
                               "libtorch-ffi"
@@ -48,28 +52,34 @@ let
                               }
                             )
                           )
-                        "--extra-include-dirs=${pkgsNew."libtorch_${postfix}"}/include/torch/csrc/api/include";
+                          "--extra-include-dirs=${pkgsNew."libtorch_${postfix}"}/include/torch/csrc/api/include"
+                        );
                     "hasktorch_${postfix}" =
-                      overrideCabal
-                        (haskellPackagesOld.callCabal2nix
-                          "hasktorch"
-                          ../hasktorch
-                          { libtorch-ffi = haskellPackagesNew."libtorch-ffi_${postfix}"; }
-                        )
-                        (old: {
-                              preConfigure = (old.preConfigure or "") + optionalString (!isDarwin) ''
-                                export LD_PRELOAD=${pkgs.mkl}/lib/libmkl_rt.so
-                              '';
-                            }
+                      enableLibraryProfiling
+                        (overrideCabal
+                           (haskellPackagesOld.callCabal2nix
+                             "hasktorch"
+                             ../hasktorch
+                             { libtorch-ffi = haskellPackagesNew."libtorch-ffi_${postfix}"; }
+                           )
+                           (old: {
+                                 preConfigure = (old.preConfigure or "") + optionalString (!isDarwin) ''
+                                   export LD_PRELOAD=${pkgs.mkl}/lib/libmkl_rt.so
+                                 '';
+                               }
+                           )
                         );
                     "hasktorch-examples_${postfix}" =
                       # failOnAllWarnings
-                        (haskellPackagesOld.callCabal2nix
-                          "examples"
-                          ../examples
-                          { libtorch-ffi = haskellPackagesNew."libtorch-ffi_${postfix}"
-                          ; hasktorch = haskellPackagesNew."hasktorch_${postfix}"
-                          ; }
+                      enableLibraryProfiling
+                        (enableExecutableProfiling
+                          (haskellPackagesOld.callCabal2nix
+                            "examples"
+                            ../examples
+                            { libtorch-ffi = haskellPackagesNew."libtorch-ffi_${postfix}"
+                              ; hasktorch = haskellPackagesNew."hasktorch_${postfix}"
+                              ; }
+                          )
                         );
                   };