libtorch-ffi 2.0.1.5 → 2.0.1.6

raw patch · 3 files changed

+346/−2 lines, 3 filesbuild-type:Customsetup-changed

Files

• README.md +134/−0
• Setup.hs +198/−0
• libtorch-ffi.cabal +14/−2

@@ -0,0 +1,134 @@+# libtorch-ffi++This package provides FFI bindings to PyTorch's libtorch C++ library.++## Setup++The package automatically downloads and configures libtorch during the build process. You can customize the setup using environment variables.++### Environment Variables++#### `LIBTORCH_VERSION`+- **Default**: `2.5.0`+- **Description**: Specifies the version of libtorch to download and use+- **Example**: `export LIBTORCH_VERSION=2.5.0`++#### `LIBTORCH_HOME`+- **Default**: XDG cache directory (`~/.cache/libtorch` on Linux/macOS)+- **Description**: Base directory where libtorch will be downloaded and stored+- **Example**: `export LIBTORCH_HOME=/opt/libtorch`++#### `LIBTORCH_CUDA_VERSION`+- **Default**: `cpu`+- **Description**: CUDA version for GPU support+- **Options**: +  - `cpu` - CPU-only version (default)+  - `cu117` - CUDA 11.7+  - `cu118` - CUDA 11.8+  - `cu121` - CUDA 12.1+  - Any other CUDA version string supported by PyTorch+- **Example**: `export LIBTORCH_CUDA_VERSION=cu118`++#### `LIBTORCH_SKIP_DOWNLOAD`+- **Default**: Not set+- **Description**: When set (to any value), skips the automatic download of libtorch+- **Use case**: When you have libtorch already installed system-wide+- **Example**: `export LIBTORCH_SKIP_DOWNLOAD=1`++### Directory Structure++The downloaded libtorch is stored in a platform-specific directory structure:+```+$LIBTORCH_HOME/+└── <version>/+    └── <platform>/+        └── <cuda-flavor>/+            ├── lib/+            ├── include/+            └── .ok+```++Where:+- `<version>` is the libtorch version (e.g., `2.5.0`)+- `<platform>` is one of:+  - `macos-arm64` - macOS on Apple Silicon+  - `macos-x86_64` - macOS on Intel+  - `linux-x86_64` - Linux on x86_64+- `<cuda-flavor>` is the CUDA version (e.g., `cpu`, `cu118`)++### Build Process++1. **Pre-configuration**: The package checks if it's running in a Nix sandbox. If not, it proceeds with the download process.++2. **Download**: If libtorch is not found in the cache directory, it will be automatically downloaded from PyTorch's official servers.++3. **Configuration**: The build system automatically:+   - Adds the libtorch library directory to the library search path+   - Adds the include directories for C++ headers+   - Sets up proper runtime library paths (rpath) for dynamic linking+   - On macOS, adds the `-ld_classic` flag for compatibility++### Platform-Specific Notes++#### macOS+- Uses rpath for dynamic library loading+- Automatically adds `-ld_classic` flag for linker compatibility+- Supports both Apple Silicon (arm64) and Intel (x86_64) architectures+- **Since libtorch-ffi's rpath is propagated, it doesn't matter whether hasktorch is a static link or a shared link**++#### Linux+- Uses rpath for dynamic library loading+- Supports x86_64 architecture+- Multiple CUDA versions available for GPU support+- **Since libtorch-ffi's rpath is not propagated, hasktorch must be a shared link**++### Linking Configuration++Due to rpath propagation differences between platforms, Linux requires shared linking. Add the following configuration:++#### For Cabal (cabal.project)+```+shared: True+executable-dynamic: True+```++#### For Stack (stack.yaml)+```yaml+configure-options:+  $targets:+    - --enable-executable-dynamic+    - --enable-shared+```++### Nix Support++The package detects when it's being built in a Nix sandbox and skips the automatic download. In this case, libtorch should be provided through Nix derivation inputs.++### Troubleshooting++1. **Download failures**: Check your internet connection and ensure the PyTorch download servers are accessible.++2. **Missing libraries**: The `.ok` marker file indicates a successful download. If this file is missing but the directory exists, delete the directory and let the setup download again.++3. **CUDA version mismatch**: Ensure your system CUDA version matches the `LIBTORCH_CUDA_VERSION` you've specified.++4. **Custom libtorch installation**: Set `LIBTORCH_SKIP_DOWNLOAD=1` and ensure your system's libtorch is properly configured in your build environment.++### Example Usage++```bash+# Use CPU-only version+cabal build libtorch-ffi++# Use CUDA 11.8 version+export LIBTORCH_CUDA_VERSION=cu118+cabal build libtorch-ffi++# Use a specific version+export LIBTORCH_VERSION=2.4.0+cabal build libtorch-ffi++# Use existing system libtorch+export LIBTORCH_SKIP_DOWNLOAD=1+cabal build libtorch-ffi+```

