diff --git a/htvm.cabal b/htvm.cabal --- a/htvm.cabal +++ b/htvm.cabal @@ -5,7 +5,7 @@ EDSL for building models with TVM. homepage: https://github.com/grwlf/htvm -version: 0.1.0.0 +version: 0.1.1 license: GPL-3 license-file: LICENSE category: Machine Learning @@ -50,6 +50,7 @@ , HTVM.EDSL.Build , HTVM.EDSL , HTVM.Runtime.FFI + , HTVM.Runtime.TVMData , HTVM.Runtime , HTVM.Prelude , HTVM diff --git a/src/HTVM/EDSL/Build.hs b/src/HTVM/EDSL/Build.hs --- a/src/HTVM/EDSL/Build.hs +++ b/src/HTVM/EDSL/Build.hs @@ -72,7 +72,7 @@ hPutStr stderr err case ec of ExitFailure ec -> do - error $ "stage failed, exit code " <> show ec + error $ "stage: model generator failed, exit code " <> show ec ExitSuccess -> do return (Assembly mod out) @@ -85,7 +85,7 @@ hPutStr stdout out case ec of ExitFailure ec -> do - error $ "compileModel failed, exit code " <> show ec + error $ "compileModel: g++ failed, exit code " <> show ec ExitSuccess -> do return (ModuleLib fp mod) @@ -110,13 +110,13 @@ printFunction :: CompileConfig -> Function -> IO Text -printFunction cc f@(Function te) = do +printFunction cc f@(Function _ te) = do withTmpf "printer" $ \f -> do ProgramBin prg <- compileProgram cc f (printPrinter te) let exec_fp = if isAbsolute prg then prg else "./" <> prg (ec,out,err) <- readProcessWithExitCode exec_fp [] [] case ec of ExitFailure ec -> do - error $ "compileModel failed, exit code " <> show ec + error $ "printFunction: compileProgram failed, exit code " <> show ec ExitSuccess -> return (tpack out) diff --git a/src/HTVM/EDSL/Monad.hs b/src/HTVM/EDSL/Monad.hs --- a/src/HTVM/EDSL/Monad.hs +++ b/src/HTVM/EDSL/Monad.hs @@ -84,11 +84,11 @@ stageTenExpr s = stage <$> runStmtT initStmtCtx s where stage (te,StmtCtx{..}) = sc_expr te -stageFunction :: (Monad m) => StmtT m Function -> m Function -stageFunction fe = Function <$> stageTenExpr (unFunction <$> fe) +stageFunctionT :: (Monad m) => StmtT m Function -> m Function +stageFunctionT fe = stage <$> runStmtT initStmtCtx fe where + stage (Function n te,StmtCtx{..}) = Function n (sc_expr te) -- | Returned module contains all its definitions. --- FIXME: Encode self-contained Modules differently. stageModuleT :: (Monad m) => StmtT m Module -> m Module stageModuleT s = stage <$> runStmtT initStmtCtx s where stage (Module funcs te,StmtCtx{..}) = Module funcs (sc_expr te) @@ -119,33 +119,46 @@ assign :: forall m a . (TensorLike a, Monad m) => a -> StmtT m a assign a = fromTenExpr <$> assignN (toPattern @a) "asgn" (toTenExpr a) -newtype Function = Function { unFunction :: TenExpr } +-- | Function represents TVM expression which is a valid `Module`-function definition +-- Note that Module-functions ate not first-class objects in TVM (TODO: check +-- that fact). +data Function = Function { funcName :: Text, unFunction :: TenExpr } deriving(Read,Show,Eq,Ord) +-- | Module contains a valid module expression and a set of module functions data Module = Module { modFuncs :: [Function] , modExpr :: TenExpr } deriving(Read,Show,Eq,Ord) -data ModuleGenSrc = ModuleGenSrc Module Text +-- | ModuleGenSrc is a C++ sources Module generator +data ModuleGenSrc = ModuleGenSrc { mgen_mod :: Module, mgen_src :: Text } deriving(Show,Read,Eq,Ord) -data ProgramSrc = ProgramSrc Text +-- | Represents C++ sources arbitrary program +data ProgramSrc = ProgramSrc { prog_src :: Text } deriving(Show,Read,Eq,Ord) +-- | Represent path to arbitrary program's binary data ProgramBin = ProgramBin FilePath deriving(Show,Read,Eq,Ord) +-- | Represent path to Module generator binary data ModuleGen = ModuleGen FilePath Module deriving(Show,Read,Eq,Ord) +-- | LLVM Assembly produced by Module generator, along with source Module data Assembly = Assembly Module String deriving(Show,Read,Eq,Ord) +-- | Path to compiled Module along with its source expression data ModuleLib = ModuleLib FilePath Module deriving(Show,Read,Eq,Ord) +-- | Define a module function. Accepts its name @n@, Placeholder definitions +-- @plh@ which become a type of arguments and a lambda function @fbody@ defining +-- the body. List passed to @fbody@ would have same length as @plh@. function :: (Monad m) => Text -> [Placeholder] -> ([Tensor] -> StmtT m Tensor) -> StmtT m Function function n plh fbody = do - Function <$> do + Function <$> pure n <*> do (\x -> assign_ (PFunc (Name n)) x >> pure (TenId (Name n))) =<< do scope $ do plhs <- forM plh $ assignN PTensor "plh" . TenPlh @@ -167,11 +180,22 @@ axis <- freshP "bcomp" assign (Tuple $ batchCompute' se axis tbody) -compute :: (Monad m) => ShapeExpr -> (Expr -> Expr) -> StmtT m Tensor -compute se ebody = do +uniCompute :: (Monad m) => ShapeExpr -> (Expr -> Expr) -> StmtT m Tensor +uniCompute se ebody = do axis <- freshP "comp" assign (Tensor $ flip TenSlice 0 $ batchCompute' se axis ((\x -> [x]) . ebody)) +-- | Specialize computes to different number of dimentsions +class Computable a where + compute :: (Monad m) => ShapeExpr -> (a -> Expr) -> StmtT m Tensor + +-- TODO: assert the number of dimentions in @se@ equals to number of elements in axis tuple +instance Computable (Expr) where compute se f = uniCompute se (\e -> f (e!0)) +instance Computable (Expr,Expr) where compute se f = uniCompute se (\e -> f (e!0,e!1)) +instance Computable (Expr,Expr,Expr) where compute se f = uniCompute se (\e -> f (e!0,e!1,e!2)) +instance Computable (Expr,Expr,Expr,Expr) where compute se f = uniCompute se (\e -> f (e!0,e!1,e!2,e!3)) +instance Computable (Expr,Expr,Expr,Expr,Expr) where compute se f = uniCompute se (\e -> f (e!0,e!1,e!2,e!3,e!4)) + -- | Version of assign where the computation rule is specified for each -- Tensor's item -- compute :: (Monad m) => ShapeExpr -> ([Expr] -> Expr) -> StmtT m Tensor @@ -220,7 +244,7 @@ reduce_axis :: (Monad m) => (DimExpr,DimExpr) -> StmtT m IterVar reduce_axis (a,b) = IterVar . (\(TenId n) -> EId n) <$> assignN PIterVar "reduce_axis" (TenCall TenReduceAxis [TenArg $ TenTuple [TenDim a, TenDim b]]) -infixr 7 ! +infixr 8 ! class Sliceable a b c | a->b, a->c where (!) :: a -> b -> c @@ -346,8 +370,15 @@ sigmoid :: Tensor -> Tensor sigmoid = elemwise1 "sigmoid" +relu :: Tensor -> Tensor +relu = elemwise1 "relu" + split :: Tensor -> [Integer] -> Integer -> Tuple split (Tensor a) indices axis = Tuple $ TenCall TenSplit [TenArg a, IntsArg indices, IntArg axis] + + +differentiate :: Tensor -> [Tensor] -> Tuple +differentiate (Tensor a) ts = Tuple $ TenCall TenDifferentiate [TenArg a, TenArg $ TenTuple [t|(Tensor t)<-ts]] {- ____ _ _ _ diff --git a/src/HTVM/EDSL/Printer.hs b/src/HTVM/EDSL/Printer.hs --- a/src/HTVM/EDSL/Printer.hs +++ b/src/HTVM/EDSL/Printer.hs @@ -108,6 +108,7 @@ TenMatMul -> "topi::matmul" TenElemwise x -> "topi::"<>x TenSplit -> "topi::split" + TenDifferentiate -> "htvm_differentiate" printLayout :: Layout -> Text printLayout l = @@ -193,6 +194,7 @@ line $ "#include " line $ "#include " line $ "#include " + line $ "#include " line $ "" line $ "static inline tvm::Array \ @@ -216,6 +218,7 @@ line $ "tvm::IterVar htvm_axis_id(tvm::Tensor t, int i) { return t->op->root_iter_vars()[i]; }" line $ "tvm::Array htvm_shape(tvm::Tensor t) { return t->shape; }" line $ "tvm::Array htvm_shape(tvm::Array t) { return t; }" + line $ "tvm::Array htvm_differentiate(tvm::Tensor t, tvm::Array a){ return tvm::ir::Differentiate(t,a)->result; }" line $ "" diff --git a/src/HTVM/EDSL/Types.hs b/src/HTVM/EDSL/Types.hs --- a/src/HTVM/EDSL/Types.hs +++ b/src/HTVM/EDSL/Types.hs @@ -5,10 +5,11 @@ import Data.Monoid import Data.Text(Text) --- | Name is the string convertable to valid C/C++ identifier +-- | Name represents valid C/C++ identifier newtype Name = Name { n_get :: Text } deriving(Show,Read,Ord,Eq,Semigroup,Monoid) +-- | Const encodes valid C/C++ constants data Const = CInt Integer | CFloat32 Float @@ -32,11 +33,6 @@ signum = error "signum is undefined for DimExpr" fromInteger = DimConst --- -- | Axis represents iterator running through the range supplied. Equivalent of --- -- `tvm::IterVar`. --- data Axis = Axis Name (DimExpr,DimExpr) --- deriving(Show,Read,Ord,Eq) - -- | Shape expressions represents the shape of a tensor, i.e. the number and -- size of its dimentions. Rough equivalent of `tvm::Array`. data ShapeExpr = @@ -49,44 +45,20 @@ -- ^ Concatenation on shapes deriving(Show,Read,Ord,Eq) --- | Return the number of dimentions of ShapeExpr which is always known at compile time. --- TODO: Move to `Eval.hs` as a generic algorithm --- shapeDim :: ShapeExpr -> Integer --- shapeDim (ShapeTen ndim _) = ndim --- shapeDim (ShapeVector _) = 1 --- shapeDim (ShapeScalar) = 0 --- shapeDim (ShapeSum se1 se2) = shapeDim se1 + shapeDim se2 - instance Semigroup ShapeExpr where (<>) a b = ShapeSum a b -shape :: [DimExpr] -> ShapeExpr -shape des = undefined - --- | Convert ShapeExpr in flattern form, where each list itme represents a --- dimention, either of known size or unknown at compile time. Empty list --- represents a shape of scalar. --- FIXME: This function is impossible --- shapeFlattern :: ShapeExpr -> [Either DimExpr Integer] --- shapeFlattern sh = --- case sh of --- ShapeId 1 n -> [Left n] --- ShapeId x n -> error "shapeFlattern: don't know how to represent multidimentional shape variables" --- ShapeVector x -> [Right x] --- ShapeScalar -> [] --- ShapeSum a b -> shapeFlattern a <> shapeFlattern b - +-- | A registry of expression-level function names data ExprFuncName = ExprOp Text | ExprSum | ESigmoid deriving(Show,Read,Ord,Eq) --- | Scalar expressions +-- | Scalar expressions, equivalent of `tvm::Expr` data Expr = EConst Const -- ^ A constant | EId Name -- ^ A variable - -- | EShape ShapeExpr -- ^ A shape expression | EShapeSlice ShapeExpr Integer -- ^ Access a certain dimention of ShapeExpr | ETenSlice TenExpr [Expr] -- ^ Accessing an individual element of a tensor @@ -104,25 +76,18 @@ signum = error "signum is undefined" fromInteger = EConst . CInt +-- | Representation of tvm type codes data Type = TypeFloat32 | TypeInt32 | TypeTensor Type ShapeExpr deriving(Show,Read,Ord,Eq) +float32 :: Type float32 = TypeFloat32 --- | Common arguments to various functions -data Args = Args { - a_name :: Maybe Name - , a_shape :: Maybe ShapeExpr - , a_type :: Maybe Type - } deriving(Show,Read,Ord,Eq) - -nullArgs :: Args -nullArgs = Args Nothing Nothing Nothing - --- | Pattern is a left-hand-side of assignments +-- | Pattern represents left-hand-side of C/C++ assignments +-- -- FIXME: Separate type codes from Name binding data Pattern = PTensor Name -- ^ Tensor @@ -137,7 +102,7 @@ | PStage Name deriving(Show,Read,Ord,Eq) --- | List of valid Tensor-Expression level function names +-- | A registry of tensor-level function names data TenFuncName = TenOp Text | TenReduceAxis @@ -149,9 +114,10 @@ | TenMatMul | TenElemwise Text | TenSplit + | TenDifferentiate deriving(Show,Read,Ord,Eq) --- | `TenCall` receive arguments of the following kinds +-- | Kinds of arguments received by `TenCall` data TenArg = TenArg TenExpr -- ^ Ordinary argument, another `TenExpr` | StrArg Text -- ^ String argument @@ -172,14 +138,17 @@ -- * We don't keep Type as a part of TenExpr since in theory we shouldn't need -- it (assuming the typechecker is present) data TenExpr = - TenId Name + TenId Name -- ^ Identifier | TenPlh Placeholder -- ^ Placeholder is a disting kind of TenExpr because it -- refers `Type` and `ShapeExpr` which are not `TenExpr` + -- + -- FIXME: Replace `TenPlh` and `TenDef` with a + -- function representation TenFun | TenTuple [TenExpr] | TenSlice TenExpr Integer -- ^ Slice `TenTuple` - | TenDim DimExpr - | TenShape ShapeExpr + | TenDim DimExpr -- ^ Dimention expression + | TenShape ShapeExpr -- ^ Shape expression | TenExpr Expr -- ^ We need TenExpr to encode `reduce_axis` results. It returns -- sliceable expressions @@ -196,14 +165,6 @@ deriving(Show,Read,Ord,Eq) +-- | Placeholder collects information about entry or exit points of TVM programs type Placeholder = (Text,Type,ShapeExpr) - --- pls_name :: Placeholder -> Name --- pls_name (nm,_,_) = nm - --- data Axis = Axis {aExpr :: TenExpr} --- deriving(Read,Show,Eq,Ord) - - - diff --git a/src/HTVM/Runtime.hs b/src/HTVM/Runtime.hs --- a/src/HTVM/Runtime.hs +++ b/src/HTVM/Runtime.hs @@ -1,5 +1,6 @@ module HTVM.Runtime ( module HTVM.Runtime.FFI + , module HTVM.Runtime.TVMData , module HTVM.Runtime , module Data.Int , module Data.Word @@ -9,3 +10,4 @@ import Data.Word (Word8,Word16,Word32,Word64) import HTVM.Runtime.FFI +import HTVM.Runtime.TVMData diff --git a/src/HTVM/Runtime/FFI.chs b/src/HTVM/Runtime/FFI.chs --- a/src/HTVM/Runtime/FFI.chs +++ b/src/HTVM/Runtime/FFI.chs @@ -1,39 +1,22 @@ --- | Module defines wrappers for DLPack messages which are used by TVM to pass --- to/from models +-- | DLPack message wrappers to pass data to/from TVM models --- {-# OPTIONS_GHC -fwarn-unused-imports #-} --- {-# OPTIONS_GHC -fwarn-missing-signatures #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FunctionalDependencies #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE NondecreasingIndentation #-} -{-# LANGUAGE FlexibleContexts #-} - module HTVM.Runtime.FFI where import qualified Data.Array as Array import Control.Exception (Exception, throwIO) -import Control.Arrow ((***)) import Control.Monad (forM_) -import Data.Array (Array(..)) -import Data.ByteString (ByteString,pack) import Data.Word (Word8,Word16,Word32,Word64) import Data.Int (Int8,Int16,Int32,Int64) -import Data.Bits (FiniteBits(..),(.&.),shiftR) -import Data.Tuple (swap) +import Data.Bits ((.&.),shiftR) import Data.Text (Text) -import Data.List (nub) import Foreign (ForeignPtr, newForeignPtr, Ptr, Storable(..), alloca, allocaArray, peek, plusPtr, poke, peekArray, pokeArray, castPtr, advancePtr, malloc, mallocArray, FunPtr(..), free, - withForeignPtr, nullPtr) -import Foreign.C.Types (CInt, CLong) + withForeignPtr, nullPtr, newForeignPtr_) +import Foreign.C.Types (CInt, CLong, CSize) import Foreign.C.String (CString, withCString, peekCAString) import System.IO.Unsafe (unsafePerformIO) @@ -49,6 +32,7 @@ | TVMFunCallFailed Int String | TVMFunCallBadType Int | TVMCopyFailed Int String + | PokeShapeMismatch [Integer] [Integer] deriving(Show,Read,Ord,Eq) instance Exception TVMError @@ -59,7 +43,6 @@ {# enum DLDataTypeCode as TVMDataTypeCode {upcaseFirstLetter} deriving(Eq) #} {# enum DLDeviceType as TVMDeviceType {upcaseFirstLetter} deriving(Eq) #} - {# enum TVMDeviceExtType {upcaseFirstLetter} deriving(Eq) #} {# enum TVMTypeCode {upcaseFirstLetter} deriving(Eq) #} @@ -75,6 +58,7 @@ -- | Representation of device identifiers type TVMDeviceId = Int +-- | TODO: document data TVMContext instance Storable TVMContext where @@ -83,7 +67,7 @@ peek = error "peek undefined" poke = error "poke undefined" --- | Tensor representation, see `DLTensor` +-- | Representation of `DLTensor` C structure data TVMTensor_Repr instance Storable TVMTensor_Repr where @@ -92,15 +76,19 @@ peek = error "peek undefined" poke = error "poke undefined" --- | Alias for `TVMArrayHandle` +-- | Alias for C type `TVMArrayHandle`, which is internally defined as a +-- pointer to DLTensor. type TVMArrayHandle = Ptr TVMTensor_Repr --- | Alias for pointer to `TVMArray` aka `DLTensor`. +-- | Main runtime representation of Tensors in TVM. Tensors contain multy- +-- dimentional arrays of numbers. Their data is stored either in main CPU memory +-- or on the computing device, e.g. in the memory of GPU card. type TVMTensor = ForeignPtr TVMTensor_Repr -- | Alias for `TVMStreamHandle`. Not supported via this FFI currently. type TVMStreamHandle = Ptr () +-- | TVMValue represents function argument, accepted by TVM functions. data TVMValue instance Storable TVMValue where @@ -109,12 +97,37 @@ peek = error "peek undefined" poke = error "poke undefined" +-- | Representation of ModuleHandle +data TVMModule_Repr + +instance Storable TVMModule_Repr where + sizeOf _ = {# sizeof TVMModuleHandle #} + alignment _ = {# alignof TVMModuleHandle #} + peek = error "peek is undefined for TVMModuleHandle" + poke = error "poke is undefined for TVMModuleHandle" + -- | Alias for void* used as Module handle -type TVMModule = Ptr () +type TVMModuleHandle = Ptr TVMModule_Repr +-- | Foreign pointer to Module handle +type TVMModule = ForeignPtr TVMModule_Repr + +-- | Representation of FunctionHandle +data TVMFunction_Repr + +instance Storable TVMFunction_Repr where + sizeOf _ = {# sizeof TVMFunctionHandle #} + alignment _ = {# alignof TVMFunctionHandle #} + peek = error "peek is undefined for TVMFunction_Repr" + poke = error "poke is undefined for TVMFunction_Repr" + -- | Alias for void* used as Function handle -type TVMFunction = Ptr () +type TVMFunctionHandle = Ptr TVMFunction_Repr +-- | Foreign pointer to function handle +type TVMFunction = ForeignPtr TVMFunction_Repr + + setTensor :: TVMTensor -> Ptr TVMValue -> Ptr TVMTypeCode -> IO () setTensor ft pv pc = do withForeignPtr ft $ \pt -> do @@ -127,15 +140,15 @@ foreign import ccall unsafe "c_runtime_api.h TVMArrayAlloc" tvmArrayAlloc :: Ptr TVMShapeIndex - -- shape - -> CInt -- ndim, - -> CInt -- dtype_code, - -> CInt -- dtype_bits, - -> CInt -- dtype_lanes, - -> CInt -- device_type, - -> CInt -- device_id, + -- ^ shape + -> CInt -- ^ ndim, + -> CInt -- ^ dtype_code, + -> CInt -- ^ dtype_bits, + -> CInt -- ^ dtype_lanes, + -> CInt -- ^ device_type, + -> CInt -- ^ device_id, -> Ptr TVMArrayHandle - -- DLTensor* out + -- ^ DLTensor* out -> IO CInt foreign import ccall unsafe "c_runtime_api.h TVMArrayFree" @@ -144,113 +157,54 @@ foreign import ccall unsafe "c_runtime_api.h &TVMArrayFree" tvmArrayFree_ :: FunPtr (TVMArrayHandle -> IO ()) - foreign import ccall unsafe "c_runtime_api.h TVMModLoadFromFile" - tvmModLoadFromFile :: CString -> CString -> Ptr TVMModule -> IO CInt + tvmModLoadFromFile :: CString -> CString -> Ptr TVMModuleHandle -> IO CInt foreign import ccall unsafe "c_runtime_api.h TVMModFree" - tvmModFree :: TVMModule -> IO CInt + tvmModFree :: TVMModuleHandle -> IO CInt +foreign import ccall unsafe "c_runtime_api.h &TVMModFree" + tvmModFree_ :: FunPtr (TVMModuleHandle -> IO ()) + foreign import ccall unsafe "c_runtime_api.h TVMModGetFunction" - tvmModGetFunction :: TVMModule -> CString -> CInt -> Ptr TVMFunction -> IO CInt + tvmModGetFunction :: TVMModuleHandle -> CString -> CInt -> Ptr TVMFunctionHandle -> IO CInt foreign import ccall unsafe "c_runtime_api.h TVMFuncFree" - tvmFuncFree :: TVMFunction -> IO CInt + tvmFuncFree :: TVMFunctionHandle -> IO CInt +foreign import ccall unsafe "c_runtime_api.h &TVMFuncFree" + tvmFuncFree_ :: FunPtr (TVMFunctionHandle -> IO ()) + foreign import ccall unsafe "c_runtime_api.h TVMGetLastError" tvmGetLastError :: IO CString foreign import ccall unsafe "c_runtime_api.h TVMFuncCall" - tvmFuncCall :: TVMFunction -> Ptr TVMValue -> Ptr TVMTypeCode -> CInt -> Ptr TVMValue -> Ptr TVMTypeCode -> IO CInt + tvmFuncCall :: TVMFunctionHandle -> Ptr TVMValue -> Ptr TVMTypeCode -> CInt -> Ptr TVMValue -> Ptr TVMTypeCode -> IO CInt foreign import ccall unsafe "c_runtime_api.h TVMArrayCopyFromTo" tvmArrayCopyFromTo :: TVMArrayHandle -> TVMArrayHandle -> TVMStreamHandle -> IO CInt -class TVMIndex i where - tvmList :: i -> [Integer] - -instance TVMIndex Integer where tvmList a = [a] -instance TVMIndex (Integer,Integer) where tvmList (a,b) = [a,b] -instance TVMIndex (Integer,Integer,Integer) where tvmList (a,b,c) = [a,b,c] -instance TVMIndex (Integer,Integer,Integer,Integer) where tvmList (a,b,c,d) = [a,b,c,d] - -tvmIndexDims :: (TVMIndex i) => i -> Integer -tvmIndexDims = ilength . tvmList - -class TVMElemType e where - tvmTypeCode :: TVMDataTypeCode - tvmTypeBits :: Integer - -- | Make a parameter of type - tvmTypeLanes :: Integer - -instance TVMElemType Int32 where tvmTypeCode = KDLInt; tvmTypeBits = 32; tvmTypeLanes = 1 -instance TVMElemType Word32 where tvmTypeCode = KDLUInt; tvmTypeBits = 32; tvmTypeLanes = 1 -instance TVMElemType Float where tvmTypeCode = KDLFloat; tvmTypeBits = 32; tvmTypeLanes = 1 -instance TVMElemType Int64 where tvmTypeCode = KDLUInt; tvmTypeBits = 64; tvmTypeLanes = 1 -instance TVMElemType Word64 where tvmTypeCode = KDLUInt; tvmTypeBits = 64; tvmTypeLanes = 1 -instance TVMElemType Double where tvmTypeCode = KDLFloat; tvmTypeBits = 64; tvmTypeLanes = 1 - --- | Data source. @d@ is type of data, @i@ is a type of index, @e@ is a type of element -class (TVMIndex i, TVMElemType e) => TVMData d i e | d -> i, d -> e where - tvmIShape :: d -> [Integer] - tvmIndex :: d -> i -> IO e - tvmPeek :: [Integer] -> Ptr e -> IO d - tvmPoke :: d -> Ptr e -> IO () - -- ^ Write the contents of data to dense memory area. - -- TODO: figure out the alignment restirctions. - -instance (Storable e, Array.Ix i, TVMIndex i, TVMElemType e) => TVMData (Array i e) i e where - tvmIShape = map (uncurry (-)) . uncurry zip . (tvmList *** tvmList) . Array.bounds - tvmIndex d i = pure $ d Array.! i - tvmPoke d ptr = pokeArray ptr (Array.elems d) - tvmPeek shape ptr = error "peek is undefined for Arrays" - -tvmIShape1 d = [ilength d] -tvmIndex1 l i = pure $ l !! (fromInteger i) -tvmPoke1 d ptr = pokeArray ptr d -tvmPeek1 [x] ptr = peekArray (fromInteger x) ptr -tvmPeek1 _ ptr = error "tvmPeek1 should be called with single-element shape" -instance TVMData [Int32] Integer Int32 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 -instance TVMData [Word32] Integer Word32 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 -instance TVMData [Float] Integer Float where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 -instance TVMData [Int64] Integer Int64 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 -instance TVMData [Word64] Integer Word64 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 -instance TVMData [Double] Integer Double where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 - -tvmIShape2 [] = [0,0] -tvmIShape2 d = [ilength d, ilength (head d)] -tvmIndex2 l (r,c) = pure $ l !! (fromInteger r) !! (fromInteger c) -tvmPoke2 d ptr - | length d == 0 = pokeArray ptr (concat d) - | length (nub (map length d)) == 1 = pokeArray ptr (concat d) - | otherwise = error "All elements should have the same length" -tvmPeek2 [0,0] ptr = pure [] -tvmPeek2 [x,y] ptr = group y <$> peekArray (fromInteger $ x*y) ptr where - group :: Integer -> [a] -> [[a]] - group _ [] = [] - group n l - | n > 0 = (take (fromInteger n) l) : (group n (drop (fromInteger n) l)) - | otherwise = error "Negative n" -tvmPeek2 x _ = error "tvmPeek2 should be called with 2-element shape" -instance TVMData [[Int32]] (Integer,Integer) Int32 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 -instance TVMData [[Word32]] (Integer,Integer) Word32 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 -instance TVMData [[Float]] (Integer,Integer) Float where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 -instance TVMData [[Int64]] (Integer,Integer) Int64 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 -instance TVMData [[Word64]] (Integer,Integer) Word64 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 -instance TVMData [[Double]] (Integer,Integer) Double where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +foreign import ccall unsafe "c_runtime_api.h TVMArrayCopyToBytes" + tvmArrayCopyToBytes :: TVMArrayHandle -> Ptr Word8 -> CSize -> IO CInt -tvmDataShape :: (TVMData d i e) => d -> [Integer] -tvmDataShape = tvmIShape +foreign import ccall unsafe "c_runtime_api.h TVMArrayCopyFromBytes" + tvmArrayCopyFromBytes :: TVMArrayHandle -> Ptr Word8 -> CSize -> IO CInt -tvmDataNDim :: (TVMData d i e) => d -> Integer -tvmDataNDim = ilength . tvmDataShape +foreign import ccall unsafe "c_runtime_api.h TVMArrayDataSize" + tvmArrayDataSize :: TVMArrayHandle -> Ptr CSize -> IO CInt +{- FIXME: check data size compatibility -} toCInt :: (Integral x) => x -> CInt toCInt = fromInteger . toInteger +{- FIXME: check data size compatibility -} fromCInt :: (Integral x) => CInt -> x fromCInt = fromInteger . toInteger +{- FIXME: check data size compatibility -} +fromCSize :: (Integral x) => CSize -> x +fromCSize = fromInteger . toInteger + tensorDevice :: TVMTensor -> TVMDeviceType tensorDevice ft = unsafePerformIO $ do withForeignPtr ft $ \pt -> do @@ -274,72 +228,18 @@ map toInteger <$> do peekArray (fromInteger $ tensorNDim ft) =<< {# get DLTensor->shape #} pt --- | Create new empty TVMTensor object. This object will be managed by Haskell --- runtime which would call free when it decides to do so. --- --- FIXME: non-CPU devices will not work, see FIXME in `pokeTensor` -newEmptyTensor :: forall e . (TVMElemType e) - => [Integer] -- ^ Shape - -> TVMDeviceType -- ^ Device type (CPU|GPU|etc) - -> TVMDeviceId -- ^ Device ID - -> IO TVMTensor -newEmptyTensor shape dt did = - let - ndim = length shape - in do - alloca $ \pt -> do - allocaArray ndim $ \pshape -> do - pokeArray pshape (map (fromInteger . toInteger) shape) - r <- - tvmArrayAlloc - pshape - (fromInteger $ toInteger $ ndim) - (toCInt $ fromEnum $ tvmTypeCode @e) - (toCInt $ tvmTypeBits @e) - (toCInt $ tvmTypeLanes @e) - (toCInt $ fromEnum dt) - (toCInt $ did) pt - case r of - 0 -> peek pt >>= newForeignPtr tvmArrayFree_ - e -> throwIO (TVMAllocFailed (fromCInt e)) - --- | Allocate new Tensor object -newTensor :: forall d i e . (TVMData d i e) - => d -- ^ TvmData tensor-like object - -> TVMDeviceType -- ^ Device type - -> TVMDeviceId -- ^ Device ID - -> IO TVMTensor -newTensor d dt did = do - ft <- newEmptyTensor @e (map fromInteger $ tvmDataShape d) dt did - withForeignPtr ft $ \pt -> do - pdata <- {# get DLTensor->data #} pt - tvmPoke d (castPtr pdata) - return ft - - -peekTensor :: forall d i e b . (TVMData d i e) - => TVMTensor -> IO d -peekTensor ft = do - withForeignPtr ft $ \pt -> do - case tensorDevice ft of - KDLCPU -> do - pdata <- {# get DLTensor->data #} pt - tvmPeek (tensorShape ft) (castPtr pdata) - x -> do - fail "Not implemented" - -pokeTensor :: forall d i e b . (TVMData d i e) - => TVMTensor -> d -> IO () -pokeTensor ft d = do - withForeignPtr ft $ \pt -> do - case tensorDevice ft of - KDLCPU -> do - pdata <- {# get DLTensor->data #} pt - {- FIXME: Use CopyFromBytes for non-CPU devices -} - tvmPoke d (castPtr pdata) - x -> do - fail "Not implemented" +-- | Access device-specific raw tensor data. In case of CPU Tensor this is a +-- pointer to raw data array +unsafeTensorData :: TVMTensor -> Ptr () +unsafeTensorData p = unsafePerformIO $ withForeignPtr p {# get DLTensor->data #} +-- | Return number of bytes required to store tensor's data +tensorSize :: TVMTensor -> CSize +tensorSize p = unsafePerformIO $ do + alloca $ \psz -> do + withForeignPtr p $ \ht -> do + tvmArrayDataSize ht psz + peek psz tensorCopy :: TVMTensor -> TVMTensor -> IO () tensorCopy dst src = do @@ -356,7 +256,20 @@ getLastError :: IO String getLastError = peekCAString =<< tvmGetLastError +-- | Load module named @modname@. Module will be freed when Haskell runtime +-- decide so. +loadModule :: FilePath -> IO TVMModule +loadModule modname = do + alloca $ \pmod -> do + withCString modname $ \cmodname -> do + withCString "so" $ \so -> do + r <- tvmModLoadFromFile cmodname so pmod + case r of + 0 -> peek pmod >>= newForeignPtr tvmModFree_ + err -> throwIO =<< (TVMModLoadFailed <$> pure (fromCInt err) <*> getLastError) + -- | Load module from dynamic library @modname@ and process it with a callback @func@ +-- Module will be freed on return from @func@ withModule :: FilePath -> (TVMModule -> IO b) -> IO b withModule modname func = alloca $ \pmod -> do @@ -366,23 +279,36 @@ case r of 0 -> do m <- peek pmod - b <- func m + b <- func =<< (newForeignPtr_ m) tvmModFree m return b err -> do - str <- getLastError - throwIO (TVMModLoadFailed (fromInteger $ toInteger err) str) + throwIO =<< (TVMModLoadFailed <$> pure (fromCInt err) <*> getLastError) +-- | Load function named @funcname@ from module @mod@. Function will be +-- freed when Haskell runtime decide so. +loadFunction :: Text -> TVMModule -> IO TVMFunction +loadFunction funcname mod = do + alloca $ \pfunc -> do + withCString (tunpack funcname) $ \cfuncname -> do + withForeignPtr mod $ \hmod -> do + r <- tvmModGetFunction hmod cfuncname 0 pfunc + case r of + 0 -> peek pfunc >>= newForeignPtr tvmFuncFree_ + err -> throwIO =<< (TVMFuncLoadFailed <$> pure (fromCInt err) <*> getLastError) + -- | Load the function named @funcname@ from module @mod@, use it in callback @func@ +-- Function will be freed on return from @func@ withFunction :: Text -> TVMModule -> (TVMFunction -> IO b) -> IO b withFunction funcname mod func = alloca $ \pfunc -> do withCString (tunpack funcname) $ \cfuncname -> do - r <- tvmModGetFunction mod cfuncname 0 pfunc + withForeignPtr mod $ \hmod -> do + r <- tvmModGetFunction hmod cfuncname 0 pfunc case r of 0 -> do f <- peek pfunc - b <- func f + b <- func =<< (newForeignPtr_ f) tvmFuncFree f return b err -> do @@ -402,11 +328,12 @@ alloca $ \pvretcode -> do allocaArray nargs $ \pvargs -> do allocaArray nargs $ \pvargcodes -> do + withForeignPtr fun $ \hfun -> do forM_ ((args<>[ret])`zip`[0..nargs-1]) $ \(farg,off) -> do case off < length args of True -> setTensor farg (advancePtr pvargs off) (advancePtr pvargcodes off) False -> setTensor ret (advancePtr pvargs off) (advancePtr pvargcodes off) - r <- tvmFuncCall fun pvargs pvargcodes clen pvret pvretcode + r <- tvmFuncCall hfun pvargs pvargcodes clen pvret pvretcode case r of 0 -> do return () diff --git a/src/HTVM/Runtime/TVMData.hs b/src/HTVM/Runtime/TVMData.hs new file mode 100644 --- /dev/null +++ b/src/HTVM/Runtime/TVMData.hs @@ -0,0 +1,185 @@ +-- | DLPack message wrappers to pass data to/from TVM models + +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE NondecreasingIndentation #-} +{-# LANGUAGE FlexibleContexts #-} + +module HTVM.Runtime.TVMData where + +import qualified Data.Array as Array + +import Control.Exception (throwIO) +import Control.Monad (when) +import Control.Arrow ((***)) +import Data.Array (Array(..)) +import Data.Word (Word8,Word16,Word32,Word64) +import Data.Int (Int8,Int16,Int32,Int64) +import Data.Text (Text) +import Data.List (nub) +import Foreign (newForeignPtr, Ptr, Storable(..), alloca, + allocaArray, peek, poke, peekArray, pokeArray, + castPtr, advancePtr, malloc, mallocArray, + withForeignPtr) + +import HTVM.Prelude +import HTVM.Runtime.FFI + + +class TVMIndex i where + tvmList :: i -> [Integer] + +instance TVMIndex Integer where tvmList a = [a] +instance TVMIndex (Integer,Integer) where tvmList (a,b) = [a,b] +instance TVMIndex (Integer,Integer,Integer) where tvmList (a,b,c) = [a,b,c] +instance TVMIndex (Integer,Integer,Integer,Integer) where tvmList (a,b,c,d) = [a,b,c,d] + +tvmIndexDims :: (TVMIndex i) => i -> Integer +tvmIndexDims = ilength . tvmList + +class TVMElemType e where + tvmTypeCode :: TVMDataTypeCode + tvmTypeBits :: Integer + -- | Make a parameter of type + tvmTypeLanes :: Integer + +instance TVMElemType Int32 where tvmTypeCode = KDLInt; tvmTypeBits = 32; tvmTypeLanes = 1 +instance TVMElemType Word32 where tvmTypeCode = KDLUInt; tvmTypeBits = 32; tvmTypeLanes = 1 +instance TVMElemType Float where tvmTypeCode = KDLFloat; tvmTypeBits = 32; tvmTypeLanes = 1 +instance TVMElemType Int64 where tvmTypeCode = KDLUInt; tvmTypeBits = 64; tvmTypeLanes = 1 +instance TVMElemType Word64 where tvmTypeCode = KDLUInt; tvmTypeBits = 64; tvmTypeLanes = 1 +instance TVMElemType Double where tvmTypeCode = KDLFloat; tvmTypeBits = 64; tvmTypeLanes = 1 + +-- | Data bundle accepted by TVM model. @d@ is type of data, @i@ is a type of +-- index, @e@ is the type of element +class (TVMIndex i, TVMElemType e) => TVMData d i e | d -> i, d -> e where + tvmIShape :: d -> [Integer] + tvmIndex :: d -> i -> IO e + tvmPeek :: [Integer] -> Ptr e -> IO d + tvmPoke :: d -> Ptr e -> IO () + -- ^ Write the contents of data to dense memory area. + -- TODO: figure out the alignment restirctions. + +instance (Storable e, Array.Ix i, TVMIndex i, TVMElemType e) => TVMData (Array i e) i e where + tvmIShape = map (uncurry (-)) . uncurry zip . (tvmList *** tvmList) . Array.bounds + tvmIndex d i = pure $ d Array.! i + tvmPoke d ptr = pokeArray ptr (Array.elems d) + tvmPeek shape ptr = error "FIXME: tvmPeek is not implemented for Data.Array" + +tvmIShape1 d = [ilength d] +tvmIndex1 l i = pure $ l !! (fromInteger i) +tvmPoke1 d ptr = pokeArray ptr d +tvmPeek1 [x] ptr = peekArray (fromInteger x) ptr +tvmPeek1 _ ptr = error "tvmPeek1 should be called with single-element shape" +instance TVMData [Int32] Integer Int32 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 +instance TVMData [Word32] Integer Word32 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 +instance TVMData [Float] Integer Float where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 +instance TVMData [Int64] Integer Int64 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 +instance TVMData [Word64] Integer Word64 where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 +instance TVMData [Double] Integer Double where tvmIShape = tvmIShape1 ; tvmIndex = tvmIndex1; tvmPoke = tvmPoke1; tvmPeek = tvmPeek1 + +tvmIShape2 [] = [0,0] +tvmIShape2 d = [ilength d, ilength (head d)] +tvmIndex2 l (r,c) = pure $ l !! (fromInteger r) !! (fromInteger c) +tvmPoke2 d ptr + | length d == 0 = pokeArray ptr (concat d) + | length (nub (map length d)) == 1 = pokeArray ptr (concat d) + | otherwise = error "All elements should have the same length" +tvmPeek2 [0,0] ptr = pure [] +tvmPeek2 [x,y] ptr = group y <$> peekArray (fromInteger $ x*y) ptr where + group :: Integer -> [a] -> [[a]] + group _ [] = [] + group n l + | n > 0 = (take (fromInteger n) l) : (group n (drop (fromInteger n) l)) + | otherwise = error "Negative n" +tvmPeek2 x _ = error "tvmPeek2 should be called with 2-element shape" +instance TVMData [[Int32]] (Integer,Integer) Int32 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +instance TVMData [[Word32]] (Integer,Integer) Word32 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +instance TVMData [[Float]] (Integer,Integer) Float where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +instance TVMData [[Int64]] (Integer,Integer) Int64 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +instance TVMData [[Word64]] (Integer,Integer) Word64 where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 +instance TVMData [[Double]] (Integer,Integer) Double where tvmIShape = tvmIShape2 ; tvmIndex = tvmIndex2; tvmPoke = tvmPoke2; tvmPeek = tvmPeek2 + +tvmDataShape :: (TVMData d i e) => d -> [Integer] +tvmDataShape = tvmIShape + +tvmDataNDim :: (TVMData d i e) => d -> Integer +tvmDataNDim = ilength . tvmDataShape + + + +-- | Create new empty TVMTensor object. This object will be managed by Haskell +-- runtime which would call free when it decides to do so. +-- +-- FIXME: non-CPU devices will not work, see FIXME in `pokeTensor` +newEmptyTensor :: forall e . (TVMElemType e) + => [Integer] -- ^ Shape + -> TVMDeviceType -- ^ Device type (CPU|GPU|etc) + -> TVMDeviceId -- ^ Device ID + -> IO TVMTensor +newEmptyTensor shape dt did = + let + ndim = length shape + in do + alloca $ \pt -> do + allocaArray ndim $ \pshape -> do + pokeArray pshape (map (fromInteger . toInteger) shape) + r <- + tvmArrayAlloc + pshape + (fromInteger $ toInteger $ ndim) + (toCInt $ fromEnum $ tvmTypeCode @e) + (toCInt $ tvmTypeBits @e) + (toCInt $ tvmTypeLanes @e) + (toCInt $ fromEnum dt) + (toCInt $ did) pt + case r of + 0 -> peek pt >>= newForeignPtr tvmArrayFree_ + e -> throwIO (TVMAllocFailed (fromCInt e)) + +-- | Allocate new Tensor object +newTensor :: forall d i e . (TVMData d i e) + => d -- ^ TvmData tensor-like object + -> TVMDeviceType -- ^ Device type + -> TVMDeviceId -- ^ Device ID + -> IO TVMTensor +newTensor d dt did = do + ft <- newEmptyTensor @e (map fromInteger $ tvmDataShape d) dt did + withForeignPtr ft $ \pt -> do + tvmPoke d (castPtr (unsafeTensorData ft)) + return ft + +-- | Transfer data from TVMTensor to TVMData instance +peekTensor :: forall d i e . (TVMData d i e) + => TVMTensor -> IO d +peekTensor ft = do + case tensorDevice ft of + KDLCPU -> do + tvmPeek (tensorShape ft) (castPtr (unsafeTensorData ft)) + x -> do + allocaArray (fromCSize $ tensorSize ft) $ \parr -> do + withForeignPtr ft $ \pt -> do + _ <- tvmArrayCopyToBytes pt parr (tensorSize ft) + tvmPeek (tensorShape ft) (castPtr parr) + +-- | Transfer data from TVMData instance to TVMTensor +pokeTensor :: forall d i e . (TVMData d i e) + => TVMTensor -> d -> IO () +pokeTensor ft d = do + when (tensorShape ft /= tvmDataShape d) $ + throwIO (PokeShapeMismatch (tensorShape ft) (tvmDataShape d)) + case tensorDevice ft of + KDLCPU -> do + tvmPoke d (castPtr (unsafeTensorData ft)) + x -> do + allocaArray (fromCSize $ tensorSize ft) $ \parr -> do + withForeignPtr ft $ \pt -> do + tvmPoke d (castPtr parr) + _ <- tvmArrayCopyFromBytes pt parr (tensorSize ft) + return () + diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -13,7 +13,7 @@ label, classify, whenFail, counterexample, elements, vectorOf, Gen, Testable, frequency, sized, Property, arbitrary, Arbitrary, listOf) -import Test.QuickCheck.Monadic (forAllM, monadicIO, run, assert) +import Test.QuickCheck.Monadic (forAllM, monadicIO, run, assert, wp) import Control.Monad (when) import Data.Functor.Foldable (Fix(..), Recursive(..), Corecursive(..)) @@ -66,6 +66,26 @@ assertEpsilonEqual :: (EpsilonEqual a, HasCallStack) => String -> Rational -> a -> a -> Assertion assertEpsilonEqual msg eps a b = assertBool msg (epsilonEqual eps a b) + +{- +testFunction :: forall d1 i1 e1 d2 i2 e2 . (TVMData d1 i1 e1, TVMData d2 i2 e2) => + [Integer] -> ([Integer] -> Stmt Function) -> (d1 -> d2) -> PropertyM IO a +testFunction ishape func_ut func_checker = + withTestModule (func_ut ishape) $ + \(ModuleLib p m) -> do + withModule p $ \hmod -> do + withFunction (funcName $ head $ modFuncs $ m) hmod $ \fmod -> do + a <- liftIO $ newEmptyTensor @e1 ishape KDLCPU 0 + c <- liftIO $ newEmptyTensor @e2 oshape KDLCPU 0 + forAllM arbitrary $ \x -> do + liftIO $ callTensorFunction c fmod [a] + c_ <- liftIO $ peekTensor c + assertEpsilonEqual "Function result" epsilon [[6.0::Float]] c_ +-} + +epsilon :: Rational +epsilon = 1e-5 + main :: IO () main = defaultMain $ testGroup "All" $ reverse [ @@ -172,13 +192,13 @@ do dump <- printFunction defaultConfig =<< do - stageFunction $ do + stageFunctionT $ do s <- shapevar [10] function "vecadd" [("A",float32,s),("B",float32,s)] $ \[a,b] -> do - compute s $ \e -> a![e!0] + b![e!0] + compute s $ \e -> a![e] + b![e] assertBool "dump should contain 'produce' keyword" $ isInfixOf "produce" dump - , testCase "Simple model should work" $ + , testCase "Simple model should work, withModule/withFunction case" $ let dim0 = 4 :: Integer fname = "vecadd" @@ -186,7 +206,7 @@ withTestModule (do s <- shapevar [fromInteger dim0] function fname [("A",float32,s),("B",float32,s)] $ \[a,b] -> do - compute s $ \e -> a![e!0] + b![e!0] + compute s $ \e -> a![e] + b![e] ) $ \(ModuleLib p _) -> do withModule p $ \hmod -> do @@ -197,13 +217,32 @@ callTensorFunction c fmod [a,b] assertEqual "Simple model result" [11,22,33,44::Float] =<< peekTensor c + , testCase "Simple model should work, loadModule/loadFunction case" $ + let + dim0 = 4 :: Integer + fname = "vecadd" + in do + withTestModule (do + s <- shapevar [fromInteger dim0] + function fname [("A",float32,s),("B",float32,s)] $ \[a,b] -> do + compute s $ \e -> a![e] + b![e] + ) $ + \(ModuleLib mod_path _) -> do + m <- loadModule mod_path + f <- loadFunction "vecadd" m + a <- newTensor @[Float] [1,2,3,4] KDLCPU 0 + b <- newTensor @[Float] [10,20,30,40] KDLCPU 0 + c <- newEmptyTensor @Float [dim0] KDLCPU 0 + callTensorFunction c f [a,b] + assertEqual "Simple model result" [11,22,33,44::Float] =<< peekTensor c + , testCase "Reduce axis operation should compile" $ withTestModule (do s <- shapevar [4] function "reduce" [("A",float32,s)] $ \[a] -> do IterVar r <- reduce_axis (0,3) - compute ShapeScalar $ \_ -> esum (a![r], [r]) + compute ShapeScalar $ \(_::Expr) -> esum (a![r], [r]) ) $ \_ -> return () , testCase "Conv2d operation should compile" $ @@ -254,7 +293,7 @@ c <- newEmptyTensor @Float [4] KDLCPU 0 callTensorFunction c fmod [a] c_ <- peekTensor c - assertEpsilonEqual "Simple model result" 1e-5 out c_ + assertEpsilonEqual "Simple model result" epsilon out c_ , testCase "Split primitive should compile" $ @@ -264,5 +303,23 @@ c <- assign $ split a [1] 0 return (c!0) ) $ \_ -> return () + + , testCase "Differentiate should work" $ + + withTestModule (do + sa <- shapevar [1] + function "difftest" [("A",float32,sa) ] $ \[a] -> do + c <- compute sa $ \i -> (a![i])*(a![i]) + dc <- assign $ differentiate c [a] + return (dc!0) + ) $ + \(ModuleLib p _) -> do + withModule p $ \hmod -> do + withFunction "difftest" hmod $ \fmod -> do + a <- newTensor @[Float] [3.0] KDLCPU 0 + c <- newEmptyTensor @Float [1,1] KDLCPU 0 + callTensorFunction c fmod [a] + c_ <- peekTensor c + assertEpsilonEqual "Differentiate result" epsilon [[6.0::Float]] c_ ]