libtorch-ffi 2.0.1.10 → 2.0.2.0

raw patch · 6 files changed

+428/−17 lines, 6 filessetup-changedPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +8/−0
• Setup.hs +269/−10
• libtorch-ffi.cabal +2/−1
• src/Torch/Internal/GC.hs +1/−1
• src/Torch/Internal/Managed/Optim.hs +27/−0
• src/Torch/Internal/Unmanaged/Optim.hs +121/−5

@@ -0,0 +1,8 @@+# Changelog for libtorch-ffi++## 2.0.2.0++- Add support for Stackage lts-23.24 (GHC 9.8); the previous release used GHC 9.6.+- Expose parameter group interface in `Torch.Internal.Managed.Optim` and `Torch.Internal.Unmanaged.Optim` (supports per-group Adam/AdamW options).+- Implement non-blocking host-to-device transfer.+- Setup: dynamic detection for CUDA package installation, download NVIDIA libs, drop `rpath-link` on macOS.

@@ -1,6 +1,9 @@ {-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-} -import Data.List (isPrefixOf)+import Data.Char (toLower)+import Data.List (isPrefixOf, isInfixOf, nubBy) import Distribution.Simple import Distribution.Simple.Program import Distribution.Simple.Setup@@ -17,8 +20,9 @@ import Control.Monad import Network.HTTP.Simple import qualified Data.ByteString.Lazy as LBS+import qualified Data.ByteString.Lazy.Char8 as LBSC import System.Environment (lookupEnv)-import Data.Maybe (fromMaybe)+import Data.Maybe (fromMaybe, listToMaybe) import GHC.IO.Exception import Control.Exception import Codec.Archive.Zip@@ -39,7 +43,7 @@       mlibtorchDir <- ensureLibtorch       case mlibtorchDir of         Nothing -> do-          putStrLn "libtorch not found, skipping configuration."+          putStrLn "libtorch not found or handled by Nix, skipping configuration."           lbi <- confHook simpleUserHooks (gpd, hbi) flags           -- For macOS, add the -ld_classic flag to the linker           case buildOS of@@ -59,7 +63,8 @@                 }           -- Call the default configuration hook with updated flags           lbi <- confHook simpleUserHooks (gpd, hbi) updatedFlags-          -- For macOS, add the -ld_classic flag to the linker+          +          -- Add RPath so the binary finds the libs at runtime           case buildOS of             OSX -> return $ lbi { withPrograms = addRPath libDir $ addLdClassicFlag (withPrograms lbi) }             Linux -> return $ lbi { withPrograms = addRPath libDir (withPrograms lbi) }@@ -97,7 +102,7 @@  getCudaFlavor :: IO String getCudaFlavor = do-  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu126" | "cu128" | "cu130"+  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu121" | "cu118"  ensureLibtorch :: IO (Maybe FilePath) ensureLibtorch = do@@ -137,10 +142,15 @@         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.++          -- Download NVIDIA libs if on Linux + CUDA+          flavor <- getCudaFlavor+          when (buildOS == Linux && "cu" `isPrefixOf` flavor) $ do+             -- Dynamically detect which NVIDIA library versions are needed+             -- and extract them to dest/lib (same directory as libtorch_cuda.so)+             downloadNvidiaLibs dest++          -- Create an idempotence marker           writeFile marker ""           pure $ Just dest @@ -160,11 +170,259 @@     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)." +-- | Downloads necessary NVIDIA wheels from PyPI and extracts libs to dest/lib+-- Dynamically detects which library versions are needed by inspecting libtorch+-- Uses iterative detection to find transitive dependencies+downloadNvidiaLibs :: FilePath -> IO ()+downloadNvidiaLibs dest = do+  let libDir = dest </> "lib"++  -- Pass 1: Detect and download direct dependencies from libtorch+  putStrLn "Pass 1: Detecting NVIDIA library dependencies from libtorch..."+  neededLibs <- detectNeededNvidiaLibs libDir++  if null neededLibs+    then putStrLn "No missing NVIDIA libraries detected (all bundled with libtorch)."+    else do+      putStrLn $ "Found " ++ show (length neededLibs) ++ " NVIDIA libraries to download:"+      forM_ neededLibs $ \(lib, ver) ->+        putStrLn $ "  - " ++ lib ++ " (version " ++ ver ++ ")"+      downloadAndExtractPackages dest neededLibs++      -- Pass 2: Check downloaded libraries for transitive dependencies+      putStrLn "\nPass 2: Checking downloaded libraries for transitive dependencies..."+      transitiveLibs <- detectNeededNvidiaLibs libDir++      if null transitiveLibs+        then putStrLn "No additional transitive dependencies found."+        else do+          putStrLn $ "Found " ++ show (length transitiveLibs) ++ " additional libraries to download:"+          forM_ transitiveLibs $ \(lib, ver) ->+            putStrLn $ "  - " ++ lib ++ " (version " ++ ver ++ ")"+          downloadAndExtractPackages dest transitiveLibs++-- | Detect which NVIDIA libraries are needed by inspecting libtorch's dependencies+-- Returns list of (library name, major version) tuples, e.g., [("cusparse", "12"), ("cufft", "11")]+detectNeededNvidiaLibs :: FilePath -> IO [(String, String)]+detectNeededNvidiaLibs libDir = do+  putStrLn $ "  Inspecting libraries in: " ++ libDir++  -- Libraries we care about (including transitive dependencies)+  let targetLibs = ["cusparse", "cufft", "curand", "cublas", "cusolver", "nvjitlink"]++  -- Check libtorch libraries first+  let libtorchFiles = ["libtorch_cuda.so", "libtorch.so"]++  -- Also check any NVIDIA libraries that were already downloaded (for transitive deps)+  libDirContents <- listDirectory libDir+  let nvidiaLibs = filter (\f -> any (\target -> ("lib" ++ target) `isPrefixOf` f && ".so" `isInfixOf` f) targetLibs) libDirContents++  let soFiles = libtorchFiles ++ nvidiaLibs++  needed <- forM soFiles $ \soFile -> do+    let fullPath = libDir </> soFile+    exists <- doesFileExist fullPath+    if not exists+      then do+        putStrLn $ "  Skipping " ++ soFile ++ " (not found)"+        return []+      else do+        putStrLn $ "  Analyzing dependencies of " ++ soFile ++ "..."+        deps <- extractNvidiaDeps fullPath targetLibs+        forM_ deps $ \(lib, ver) ->+          putStrLn $ "    Found dependency: lib" ++ lib ++ ".so." ++ ver+        return deps++  -- Flatten results and remove duplicates+  let allNeeded = concat needed+      uniqueNeeded = nubBy (\(a,_) (b,_) -> a == b) allNeeded++  putStrLn $ "  Total unique NVIDIA dependencies found: " ++ show (length uniqueNeeded)++  -- Filter out libraries that are already bundled+  filterM (notBundled libDir) uniqueNeeded++-- | Check if a library is already bundled in libtorch+notBundled :: FilePath -> (String, String) -> IO Bool+notBundled libDir (libName, _ver) = do+  files <- listDirectory libDir+  -- Match exactly "libNAME.so" or "libNAME-*.so" to avoid false matches+  -- (e.g., "libcusparse" should not match "libcusparseLt")+  let pattern = "lib" ++ libName+  let bundled = any (\f ->+        (f == pattern ++ ".so" ||+         (pattern ++ ".so.") `isPrefixOf` f ||+         (pattern ++ "-") `isPrefixOf` f && ".so" `isInfixOf` f)) files+  return (not bundled)++-- | Extract NVIDIA library dependencies using readelf+extractNvidiaDeps :: FilePath -> [String] -> IO [(String, String)]+extractNvidiaDeps soFile targetLibs = do+  -- Try readelf first+  result <- tryReadelf soFile+  case result of+    Just output -> return $ parseNvidiaDeps output targetLibs+    Nothing -> do+      -- Fallback: try objdump+      result2 <- tryObjdump soFile+      case result2 of+        Just output -> return $ parseNvidiaDeps output targetLibs+        Nothing -> return []++tryReadelf :: FilePath -> IO (Maybe String)+tryReadelf soFile = do+  result <- try $ readProcess "readelf" ["-d", soFile] ""+  case result of+    Right output -> return (Just output)+    Left (_ :: IOException) -> return Nothing++tryObjdump :: FilePath -> IO (Maybe String)+tryObjdump soFile = do+  result <- try $ readProcess "objdump" ["-p", soFile] ""+  case result of+    Right output -> return (Just output)+    Left (_ :: IOException) -> return Nothing++-- | Parse readelf/objdump output to extract NVIDIA library dependencies+-- Example: "libcusparse.so.12" -> ("cusparse", "12")+-- Case-insensitive matching to handle libraries like "libnvJitLink" vs "libnvjitlink"+parseNvidiaDeps :: String -> [String] -> [(String, String)]+parseNvidiaDeps output targetLibs =+  [ (lib, ver)+  | line <- lines output+  , lib <- targetLibs+  , let pattern = "lib" ++ lib ++ ".so."+  , let lineLower = map toLower line+  , let patternLower = map toLower pattern+  , patternLower `isInfixOf` lineLower+  , let ver = extractVersion line lib+  , not (null ver)+  ]++extractVersion :: String -> String -> String+extractVersion line libName =+  case findSubstring pattern line of+    Just idx ->+      let afterPattern = drop (idx + length pattern) line+      in takeWhile isDigit afterPattern+    Nothing -> ""+  where+    pattern = "lib" ++ libName ++ ".so."+    isDigit c = c >= '0' && c <= '9'++    -- Case-insensitive substring search+    findSubstring :: String -> String -> Maybe Int+    findSubstring needle haystack = go 0 haystack+      where+        needleLower = map toLower needle+        go _ [] = Nothing+        go idx str+          | length str >= length needle &&+            map toLower (take (length needle) str) == needleLower = Just idx+          | otherwise = go (idx + 1) (tail str)++downloadAndExtractPackages :: FilePath -> [(String, String)] -> IO ()+downloadAndExtractPackages dest neededLibs =+  withSystemTempDirectory "nvidia-libs-download" $ \tmpDir -> do+    forM_ neededLibs $ \(libName, majorVer) -> do+      -- Special case: nvidia-cufft-cu12 provides .so.11, but we need .so.12+      -- In this case, use the generic nvidia-cufft package which has 12.x+      let useGenericFirst = (libName == "cufft" && majorVer == "12")++      -- Map major version to PyPI suffix (12 -> cu12, 11 -> cu11, etc.)+      let pypiSuffix = "cu" ++ majorVer+      let pkgWithSuffix = "nvidia-" ++ libName ++ "-" ++ pypiSuffix+      let pkgGeneric = "nvidia-" ++ libName++      -- Try generic first for known mismatches, otherwise try CUDA-specific first+      (finalPkg, finalUrl) <- if useGenericFirst+        then do+          putStrLn $ "Fetching metadata for " ++ pkgGeneric ++ " (known version mismatch)..."+          mUrl <- getPyPiWheelUrl pkgGeneric+          case mUrl of+            Just url -> return (pkgGeneric, Just url)+            Nothing -> do+              putStrLn $ "  Not found, trying: " ++ pkgWithSuffix ++ "..."+              url <- getPyPiWheelUrl pkgWithSuffix+              return (pkgWithSuffix, url)+        else do+          putStrLn $ "Fetching metadata for " ++ pkgWithSuffix ++ "..."+          mUrl <- getPyPiWheelUrl pkgWithSuffix+          -- Fallback to generic package name if versioned one not found+          -- (e.g., nvidia-curand instead of nvidia-curand-cu10)+          case mUrl of+            Just url -> return (pkgWithSuffix, Just url)+            Nothing -> do+              putStrLn $ "  Not found, trying fallback: " ++ pkgGeneric ++ "..."+              url <- getPyPiWheelUrl pkgGeneric+              return (pkgGeneric, url)++      case finalUrl of+        Nothing -> putStrLn $ "Warning: Could not find manylinux wheel for " ++ finalPkg+        Just url -> do+          let fileName = takeFileName url+          let downloadPath = tmpDir </> fileName+          putStrLn $ "Downloading " ++ fileName ++ "..."+          +          request  <- parseRequest url+          response <- httpLBS request+          LBS.writeFile downloadPath (getResponseBody response)+          +          putStrLn $ "Extracting libraries from " ++ fileName ++ "..."+          archive <- toArchive <$> LBS.readFile downloadPath+          +          -- Iterate over entries and extract only .so files to dest/lib+          let libDest = dest </> "lib"+          createDirectoryIfMissing True libDest+          +          forM_ (zEntries archive) $ \entry -> do+            let path = eRelativePath entry+                isSharedObj = ".so" `isInfixOf` path+                -- Wheels put libs in nvidia/<pkg>/lib/ or lib/+                isLibDir = "lib/" `isInfixOf` path +            +            when (isSharedObj && isLibDir) $ do+              let entryFileName = takeFileName path+              let targetPath = libDest </> entryFileName+              putStrLn $ "  Extracting: " ++ entryFileName ++ " -> " ++ targetPath+              -- Write the file+              let entryData = fromEntry entry+              LBS.writeFile targetPath entryData+              -- Set readable and executable permissions for shared libraries+              -- Using only owner permissions (older directory library compatibility)+              let perms = foldl (flip ($)) emptyPermissions+                          [ setOwnerReadable True+                          , setOwnerExecutable True+                          ]+              setPermissions targetPath perms++-- | Quick and dirty PyPI JSON parser to find manylinux x86_64 url+-- Avoids adding Aeson dependency to Setup.hs+getPyPiWheelUrl :: String -> IO (Maybe String)+getPyPiWheelUrl pkg = do+  let jsonUrl = "https://pypi.org/pypi/" ++ pkg ++ "/json"+  request <- parseRequest jsonUrl+  response <- httpLBS request+  let body = getResponseBody response+  +  -- We look for the "url" field inside an object that also has "manylinux" and "x86_64" in the filename+  -- This is a heuristic search on the raw JSON string+  let urls = extractUrls body+  return $ listToMaybe [ u | u <- urls, "manylinux" `isInfixOf` u, "x86_64" `isInfixOf` u ]++-- | Helper to extract all strings that look like URLs from JSON+extractUrls :: LBS.ByteString -> [String]+extractUrls content = +  let parts = LBSC.split '"' content+      -- Filter for https URLs+  in [ LBSC.unpack p | p <- parts, "https://" `LBSC.isPrefixOf` p, ".whl" `LBSC.isSuffixOf` p ]+ copyTree :: FilePath -> FilePath -> IO () copyTree src dest = do   createDirectoryIfMissing True dest@@ -206,5 +464,6 @@ addRPath :: FilePath -> ProgramDb -> ProgramDb addRPath libDir progDb =   userSpecifyArgs (programName ldProgram)-  ["-Wl,-rpath," ++ libDir]+  [ "-Wl,-rpath," ++ libDir          -- for runtime library search+  ]   progDb