@@ -0,0 +1,198 @@+import Data.List (isPrefixOf)+import Distribution.Simple+import Distribution.Simple.Program+import Distribution.Simple.Setup+import Distribution.Simple.LocalBuildInfo+import Distribution.Types.LocalBuildInfo+import Distribution.Types.GenericPackageDescription+import Distribution.Types.HookedBuildInfo+import Distribution.PackageDescription+import Distribution.System+import System.Directory+import System.FilePath+import System.Process+import System.IO.Temp+import Control.Monad+import Network.HTTP.Simple+import qualified Data.ByteString.Lazy as LBS+import System.Environment (lookupEnv)+import Data.Maybe (fromMaybe)+import GHC.IO.Exception+import Control.Exception++main :: IO ()+main = defaultMainWithHooks $ simpleUserHooks+  { preConf = \_ _ -> do+      pure emptyHookedBuildInfo+  , confHook = \(gpd, hbi) flags -> do+      mlibtorchDir <- ensureLibtorch+      case mlibtorchDir of+        Nothing -> do+          putStrLn "libtorch not found, skipping configuration."+          lbi <- confHook simpleUserHooks (gpd, hbi) flags+          -- For macOS, add the -ld_classic flag to the linker+          case buildOS of+            OSX -> return $ lbi { withPrograms = addLdClassicFlag (withPrograms lbi) }+            _ -> return $ lbi+        Just libtorchDir -> do+          libtorchDir <- getLocalUserLibtorchDir+          let libDir     = libtorchDir </> "lib"+              includeDir = libtorchDir </> "include"++          let updatedFlags = flags+                { configExtraLibDirs      = libDir : configExtraLibDirs flags+                , configExtraIncludeDirs  =+                    includeDir+                    : (includeDir </> "torch" </> "csrc" </> "api" </> "include")+                    : configExtraIncludeDirs flags+                }+          -- Call the default configuration hook with updated flags+          lbi <- confHook simpleUserHooks (gpd, hbi) updatedFlags+          -- For macOS, add the -ld_classic flag to the linker+          case buildOS of+            OSX -> return $ lbi { withPrograms = addRPath libDir $ addLdClassicFlag (withPrograms lbi) }+            Linux -> return $ lbi { withPrograms = addRPath libDir (withPrograms lbi) }+            _ -> return $ lbi+  }++getLibtorchVersion :: IO String+getLibtorchVersion = do+  mVersion <- lookupEnv "LIBTORCH_VERSION"+  case mVersion of+    Nothing -> return "2.5.0"+    Just other -> return other++getLocalUserLibtorchDir :: IO FilePath+getLocalUserLibtorchDir = do+  mHome <- lookupEnv "LIBTORCH_HOME"+  libtorchVersion <- getLibtorchVersion+  base <- case mHome of+    Just h  -> pure h+    Nothing -> do+      -- XDG cache (Linux/macOS). Falls back to ~/.cache+      cache <- getXdgDirectory XdgCache "libtorch"+      pure cache+  flavor <- getCudaFlavor+  pure $ base </> libtorchVersion </> platformTag </> flavor++platformTag :: FilePath+platformTag =+  case (buildOS, buildArch) of+    (OSX,    AArch64) -> "macos-arm64"+    (OSX,    X86_64)  -> "macos-x86_64"+    (Linux,  X86_64)  -> "linux-x86_64"+    -- add more as needed+    _ -> error $ "Unsupported platform: " <> show (buildOS, buildArch)++getCudaFlavor :: IO String+getCudaFlavor = do+  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu117" | "cu118" | "cu121"++ensureLibtorch :: IO (Maybe FilePath)+ensureLibtorch = do+  isSandbox <- isNixSandbox+  if isSandbox+    then return Nothing+    else downloadLibtorch++isNixSandbox :: IO Bool+isNixSandbox = do+  nix <- lookupEnv "NIX_BUILD_TOP"+  case nix of+    Just path -> do+      let isNixPath = any (`isPrefixOf` path) ["/build", "/private/tmp/nix-build"]+      if isNixPath+        then do+          putStrLn "Nix sandbox detected; skipping libtorch download."+          return True+        else do+          return False+    Nothing -> return False++downloadLibtorch :: IO (Maybe FilePath)+downloadLibtorch = do+  skip <- lookupEnv "LIBTORCH_SKIP_DOWNLOAD"+  case skip of+    Just _ -> do+      putStrLn "LIBTORCH_SKIP_DOWNLOAD set; assuming libtorch exists globally."+      return Nothing+    Nothing -> do+      dest <- getLocalUserLibtorchDir+      let marker = dest </> ".ok"+      exists <- doesFileExist marker+      present <- doesDirectoryExist dest+      if present && exists+        then pure $ Just dest+        else do+          putStrLn $ "libtorch not found in local cache, installing to " <> dest+          downloadAndExtractLibtorchTo dest+          -- Create an idempotence marker that checks+          -- if we've already downloaded torch.+          -- Since we'll be moving everything this will+          -- be the our main reference.+          writeFile marker ""+          pure $ Just dest++downloadAndExtractLibtorchTo :: FilePath -> IO ()+downloadAndExtractLibtorchTo dest = do+  createDirectoryIfMissing True dest+  (url, fileName) <- computeURL+  putStrLn $ "Downloading libtorch from: " ++ url+  withSystemTempDirectory "libtorch-download" $ \tmpDir -> do+    let downloadPath = tmpDir </> fileName+    request  <- parseRequest url+    response <- httpLBS request+    LBS.writeFile downloadPath (getResponseBody response)+    putStrLn "Download complete. Extracting..."+    callProcess "unzip" ["-q", downloadPath, "-d", tmpDir]+    let unpacked = tmpDir </> "libtorch"+    exists <- doesDirectoryExist unpacked+    let src = if exists then unpacked else tmpDir+    -- We want to move the directory since this operation is atomic.+    -- If that doesn't work we fall back to copying.+    (renameDirectory src dest) `catch` (\(_::IOException) -> copyTree src dest)+    putStrLn "libtorch extracted successfully (global cache)."++copyTree :: FilePath -> FilePath -> IO ()+copyTree src dest = do+  createDirectoryIfMissing True dest+  entries <- listDirectory src+  mapM_ (\e -> do+          let s = src  </> e+              d = dest </> e+          isDir <- doesDirectoryExist s+          if isDir then copyTree s d else copyFile s d+        ) entries++computeURL :: IO (String, String)+computeURL = do+  flavor <- getCudaFlavor+  v <- getLibtorchVersion+  pure $ case buildOS of+    OSX -> case buildArch of+      AArch64 -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-" ++ v ++ ".zip"+                 , "libtorch-macos-arm64.zip" )+      X86_64  -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-macos-x86_64-" ++ v ++ ".zip"+                 , "libtorch-macos-x86_64.zip" )+      _       -> error "Unsupported macOS arch"+    Linux -> case flavor of+      "cpu"  -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-" ++ v ++ "%2Bcpu.zip"+                , "libtorch-linux.zip" )+      cudaVersion -> ( "https://download.pytorch.org/libtorch/" ++ cudaVersion ++"/libtorch-cxx11-abi-shared-with-deps-" ++ v ++ "%2B" ++ cudaVersion ++ ".zip"+                , "libtorch-linux-" ++ cudaVersion ++ ".zip" )+    Windows -> error "Windows not supported by this setup"++-- Add -ld_classic flag to GHC program arguments for macOS+addLdClassicFlag :: ProgramDb -> ProgramDb+addLdClassicFlag progDb = +  case lookupProgram ghcProgram progDb of+    Just ghc ->+      let ghc' = ghc { programOverrideArgs = ["-optl-ld_classic"] ++ programOverrideArgs ghc }+      in updateProgram ghc' progDb+    Nothing -> progDb++addRPath :: FilePath -> ProgramDb -> ProgramDb+addRPath libDir progDb =+  userSpecifyArgs (programName ldProgram)+  ["-Wl,-rpath," ++ libDir]+  progDb

@@ -1,6 +1,6 @@ cabal-version:       3.0 name:                libtorch-ffi-version:             2.0.1.5+version:             2.0.1.6 -- The prefix(2.0) of this version("2.0.0.0") is the same as libtorch's one. synopsis:            Haskell bindings for PyTorch description:         This package provides Haskell bindings to libtorch, the C++ library underlying PyTorch, specifically designed for the Hasktorch ecosystem.@@ -10,7 +10,19 @@ maintainer:          hasktorch@gmail.com copyright:           2018 Austin Huang category:            Codegen-build-type:          Simple+build-type:          Custom+extra-source-files:  README.md++custom-setup+  setup-depends:+    base >= 4.7 && < 5,+    Cabal >= 3.0 && < 4,+    directory >= 1.3 && < 2,+    filepath >= 1.4 && < 2,+    process >= 1.6 && < 2,+    temporary >= 1.3 && < 2,+    http-conduit >= 2.3 && < 3,+    bytestring >= 0.10 && < 1  Flag cuda  Description: A flag to link libtorch_cuda.