htvm 0.1.0.0 → 0.1.1
raw patch · 9 files changed
+434/−267 lines, 9 files
Files
- htvm.cabal +2/−1
- src/HTVM/EDSL/Build.hs +4/−4
- src/HTVM/EDSL/Monad.hs +41/−10
- src/HTVM/EDSL/Printer.hs +3/−0
- src/HTVM/EDSL/Types.hs +18/−57
- src/HTVM/Runtime.hs +2/−0
- src/HTVM/Runtime/FFI.chs +115/−188
- src/HTVM/Runtime/TVMData.hs +185/−0
- test/Main.hs +64/−7
htvm.cabal view
@@ -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
src/HTVM/EDSL/Build.hs view
@@ -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)
src/HTVM/EDSL/Monad.hs view
@@ -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]] {- ____ _ _ _
src/HTVM/EDSL/Printer.hs view
@@ -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 <topi/nn.h>" line $ "#include <topi/elemwise.h>" line $ "#include <topi/transform.h>"+ line $ "#include <tvm/autodiff.h>" line $ "" line $ "static inline tvm::Array<tvm::Expr> \@@ -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<tvm::Expr> htvm_shape(tvm::Tensor t) { return t->shape; }" line $ "tvm::Array<tvm::Expr> htvm_shape(tvm::Array<tvm::Expr> t) { return t; }"+ line $ "tvm::Array<tvm::Tensor> htvm_differentiate(tvm::Tensor t, tvm::Array<tvm::Tensor> a){ return tvm::ir::Differentiate(t,a)->result; }" line $ ""
src/HTVM/EDSL/Types.hs view
@@ -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<Expr>`. 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)---
src/HTVM/Runtime.hs view
@@ -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
src/HTVM/Runtime/FFI.chs view
@@ -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 ()
+ src/HTVM/Runtime/TVMData.hs view
@@ -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 ()+
test/Main.hs view
@@ -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_ ]