@@ -1,6 +1,6 @@ cabal-version:       3.0 name:                libtorch-ffi-version:             2.0.1.10+version:             2.0.2.0 -- 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.@@ -12,6 +12,7 @@ category:            Codegen build-type:          Custom extra-source-files:  README.md+                     CHANGELOG.md  custom-setup   setup-depends:

@@ -96,7 +96,7 @@ #ifdef darwin_HOST_OS     msgOutOfMemory = "MPS backend out of memory" #else-    msgOutOfMemory = "Exception: CUDA out of memory."+    msgOutOfMemory = "CUDA out of memory." #endif {-# INLINE retryWithGC' #-} 

@@ -148,3 +148,30 @@  load :: ForeignPtr Optimizer -> ForeignPtr StdString -> IO () load = _cast2 Unmanaged.load++adamwWithParamGroups+  :: CDouble+  -> CDouble+  -> CDouble+  -> CDouble+  -> CDouble+  -> CBool+  -> ForeignPtr TensorList+  -> ForeignPtr TensorList+  -> IO (ForeignPtr Optimizer)+adamwWithParamGroups = _cast8 Unmanaged.adamwWithParamGroups++getAllParams :: ForeignPtr Optimizer -> IO (ForeignPtr TensorList)+getAllParams = _cast1 Unmanaged.getAllParams++stepOnly :: ForeignPtr Optimizer -> IO ()+stepOnly = _cast1 Unmanaged.stepOnly++zeroGrad :: ForeignPtr Optimizer -> IO ()+zeroGrad = _cast1 Unmanaged.zeroGrad++setParamGrads :: ForeignPtr Optimizer -> ForeignPtr TensorList -> IO ()+setParamGrads = _cast2 Unmanaged.setParamGrads++setLr :: ForeignPtr Optimizer -> CDouble -> IO ()+setLr = _cast2 Unmanaged.setLr

