hasktorch 0.2.0.1 → 0.2.1.0
raw patch · 10 files changed
+108/−13 lines, 10 filesdep ~libtorch-ffidep ~libtorch-ffi-helper
Dependency ranges changed: libtorch-ffi, libtorch-ffi-helper
Files
- hasktorch.cabal +3/−3
- src/Torch/Backend.hs +3/−1
- src/Torch/Device.hs +3/−1
- src/Torch/Random.hs +16/−1
- src/Torch/Tensor.hs +19/−3
- src/Torch/TensorOptions.hs +10/−4
- src/Torch/Typed/Auxiliary.hs +4/−0
- src/Torch/Typed/Factories.hs +1/−0
- src/Torch/Typed/Functional.hs +32/−0
- src/Torch/Typed/Tensor.hs +17/−0
hasktorch.cabal view
@@ -1,5 +1,5 @@ name: hasktorch-version: 0.2.0.1+version: 0.2.1.0 synopsis: Haskell bindings to libtorch, supporting both typed and untyped tensors. description: Hasktorch is a library for tensors and neural networks in Haskell. It is an independent open source community project which leverages the core C++ libraries shared by PyTorch. homepage: https://github.com/hasktorch/hasktorch#readme@@ -102,8 +102,8 @@ ghc-options: -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver -fplugin GHC.TypeLits.Extra.Solver -fconstraint-solver-iterations=0 -fplugin GHC.NotExport.Plugin build-depends: async >= 2.2.5 && < 2.3 , base >= 4.7 && < 5- , libtorch-ffi == 2.0.*- , libtorch-ffi-helper == 2.0.*+ , libtorch-ffi == 2.0.1.*+ , libtorch-ffi-helper == 2.0.0.* , finite-typelits >= 0.1 && < 0.3 , ghc-typelits-extra >= 0.4.6 && < 0.5 , ghc-typelits-knownnat >= 0.7.9 && < 0.8
src/Torch/Backend.hs view
@@ -7,12 +7,13 @@ import qualified Torch.Internal.Const as ATen import qualified Torch.Internal.Type as ATen -data Backend = CPU | CUDA | HIP | SparseCPU | SparseCUDA | XLA+data Backend = CPU | CUDA | HIP | SparseCPU | SparseCUDA | XLA | MPS deriving (Eq, Show) instance Castable Backend ATen.Backend where cast CPU f = f ATen.bCPU cast CUDA f = f ATen.bCUDA+ cast MPS f = f ATen.bMPS cast HIP f = f ATen.bHIP cast SparseCPU f = f ATen.bSparseCPU cast SparseCUDA f = f ATen.bSparseCUDA@@ -21,6 +22,7 @@ uncast x f | x == ATen.bCPU = f CPU | x == ATen.bCUDA = f CUDA+ | x == ATen.bMPS = f MPS | x == ATen.bHIP = f HIP | x == ATen.bSparseCPU = f SparseCPU | x == ATen.bSparseCUDA = f SparseCUDA
src/Torch/Device.hs view
@@ -8,7 +8,7 @@ import qualified Torch.Internal.Const as ATen import qualified Torch.Internal.Type as ATen -data DeviceType = CPU | CUDA+data DeviceType = CPU | CUDA | MPS deriving (Eq, Ord, Show) data Device = Device {deviceType :: DeviceType, deviceIndex :: I.Int16}@@ -17,7 +17,9 @@ instance Castable DeviceType ATen.DeviceType where cast CPU f = f ATen.kCPU cast CUDA f = f ATen.kCUDA+ cast MPS f = f ATen.kMPS uncast x f | x == ATen.kCPU = f CPU | x == ATen.kCUDA = f CUDA+ | x == ATen.kMPS = f MPS
src/Torch/Random.hs view
@@ -54,6 +54,11 @@ ATen.generator_set_current_seed genPtr seed genenerator <- newTVarIO (Right genPtr) return $ UnsafeGenerator genenerator+ Device MPS _ -> do+ genPtr <- ATen.newMPSGenerator+ ATen.generator_set_current_seed genPtr seed+ genenerator <- newTVarIO (Right genPtr)+ return $ UnsafeGenerator genenerator type RandomGenFunc = ForeignPtr ATen.IntArray -> ForeignPtr ATen.Generator -> ForeignPtr ATen.TensorOptions -> IO (ForeignPtr ATen.Tensor) @@ -67,7 +72,10 @@ let device = if generatorIsCuda v' then Device {deviceType = CUDA, deviceIndex = fromIntegral $ generatorDevice v'}- else Device {deviceType = CPU, deviceIndex = 0}+ else+ if generatorIsMps v'+ then Device {deviceType = MPS, deviceIndex = 0}+ else Device {deviceType = CPU, deviceIndex = 0} seed = generatorSeed v' writeTVar generator $ seed `seq` deviceType device `seq` deviceIndex device `seq` Left (seed, device) return $ Right v'@@ -80,6 +88,10 @@ gen <- ATen.newCUDAGenerator (fromIntegral idx) ATen.generator_set_current_seed gen seed return gen+ Device MPS _ -> do+ gen <- ATen.newMPSGenerator+ ATen.generator_set_current_seed gen seed+ return gen tensor <- cast3 func size genPtr options nextGenenerator <- newTVarIO (Right genPtr) return (tensor, UnsafeGenerator nextGenenerator)@@ -89,6 +101,9 @@ generatorIsCuda :: ForeignPtr ATen.Generator -> Bool generatorIsCuda gen = unsafePerformIO $ cast1 ATen.generator_is_cuda gen++ generatorIsMps :: ForeignPtr ATen.Generator -> Bool+ generatorIsMps gen = unsafePerformIO $ cast1 ATen.generator_is_mps gen generatorDevice :: ForeignPtr ATen.Generator -> Int generatorDevice gen = unsafePerformIO $ cast1 ATen.generator_get_device gen
src/Torch/Tensor.hs view
@@ -134,11 +134,19 @@ then do isCUDA <- cast1 ATen.tensor_is_cuda t :: IO Bool if isCUDA then cuda <$> cast1 ATen.tensor_get_device t else pure cpu- else pure cpu+ else do+ hasMPS <- cast0 ATen.hasMPS :: IO Bool+ if hasMPS+ then do+ isMPS <- cast1 ATen.tensor_is_mps t :: IO Bool+ if isMPS then pure mps else pure cpu+ else+ pure cpu where cpu = Device {deviceType = CPU, deviceIndex = 0} cuda :: Int -> Device cuda di = Device {deviceType = CUDA, deviceIndex = fromIntegral di}+ mps = Device {deviceType = MPS, deviceIndex = 0} -- | Returns the data type of the input tensor dtype ::@@ -221,7 +229,10 @@ -- | output Tensor _toDevice device' t = unsafePerformIO $ do- hasCUDA <- cast0 ATen.hasCUDA :: IO Bool+ hasDevice <- case deviceType device' of+ CPU -> pure True+ CUDA -> cast0 ATen.hasCUDA+ MPS -> cast0 ATen.hasMPS let device = Torch.Tensor.device t t' <- toDevice'@@ -229,7 +240,7 @@ (deviceType device') (deviceIndex device) (deviceIndex device')- hasCUDA+ hasDevice check (deviceType device') (deviceType $ Torch.Tensor.device t')@@ -241,6 +252,8 @@ toDevice' CUDA CUDA di di' True | di /= di' = getOpts t >>= withDeviceIndex di' >>= to t -- copy from di to di' toDevice' CPU CUDA 0 di' True | di' >= 0 = getOpts t >>= withDeviceIndex di' >>= to t -- copy from cpu:0 to cuda:di' toDevice' CUDA CPU di 0 True | di >= 0 = getOpts t >>= withDeviceType CPU >>= to t -- copy from cuda:di to cpu:0+ toDevice' CPU MPS 0 0 True = getOpts t >>= withDeviceType MPS >>= to t -- copy from cpu:0 to mps:0'+ toDevice' MPS CPU 0 0 True = getOpts t >>= withDeviceType CPU >>= to t -- copy from mps:0 to cpu:0 toDevice' dt dt' di di' _ = error $ "cannot move tensor from \""@@ -364,6 +377,9 @@ toCUDA :: Tensor -> Tensor toCUDA t = unsafePerformIO $ (cast1 ATen.tensor_cuda) t++toMPS :: Tensor -> Tensor+toMPS t = unsafePerformIO $ (cast1 ATen.tensor_mps) t -------------------------------------------------------------------------------- -- Indexing support
src/Torch/TensorOptions.hs view
@@ -35,8 +35,14 @@ withDevice :: Device -> TensorOptions -> TensorOptions withDevice Device {..} opts = unsafePerformIO $ do- hasCUDA <- cast0 ATen.hasCUDA- withDevice' deviceType deviceIndex hasCUDA opts+ case deviceType of+ CPU -> pure opts+ CUDA -> do+ hasCUDA <- cast0 ATen.hasCUDA+ withDevice' deviceType deviceIndex hasCUDA opts+ MPS -> do+ hasMPS <- cast0 ATen.hasMPS+ withDevice' deviceType deviceIndex hasMPS opts where withDeviceType :: DeviceType -> TensorOptions -> IO TensorOptions withDeviceType dt opts = cast2 ATen.tensorOptions_device_D opts dt@@ -44,9 +50,9 @@ withDeviceIndex di opts = cast2 ATen.tensorOptions_device_index_s opts di -- careful, this somehow implies deviceType = CUDA withDevice' :: DeviceType -> Int16 -> Bool -> TensorOptions -> IO TensorOptions- withDevice' CPU 0 False opts = pure opts- withDevice' CPU 0 True opts = pure opts >>= withDeviceType CPU+ withDevice' CPU 0 _ opts = pure opts >>= withDeviceType CPU withDevice' CUDA di True opts | di >= 0 = pure opts >>= withDeviceIndex di+ withDevice' MPS 0 True opts = pure opts >>= withDeviceType MPS withDevice' dt di _ _ = error $ "cannot move tensor to \"" <> show dt <> ":" <> show di <> "\""
src/Torch/Typed/Auxiliary.hs view
@@ -347,6 +347,10 @@ ( DTypeIsFloatingPoint '( 'D.CPU, 0) dtype, DTypeIsNotHalf '( 'D.CPU, 0) dtype )+ StandardFloatingPointDTypeValidation '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype+ ) StandardFloatingPointDTypeValidation '( 'D.CUDA, deviceIndex) dtype = DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype StandardFloatingPointDTypeValidation '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype
src/Torch/Typed/Factories.hs view
@@ -90,6 +90,7 @@ DTypeIsNotHalf '( 'D.CPU, 0) dtype ) RandDTypeIsValid '( 'D.CUDA, _) dtype = ()+ RandDTypeIsValid '( 'D.MPS, 0) dtype = () RandDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype rand ::
src/Torch/Typed/Functional.hs view
@@ -157,6 +157,7 @@ type family SumDTypeIsValid (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where SumDTypeIsValid '( 'D.CPU, 0) dtype = DTypeIsNotHalf '( 'D.CPU, 0) dtype SumDTypeIsValid '( 'D.CUDA, _) dtype = ()+ SumDTypeIsValid '( 'D.MPS, 0) dtype = () SumDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | sumAll@@ -249,6 +250,7 @@ type family MinMaxDTypeIsValid (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where MinMaxDTypeIsValid '( 'D.CPU, 0) dtype = DTypeIsNotHalf '( 'D.CPU, 0) dtype MinMaxDTypeIsValid '( 'D.CUDA, _) dtype = ()+ MinMaxDTypeIsValid '( 'D.MPS, 0) dtype = () MinMaxDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | min@@ -939,6 +941,7 @@ type family InverseShapeIsValid (device :: (D.DeviceType, Nat)) (shape :: [Nat]) :: Constraint where InverseShapeIsValid '( 'D.CPU, 0) _ = () InverseShapeIsValid '( 'D.CUDA, _) shape = AllDimsPositive shape+ InverseShapeIsValid '( 'D.MPS, 0) shape = AllDimsPositive shape type family InverseDTypeIsValid (device :: (D.DeviceType, Nat)) (dtype :: D.DType) :: Constraint where InverseDTypeIsValid '( 'D.CPU, 0) dtype =@@ -949,6 +952,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ InverseDTypeIsValid '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype+ ) InverseDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | inverse@@ -981,6 +988,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ SymeigDTypeIsValid '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype+ ) SymeigDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | symeig@@ -1086,6 +1097,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ EigDTypeIsValid '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype+ ) EigDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype type family ToComplexNumber (dtype :: D.DType) :: D.DType where@@ -1168,6 +1183,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ SVDDTypeIsValid '(D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '(D.MPS, 0) dtype,+ DTypeIsNotHalf '(D.MPS, 0) dtype+ ) SVDDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | Singular Value Decomposition@@ -1233,6 +1252,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ CholeskyDTypeIsValid '(D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '(D.MPS, 0) dtype,+ DTypeIsNotHalf '(D.MPS, 0) dtype+ ) CholeskyDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | cholesky@@ -1343,6 +1366,10 @@ ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype )+ SolveDTypeIsValid '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype+ ) SolveDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | solve@@ -3030,6 +3057,7 @@ DTypeIsNotHalf '(D.CPU, 0) dtype ) DotDTypeIsValid '( 'D.CUDA, deviceIndex) dtype = DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype+ DotDTypeIsValid '( 'D.MPS, 0) dtype = DTypeIsFloatingPoint '( 'D.MPS, 0) dtype DotDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | dot product@@ -4168,6 +4196,10 @@ GeluDTypeIsValid '( 'D.CUDA, deviceIndex) dtype = ( DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype, DTypeIsNotHalf '( 'D.CUDA, deviceIndex) dtype+ )+ GeluDTypeIsValid '( 'D.MPS, 0) dtype =+ ( DTypeIsFloatingPoint '( 'D.MPS, 0) dtype,+ DTypeIsNotHalf '( 'D.MPS, 0) dtype ) GeluDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype
src/Torch/Typed/Tensor.hs view
@@ -140,6 +140,9 @@ instance (KnownNat n) => KnownDevice '( 'D.CUDA, n) where deviceVal = D.Device D.CUDA (natValInt16 @n) +instance (KnownNat n) => KnownDevice '( 'D.MPS, n) where+ deviceVal = D.Device D.MPS (natValInt16 @n)+ type Size = Type -> Type type Shape = [Type -> Type]@@ -207,6 +210,8 @@ type CUDATensor deviceIndex = Tensor '( 'D.CUDA, deviceIndex) +type MPSTensor deviceIndex = Tensor '( 'D.MPS, 0)+ data UnknownShapeTensor device dtype = forall shape. UnknownShapeTensor (Tensor device dtype shape) type family ComputeHaskellType (dtype :: D.DType) :: Type where@@ -310,6 +315,7 @@ Just (SomeNat (Proxy :: Proxy n)) -> case deviceType of D.CPU -> SomeDevice $ Proxy @'( 'D.CPU, n) D.CUDA -> SomeDevice $ Proxy @'( 'D.CUDA, n)+ D.MPS -> SomeDevice $ Proxy @'( 'D.MPS, n) withTensor :: D.Tensor ->@@ -381,6 +387,7 @@ DTypeIsNotHalf '( 'D.CPU, 0) dtype ) BasicArithmeticDTypeIsValid '( 'D.CUDA, _) dtype = ()+ BasicArithmeticDTypeIsValid '( 'D.MPS, 0) dtype = () BasicArithmeticDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype add,@@ -408,6 +415,7 @@ DTypeIsNotHalf '( 'D.CPU, 0) dtype ) ComparisonDTypeIsValid '( 'D.CUDA, _) dtype = ()+ ComparisonDTypeIsValid '( 'D.MPS, 0) dtype = () ComparisonDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype gt,@@ -468,6 +476,7 @@ DTypeIsNotHalf '( 'D.CPU, 0) dtype ) MatMulDTypeIsValid '( 'D.CUDA, deviceIndex) dtype = DTypeIsFloatingPoint '( 'D.CUDA, deviceIndex) dtype+ MatMulDTypeIsValid '( 'D.MPS, 0) dtype = DTypeIsFloatingPoint '( 'D.MPS, 0) dtype MatMulDTypeIsValid '(deviceType, _) dtype = UnsupportedDTypeForDevice deviceType dtype -- | matrix multiplication@@ -635,6 +644,14 @@ Tensor device dtype shape -> CUDATensor 0 dtype shape toCUDA t = UnsafeMkTensor $ D.toCUDA (toDynamic t)++-- | move tensor to the first MPS device+-- TODO: what if this fails?+toMPS ::+ forall device' device shape dtype.+ Tensor device dtype shape ->+ MPSTensor 0 dtype shape+toMPS t = UnsafeMkTensor $ D.toMPS (toDynamic t) -- | move tensor to device -- TODO: what if this fails?