@@ -261,10 +261,17 @@     return dynamic_cast<torch::optim::Optimizer*>(optimizer);   }|] -getParams :: Ptr Optimizer -> IO (Ptr TensorList) +getParams :: Ptr Optimizer -> IO (Ptr TensorList) getParams optimizer =   [C.throwBlock| std::vector<at::Tensor>* {-    return new std::vector<at::Tensor>($(torch::optim::Optimizer* optimizer)->param_groups().at(0).params());+    auto optimizer = $(torch::optim::Optimizer* optimizer);+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();+    for(auto& group : optimizer->param_groups()){+      for(auto& p : group.params()){+        result->push_back(p);+      }+    }+    return result;   }|]  step :: Ptr Optimizer -> (Ptr TensorList -> IO (Ptr Tensor)) -> IO (Ptr Tensor)@@ -280,7 +287,13 @@         auto func = (Func)tfunc;         auto v = optimizer->step([&]{           optimizer->zero_grad();-          auto loss = func(&(optimizer->param_groups().at(0).params()));+          std::vector<at::Tensor> all_params;+          for(auto& group : optimizer->param_groups()){+            for(auto& p : group.params()){+              all_params.push_back(p);+            }+          }+          auto loss = func(&all_params);           loss->backward();           return *loss;         });@@ -304,7 +317,13 @@         auto func = (Func)tfunc;         auto v = optimizer->step([&]{           optimizer->zero_grad();-          auto lossWithGenerator = func(&(optimizer->param_groups().at(0).params()),&generator);+          std::vector<at::Tensor> all_params;+          for(auto& group : optimizer->param_groups()){+            for(auto& p : group.params()){+              all_params.push_back(p);+            }+          }+          auto lossWithGenerator = func(&all_params,&generator);           auto loss = std::get<0>(*lossWithGenerator);           generator = std::get<1>(*lossWithGenerator);           loss.backward();@@ -326,7 +345,13 @@     optimizer->zero_grad();     loss->backward();     optimizer->step();-    return new std::vector<at::Tensor>(optimizer->param_groups().at(0).params());+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();+    for(auto& group : optimizer->param_groups()){+      for(auto& p : group.params()){+        result->push_back(p);+      }+    }+    return result;   }|]  save :: Ptr Optimizer -> Ptr StdString -> IO ()@@ -341,4 +366,95 @@   [C.throwBlock| void {     std::ifstream input(*$(std::string* filename));     torch::load(*$(torch::optim::Optimizer* optimizer),input);+  }|]++adamwWithParamGroups+  :: CDouble+  -> CDouble+  -> CDouble+  -> CDouble+  -> CDouble+  -> CBool+  -> Ptr TensorList+  -> Ptr TensorList+  -> IO (Ptr Optimizer)+adamwWithParamGroups adamLr adamBetas0 adamBetas1 adamEps adamWeightDecay adamAmsgrad decayParams noDecayParams =+  [C.throwBlock| torch::optim::Optimizer* {+    std::vector<at::Tensor>* decay_params = $(std::vector<at::Tensor>* decayParams);+    std::vector<at::Tensor>* no_decay_params = $(std::vector<at::Tensor>* noDecayParams);+    std::vector<at::Tensor> dp;+    for(int i=0;i<decay_params->size();i++){+      dp.push_back((*decay_params)[i].detach().set_requires_grad(true));+    }+    std::vector<at::Tensor> ndp;+    for(int i=0;i<no_decay_params->size();i++){+      ndp.push_back((*no_decay_params)[i].detach().set_requires_grad(true));+    }+    auto options_decay = torch::optim::AdamWOptions()+      .lr($(double adamLr))+      .betas(std::make_tuple($(double adamBetas0),$(double adamBetas1)))+      .eps($(double adamEps))+      .weight_decay($(double adamWeightDecay))+      .amsgrad($(bool adamAmsgrad));+    auto options_no_decay = torch::optim::AdamWOptions()+      .lr($(double adamLr))+      .betas(std::make_tuple($(double adamBetas0),$(double adamBetas1)))+      .eps($(double adamEps))+      .weight_decay(0.0)+      .amsgrad($(bool adamAmsgrad));+    std::vector<torch::optim::OptimizerParamGroup> param_groups;+    param_groups.emplace_back(torch::optim::OptimizerParamGroup(ndp, std::make_unique<torch::optim::AdamWOptions>(options_no_decay)));+    param_groups.emplace_back(torch::optim::OptimizerParamGroup(dp, std::make_unique<torch::optim::AdamWOptions>(options_decay)));+    torch::optim::AdamW* optimizer = new torch::optim::AdamW(param_groups);+    optimizer->zero_grad();+    return dynamic_cast<torch::optim::Optimizer*>(optimizer);+  }|]++getAllParams :: Ptr Optimizer -> IO (Ptr TensorList)+getAllParams optimizer =+  [C.throwBlock| std::vector<at::Tensor>* {+    auto optimizer = $(torch::optim::Optimizer* optimizer);+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();+    for(auto& group : optimizer->param_groups()){+      for(auto& p : group.params()){+        result->push_back(p);+      }+    }+    return result;+  }|]++stepOnly :: Ptr Optimizer -> IO ()+stepOnly optimizer =+  [C.throwBlock| void {+    $(torch::optim::Optimizer* optimizer)->step();+  }|]++zeroGrad :: Ptr Optimizer -> IO ()+zeroGrad optimizer =+  [C.throwBlock| void {+    $(torch::optim::Optimizer* optimizer)->zero_grad();+  }|]++setParamGrads :: Ptr Optimizer -> Ptr TensorList -> IO ()+setParamGrads optimizer grads =+  [C.throwBlock| void {+    auto optimizer = $(torch::optim::Optimizer* optimizer);+    auto grads = $(std::vector<at::Tensor>* grads);+    int idx = 0;+    for(auto& group : optimizer->param_groups()){+      for(auto& p : group.params()){+        p.mutable_grad() = (*grads)[idx];+        idx++;+      }+    }+  }|]++setLr :: Ptr Optimizer -> CDouble -> IO ()+setLr optimizer newLr =+  [C.throwBlock| void {+    auto optimizer = $(torch::optim::Optimizer* optimizer);+    for(auto& group : optimizer->param_groups()){+      auto& options = static_cast<torch::optim::AdamWOptions&>(group.options());+      options.lr($(double newLr));+    }   }|]