halide-haskell (empty) → 0.0.1.0
raw patch · 30 files changed
+6345/−0 lines, 30 filesdep +HUnitdep +QuickCheckdep +Win32
Dependencies added: HUnit, QuickCheck, Win32, base, bytestring, constraints, filepath, halide-haskell, hspec, inline-c, inline-c-cpp, primitive, template-haskell, temporary, text, unix, vector
Files
- CHANGELOG.md +11/−0
- LICENSE +29/−0
- README.md +80/−0
- example/Example01.hs +13/−0
- example/GettingStarted.hs +32/−0
- halide-haskell.cabal +148/−0
- src/Language/Halide.hs +223/−0
- src/Language/Halide/Buffer.hs +419/−0
- src/Language/Halide/Context.hs +139/−0
- src/Language/Halide/Dimension.hs +153/−0
- src/Language/Halide/Expr.hs +657/−0
- src/Language/Halide/Func.hs +1143/−0
- src/Language/Halide/Kernel.hs +365/−0
- src/Language/Halide/LoopLevel.hs +177/−0
- src/Language/Halide/Prelude.hs +41/−0
- src/Language/Halide/RedundantConstraints.hs +18/−0
- src/Language/Halide/Schedule.hs +562/−0
- src/Language/Halide/Target.hs +518/−0
- src/Language/Halide/Trace.hs +181/−0
- src/Language/Halide/Type.hs +414/−0
- src/Language/Halide/Utils.hs +37/−0
- test/Language/Halide/BufferSpec.hs +54/−0
- test/Language/Halide/ExprSpec.hs +112/−0
- test/Language/Halide/FuncSpec.hs +252/−0
- test/Language/Halide/KernelSpec.hs +66/−0
- test/Language/Halide/LoopLevelSpec.hs +120/−0
- test/Language/Halide/ScheduleSpec.hs +229/−0
- test/Language/Halide/TargetSpec.hs +10/−0
- test/Spec.hs +1/−0
- test/Utils.hs +141/−0
+ CHANGELOG.md view
@@ -0,0 +1,11 @@+# Changelog++`halide-haskell` uses [PVP Versioning][1].+The changelog is available [on GitHub][2].++## 0.0.0.0++* Initially created.++[1]: https://pvp.haskell.org+[2]: https://github.com/twesterhout/halide-haskell/releases
+ LICENSE view
@@ -0,0 +1,29 @@+BSD 3-Clause License++Copyright (c) 2021, Tom Westerhout+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++1. Redistributions of source code must retain the above copyright notice, this+ list of conditions and the following disclaimer.++2. Redistributions in binary form must reproduce the above copyright notice,+ this list of conditions and the following disclaimer in the documentation+ and/or other materials provided with the distribution.++3. Neither the name of the copyright holder nor the names of its+ contributors may be used to endorse or promote products derived from+ this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,80 @@+# halide-haskell++[](https://github.com/twesterhout/halide-haskell/actions/workflows/ci.yml)+[](https://hackage.haskell.org/package/halide-haskell-0.0.1.0/candidate)+[](LICENSE)++[Halide](https://halide-lang.org/) is a programming language designed to make+it easier to write high-performance image and array processing code on modern+machines. Rather than being a standalone programming language, Halide is+embedded in C++. This means you write C++ code that builds an in-memory+representation of a Halide pipeline using Halide's C++ API. You can then+compile this representation to an object file, or JIT-compile it and run it in+the same process.++**This package provides Haskell bindings that allow to write Halide embedded in+Haskell without C++** 😋.++ - [Tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials)+ - [Reference documentation](https://hackage.haskell.org/package/halide-haskell-0.0.1.0)++## 🚀 Getting started++As a simple example, here's how you could implement array addition with halide-haskell:++```haskell+{-# LANGUAGE AllowAmbiguousTypes, DataKinds, OverloadedStrings #-}+import Language.Halide++-- The algorithm+mkArrayPlus = compile $ \a b -> do+ -- Create an index variable+ i <- mkVar "i"+ -- Define the resulting function. We call it "out".+ -- In pseudocode it's equivalent to the following: out[i] = a[i] + b[i]+ out <- define "out" i $ a ! i + b ! i+ -- Perform a fancy optimization and use SIMD: we split the loop over i into+ -- an inner and an outer loop and then vectorize the inner loop+ inner <- mkVar "inner"+ split TailAuto i (i, inner) 4 out >>= vectorize inner++-- Example usage of our Halide pipeline+main :: IO ()+main = do+ let a, b :: [Float]+ a = [1, 2, 3, 4, 5]+ b = [6, 7, 8, 9, 10]+ -- Compile the code+ arrayPlus <- mkArrayPlus+ -- We tell Halide to treat our list as a one-dimensional buffer+ withHalideBuffer @1 @Float a $ \a' ->+ withHalideBuffer b $ \b' ->+ -- allocate a temporary buffer for the output+ allocaCpuBuffer [length a] $ \out' -> do+ -- execute the kernel -- it is a normal function call!+ arrayPlus a' b' out'+ -- print the result+ print =<< peekToList out'+```++For more examples, have a look a the [tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials).++## 🔨 Contributing++Currently, the best way to get started is to use Nix:++```sh+nix develop+```++This will drop you into a shell with all the necessary tools to build the code such that you can do++```sh+cabal build+```++and++```sh+cabal test+```
+ example/Example01.hs view
@@ -0,0 +1,13 @@+import Control.Monad (when)+import Language.Halide++main :: IO ()+main = do+ let !host = hostTarget+ putStrLn $ "[+] host target is " <> show host+ when (hostSupportsTargetDevice (setFeature FeatureOpenCL host)) $ do+ putStrLn "[+] OpenCL is supported! Testing ..."+ testOpenCL+ when (hostSupportsTargetDevice (setFeature FeatureCUDA host)) $ do+ putStrLn "[+] CUDA is supported! Testing ..."+ testCUDA
+ example/GettingStarted.hs view
@@ -0,0 +1,32 @@+module Main (main) where++import qualified Data.Vector.Storable as S+import qualified Data.Vector.Storable.Mutable as SM+import Language.Halide+import System.IO.Unsafe (unsafePerformIO)++mkVectorPlus :: forall a. (IsHalideType a, Num a) => IO (S.Vector a -> S.Vector a -> S.Vector a)+mkVectorPlus = do+ -- First, compile the kernel+ kernel <- compile $ \a b -> do+ -- Create an index variable+ i <- mkVar "i"+ -- Define the resulting function. We call it "out".+ -- In pseudocode it's equivalent to the following: out[i] = a[i] + b[i]+ define "out" i $ a ! i + b ! i+ -- Create a Haskell function that will invoke the kernel+ pure $ \v1 v2 -> unsafePerformIO $ do+ out <- SM.new (S.length v1)+ withHalideBuffer @1 @a v1 $ \a ->+ withHalideBuffer @1 @a v2 $ \b ->+ withHalideBuffer @1 @a out $ \out' ->+ kernel a b out'+ S.unsafeFreeze out++main :: IO ()+main = do+ let a, b :: S.Vector Float+ a = S.fromList [1, 2, 3]+ b = S.fromList [4, 5, 6]+ vectorPlus <- mkVectorPlus+ print (vectorPlus a b)
+ halide-haskell.cabal view
@@ -0,0 +1,148 @@+cabal-version: 3.0+name: halide-haskell+version: 0.0.1.0+synopsis: Haskell bindings to Halide+description:+ Halide is a programming language designed to make it easier to write+ high-performance image and array processing code on modern machines. Rather+ than being a standalone programming language, Halide is embedded in C++. This+ means you write C++ code that builds an in-memory representation of a Halide+ pipeline using Halide's C++ API. You can then compile this representation to+ an object file, or JIT-compile it and run it in the same process.++ This package provides Haskell bindings that allow to write Halide embedded in+ Haskell without C++.++ The best way to learn Halide is to have a look at the+ [tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials).+ Reference documentation is provided by the haddocks of the 'Language.Halide'+ module.++homepage: https://github.com/twesterhout/halide-haskell+bug-reports: https://github.com/twesterhout/halide-haskell/issues+license: BSD-3-Clause+license-file: LICENSE+author: Tom Westerhout+maintainer:+ Tom Westerhout <14264576+twesterhout@users.noreply.github.com>++category: Language+copyright: 2022-2023 Tom Westerhout+build-type: Simple+extra-doc-files:+ CHANGELOG.md+ README.md++tested-with: GHC ==9.2.4 || ==9.2.5 || ==9.4.4++source-repository head+ type: git+ location: https://github.com/twesterhout/halide-haskell.git++common common-options+ build-depends: base >=4.16.0.0 && <5+ ghc-options:+ -W -Wall -Wcompat -Widentities -Wincomplete-uni-patterns+ -Wincomplete-record-updates -Wredundant-constraints+ -fhide-source-paths -Wmissing-export-lists -Wpartial-fields+ -Wmissing-deriving-strategies++ default-language: GHC2021+ default-extensions:+ DataKinds+ DerivingStrategies+ FunctionalDependencies+ LambdaCase+ OverloadedRecordDot+ OverloadedStrings+ TypeFamilies+ ViewPatterns++library+ import: common-options+ hs-source-dirs: src+ exposed-modules: Language.Halide+ other-modules:+ Language.Halide.Buffer+ Language.Halide.Context+ Language.Halide.Dimension+ Language.Halide.Expr+ Language.Halide.Func+ Language.Halide.Kernel+ Language.Halide.LoopLevel+ Language.Halide.Prelude+ Language.Halide.RedundantConstraints+ Language.Halide.Schedule+ Language.Halide.Target+ Language.Halide.Trace+ Language.Halide.Type+ Language.Halide.Utils++ build-depends:+ , bytestring >=0.11.1.0 && <0.12+ , constraints >=0.13.4 && <0.14+ , filepath >=1.4.2.1 && <2.0+ , inline-c >=0.9.1.6 && <0.10+ , inline-c-cpp >=0.5.0.0 && <0.6+ , primitive >=0.7.3.0 && <0.8+ , template-haskell >=2.18.0.0 && <3.0+ , temporary >=1.3 && <2.0+ , text >=1.2.5.0 && <3.0+ , vector >=0.12.3.0 && <0.13++ if os(windows)+ cpp-options: -DUSE_DLOPEN=0+ build-depends: Win32++ else+ cpp-options: -DUSE_DLOPEN=1+ build-depends: unix >=2.7.2.2 && <3.0++ extra-libraries:+ Halide+ stdc++++executable halide-haskell+ import: common-options+ hs-source-dirs: example+ main-is: Example01.hs+ build-depends:+ , halide-haskell+ , vector++executable getting-started+ import: common-options+ hs-source-dirs: example+ main-is: GettingStarted.hs+ build-depends:+ , halide-haskell+ , vector++test-suite halide-haskell-test+ import: common-options+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Spec.hs+ other-modules:+ Language.Halide.BufferSpec+ Language.Halide.ExprSpec+ Language.Halide.FuncSpec+ Language.Halide.KernelSpec+ Language.Halide.LoopLevelSpec+ Language.Halide.ScheduleSpec+ Language.Halide.TargetSpec+ Utils++ build-depends:+ , halide-haskell+ , hspec+ , HUnit+ , inline-c+ , inline-c-cpp+ , QuickCheck+ , text+ , vector++ ghc-options: -threaded -rtsopts -with-rtsopts=-N++-- build-tools-depends: hspec-discover:hspec-discover
+ src/Language/Halide.hs view
@@ -0,0 +1,223 @@+-- |+-- Module : Language.Halide+-- Copyright : (c) Tom Westerhout, 2023+--+-- This package provides Haskell bindings that allow to write Halide embedded in Haskell without C++.+--+-- This module contains the reference documentation for Halide. If you're new, the best way to learn Halide is to have a look at the [tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials).+module Language.Halide+ ( -- * Scalar expressions++ -- | The basic building block of Halide pipelines is 'Expr'. @Expr a@ represents a scalar expression of+ -- type @a@, where @a@ must be an instance of 'IsHalideType'.+ Expr (..)+ , Var+ , RVar+ , VarOrRVar+ , IsHalideType++ -- ** Creating+ , mkExpr+ , mkVar+ , mkRVar+ , undef+ , cast+ , bool++ -- ** Inspecting+ , toIntImm+ , printed+ , evaluate++ -- ** Comparisons++ -- | We can't use 'Eq' and 'Ord' instances here, because we want the comparison to happen+ -- when the pipeline is run rather than when it's built. Hence, we define lifted version of+ -- various comparison operators. Note, that infix versions of the these functions have the+ -- same precedence as the normal comparison operators.+ , eq+ , neq+ , lt+ , lte+ , gt+ , gte++ -- * Functions+ , Func (..)+ , FuncTy (..)+ , Stage (..)++ -- ** Creating+ , define+ , update+ , (!)++ -- ** Inspecting+ , getArgs+ , hasUpdateDefinitions+ , getUpdateStage++ -- * Buffers++ -- | In the C interface of Halide, buffers are described by the C struct+ -- [@halide_buffer_t@](https://halide-lang.org/docs/structhalide__buffer__t.html). On the Haskell side,+ -- we have 'HalideBuffer'.+ , HalideBuffer (..)+ -- | To easily test out your pipeline, there are helper functions to create 'HalideBuffer's without+ -- worrying about the low-level representation.+ , allocaCpuBuffer+ -- | Buffers can also be converted to lists to easily print them for debugging.+ , IsListPeek (..)+ -- | For production usage however, you don't want to work with lists. Instead, you probably want Halide+ -- to work with your existing array data types. For this, we define 'IsHalideBuffer' typeclass that+ -- teaches Halide how to convert your data into a 'HalideBuffer'. Depending on how you implement the+ -- instance, this can be very efficient, because it need not involve any memory copying.+ , IsHalideBuffer (..)+ , withHalideBuffer+ -- | There are also helper functions to simplify writing instances of 'IsHalideBuffer'.+ , bufferFromPtrShapeStrides+ , bufferFromPtrShape++ -- * Running the pipelines++ -- | There are a few ways how one can run a Halide pipeline.+ --+ -- The simplest way to build a t'Func' and then call 'realize' to evaluate it over a rectangular domain.+ , realize+ , asBufferParam+ -- | The drawback of calling 'realize' all the time is that it's impossible to pass parameters to pipelines.+ -- We can define pipelines that operate on buffers using 'asBufferParam', but we have to recompile the+ -- pipeline for every new buffer.+ --+ -- A better way to handle pipeline parameters is to define a /Haskell/ function that accepts t'Expr's+ -- and t'Func's as arguments and returns a 'Func'. We can then pass this function to 'compile'+ -- (or 'compileForTarget'), and it compile it into a /Haskell/ function that can now be invoked with+ -- normal scalars instead of t'Expr's and @Ptr 'HalideBuffer'@s instead of 'Func's.+ , compile++ -- ** Parameters++ -- | Similar to how we can specify the name of a variable in 'mkVar' (or 'mkRVar') or function in 'define',+ -- one can also specify the name of a pipeline parameter. This is achieved by using the @ViewPatterns@+ -- extension together with the 'scalar' and 'buffer' helper functions.+ , buffer+ , scalar+ -- | Another common thing to do with the parameters is to explicitly specify their shapes. For this, we expose the 'Dimension' type:+ , Dimension (..)+ , dim+ , setMin+ , setExtent+ , setStride+ , setEstimate++ -- ** Targets+ , Target (..)+ , hostTarget+ , gpuTarget+ , compileForTarget+ , DeviceAPI (..)+ , TargetFeature (..)+ , setFeature+ , hasGpuFeature+ , hostSupportsTargetDevice++ -- * Scheduling+ , Schedulable (..)+ , TailStrategy (..)+ , LoopLevel (..)+ , LoopLevelTy (..)+ , LoopAlignStrategy (..)+ , computeRoot+ , getStage+ , getLoopLevel+ , getLoopLevelAtStage+ , asUsed+ , asUsedBy+ , copyToDevice+ , copyToHost+ , storeAt+ , computeAt+ , estimate+ , bound++ -- * Debugging / Tracing++ -- | For debugging, it's often useful to observe the value of an expression when it's evaluated. If you+ -- have a complex expression that does not depend on any buffers or indices, you can 'evaluate' it.+ -- | However, often an expression is only used within a definition of a pipeline, and it's impossible to+ -- call 'evaluate' on it. In such cases, it can be wrapped with 'printed' to indicate to Halide that the+ -- value of the expression should be dumped to screen when it's computed.+ , prettyLoopNest+ , compileToLoweredStmt+ , StmtOutputFormat (..)+ , TraceEvent (..)+ , TraceEventCode (..)+ , TraceLoadStoreContents (..)+ , setCustomTrace+ , traceStores+ , traceLoads+ , collectIterationOrder++ -- * Type helpers+ , IsTuple (..)+ , ToTuple+ , FromTuple+ , IndexTuple+ , Length+ , All++ -- * Internal+ , compileToCallable+ , testCUDA+ , testOpenCL+ , SomeLoopLevel (..)+ , RawHalideBuffer (..)+ , HalideDimension (..)+ , HalideDeviceInterface+ , rowMajorStrides+ , colMajorStrides+ , isDeviceDirty+ , isHostDirty+ , bufferCopyToHost+ , module Language.Halide.Schedule+ , IsFuncBuilder+ , ReturnsFunc+ , FunctionArguments+ , FunctionReturn+ , Curry (..)+ , UnCurry (..)+ , Lowered++ -- ** inline-c helpers+ , importHalide+ , CxxExpr+ , CxxVar+ , CxxRVar+ , CxxParameter+ , CxxFunc+ , CxxImageParam+ , CxxStage+ , CxxDimension+ , CxxTarget+ , CxxLoopLevel++ -- * Convenience re-exports+ , Int32+ , Ptr+ , KnownNat+ )+where++import Foreign.Ptr (Ptr)+import GHC.TypeLits (KnownNat)+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Dimension+import Language.Halide.Expr+import Language.Halide.Func+import Language.Halide.Kernel+import Language.Halide.LoopLevel+import Language.Halide.Schedule+import Language.Halide.Target+import Language.Halide.Trace+import Language.Halide.Type
+ src/Language/Halide/Buffer.hs view
@@ -0,0 +1,419 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}++-- |+-- Module : Language.Halide.Buffer+-- Description : Buffers+-- Copyright : (c) Tom Westerhout, 2021-2023+--+-- A buffer in Halide is a __view__ of some multidimensional array. Buffers can reference data that's+-- located on a CPU, GPU, or another device. Halide pipelines use buffers for both input and output arguments.+module Language.Halide.Buffer+ ( -- * Buffers++ --++ -- | In the C interface of Halide, buffers are described by the C struct+ -- [@halide_buffer_t@](https://halide-lang.org/docs/structhalide__buffer__t.html). On the Haskell side,+ -- we have 'HalideBuffer'.+ HalideBuffer (..)+ -- | To easily test out your pipeline, there are helper functions to create 'HalideBuffer's without+ -- worrying about the low-level representation.+ , allocaCpuBuffer+ -- | Buffers can also be converted to lists to easily print them for debugging.+ , IsListPeek (..)+ -- | For production usage however, you don't want to work with lists. Instead, you probably want Halide+ -- to work with your existing array data types. For this, we define 'IsHalideBuffer' typeclass that+ -- teaches Halide how to convert your data into a 'HalideBuffer'. Depending on how you implement the+ -- instance, this can be very efficient, because it need not involve any memory copying.+ , IsHalideBuffer (..)+ , withHalideBuffer+ -- | There are also helper functions to simplify writing instances of 'IsHalideBuffer'.+ , bufferFromPtrShapeStrides+ , bufferFromPtrShape++ -- * Internals+ , RawHalideBuffer (..)+ , HalideDimension (..)+ , HalideDeviceInterface+ , rowMajorStrides+ , colMajorStrides+ , isDeviceDirty+ , isHostDirty+ , bufferCopyToHost+ )+where++import Control.Monad (forM, unless, when)+import Control.Monad.ST (RealWorld)+import Data.Foldable (foldl')+import Data.Int+import Data.Kind (Type)+import qualified Data.List as List+import Data.Proxy+import qualified Data.Vector.Storable as S+import qualified Data.Vector.Storable.Mutable as SM+import Data.Word+import Foreign.Marshal.Array+import Foreign.Marshal.Utils+import Foreign.Ptr+import Foreign.Storable+import GHC.Stack (HasCallStack)+import GHC.TypeNats+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Cpp.Exception as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Context+import Language.Halide.Type++-- | Information about a dimension in a buffer.+--+-- It is the Haskell analogue of [@halide_dimension_t@](https://halide-lang.org/docs/structhalide__dimension__t.html).+data HalideDimension = HalideDimension+ { halideDimensionMin :: {-# UNPACK #-} !Int32+ -- ^ Starting index.+ , halideDimensionExtent :: {-# UNPACK #-} !Int32+ -- ^ Length of the dimension.+ , halideDimensionStride :: {-# UNPACK #-} !Int32+ -- ^ Stride along this dimension.+ , halideDimensionFlags :: {-# UNPACK #-} !Word32+ -- ^ Extra flags.+ }+ deriving stock (Read, Show, Eq)++instance Storable HalideDimension where+ sizeOf _ = 16+ {-# INLINE sizeOf #-}+ alignment _ = 4+ {-# INLINE alignment #-}+ peek p =+ HalideDimension+ <$> peekByteOff p 0+ <*> peekByteOff p 4+ <*> peekByteOff p 8+ <*> peekByteOff p 12+ {-# INLINE peek #-}+ poke p x = do+ pokeByteOff p 0 (halideDimensionMin x)+ pokeByteOff p 4 (halideDimensionExtent x)+ pokeByteOff p 8 (halideDimensionStride x)+ pokeByteOff p 12 (halideDimensionFlags x)+ {-# INLINE poke #-}++-- | @simpleDimension extent stride@ creates a @HalideDimension@ of size @extent@ separated by+-- @stride@.+simpleDimension :: Int -> Int -> HalideDimension+simpleDimension extent stride = HalideDimension 0 (fromIntegral extent) (fromIntegral stride) 0+{-# INLINE simpleDimension #-}++-- | Get strides corresponding to row-major ordering+rowMajorStrides+ :: Integral a+ => [a]+ -- ^ Extents+ -> [a]+rowMajorStrides = drop 1 . scanr (*) 1++-- | Get strides corresponding to column-major ordering.+colMajorStrides+ :: Integral a+ => [a]+ -- ^ Extents+ -> [a]+colMajorStrides = scanl (*) 1 . init++-- | Haskell analogue of [@halide_device_interface_t@](https://halide-lang.org/docs/structhalide__device__interface__t.html).+data HalideDeviceInterface++-- | The low-level untyped Haskell analogue of [@halide_buffer_t@](https://halide-lang.org/docs/structhalide__buffer__t.html).+--+-- It's quite difficult to use 'RawHalideBuffer' correctly, and misusage can result in crashes and+-- segmentation faults. Hence, prefer the higher-level 'HalideBuffer' wrapper for all your code+data RawHalideBuffer = RawHalideBuffer+ { halideBufferDevice :: !Word64+ , halideBufferDeviceInterface :: !(Ptr HalideDeviceInterface)+ , halideBufferHost :: !(Ptr Word8)+ , halideBufferFlags :: !Word64+ , halideBufferType :: !HalideType+ , halideBufferDimensions :: !Int32+ , halideBufferDim :: !(Ptr HalideDimension)+ , halideBufferPadding :: !(Ptr ())+ }+ deriving stock (Show, Eq)++-- | An @n@-dimensional buffer of elements of type @a@.+--+-- Most pipelines use @'Ptr' ('HalideBuffer' n a)@ for input and output array arguments.+newtype HalideBuffer (n :: Nat) (a :: Type) = HalideBuffer {unHalideBuffer :: RawHalideBuffer}+ deriving stock (Show, Eq)++importHalide++instance Storable RawHalideBuffer where+ sizeOf _ = 56+ alignment _ = 8+ peek p =+ RawHalideBuffer+ <$> peekByteOff p 0 -- device+ <*> peekByteOff p 8 -- interface+ <*> peekByteOff p 16 -- host+ <*> peekByteOff p 24 -- flags+ <*> peekByteOff p 32 -- type+ <*> peekByteOff p 36 -- dimensions+ <*> peekByteOff p 40 -- dim+ <*> peekByteOff p 48 -- padding+ poke p x = do+ pokeByteOff p 0 (halideBufferDevice x)+ pokeByteOff p 8 (halideBufferDeviceInterface x)+ pokeByteOff p 16 (halideBufferHost x)+ pokeByteOff p 24 (halideBufferFlags x)+ pokeByteOff p 32 (halideBufferType x)+ pokeByteOff p 36 (halideBufferDimensions x)+ pokeByteOff p 40 (halideBufferDim x)+ pokeByteOff p 48 (halideBufferPadding x)++-- | Construct a 'HalideBuffer' from a pointer to the data, a list of extents,+-- and a list of strides, and use it in an 'IO' action.+--+-- This function throws a runtime error if the number of dimensions does not+-- match @n@.+bufferFromPtrShapeStrides+ :: forall n a b+ . (HasCallStack, KnownNat n, IsHalideType a)+ => Ptr a+ -- ^ CPU pointer to the data+ -> [Int]+ -- ^ Extents (in number of elements, __not__ in bytes)+ -> [Int]+ -- ^ Strides (in number of elements, __not__ in bytes)+ -> (Ptr (HalideBuffer n a) -> IO b)+ -- ^ Action to run+ -> IO b+bufferFromPtrShapeStrides p shape stride action =+ withArrayLen (zipWith simpleDimension shape stride) $ \n dim -> do+ unless (n == fromIntegral (natVal (Proxy @n))) $+ error $+ "specified wrong number of dimensions: "+ <> show n+ <> "; expected "+ <> show (natVal (Proxy @n))+ <> " from the type declaration"+ let !buffer =+ RawHalideBuffer+ { halideBufferDevice = 0+ , halideBufferDeviceInterface = nullPtr+ , halideBufferHost = castPtr p+ , halideBufferFlags = 0+ , halideBufferType = halideTypeFor (Proxy :: Proxy a)+ , halideBufferDimensions = fromIntegral n+ , halideBufferDim = dim+ , halideBufferPadding = nullPtr+ }+ with buffer $ \bufferPtr -> do+ r <- action (castPtr bufferPtr)+ hasDataOnDevice <-+ toEnum . fromIntegral+ <$> [CU.exp| bool { $(halide_buffer_t* bufferPtr)->device } |]+ when hasDataOnDevice $+ error "the Buffer still references data on the device; did you forget to call copyToHost?"+ pure r++-- | Similar to 'bufferFromPtrShapeStrides', but assumes column-major ordering of data.+bufferFromPtrShape+ :: (HasCallStack, KnownNat n, IsHalideType a)+ => Ptr a+ -- ^ CPU pointer to the data+ -> [Int]+ -- ^ Extents (in number of elements, __not__ in bytes)+ -> (Ptr (HalideBuffer n a) -> IO b)+ -> IO b+bufferFromPtrShape p shape = bufferFromPtrShapeStrides p shape (colMajorStrides shape)++-- | Specifies that a type @t@ can be used as an @n@-dimensional Halide buffer with elements of type @a@.+class (KnownNat n, IsHalideType a) => IsHalideBuffer t n a where+ withHalideBufferImpl :: t -> (Ptr (HalideBuffer n a) -> IO b) -> IO b++-- | Treat a type @t@ as a 'HalideBuffer' and use it in an 'IO' action.+--+-- This function is a simple wrapper around 'withHalideBufferImpl', except that the order of type parameters+-- is reversed. If you have @TypeApplications@ extension enabled, this allows you to write+-- @withHalideBuffer @3 @Float yourBuffer@ to specify that you want a 3-dimensional buffer of @Float@.+withHalideBuffer :: forall n a t b. IsHalideBuffer t n a => t -> (Ptr (HalideBuffer n a) -> IO b) -> IO b+withHalideBuffer = withHalideBufferImpl @t @n @a++-- | Storable vectors are one-dimensional buffers. This involves no copying.+instance IsHalideType a => IsHalideBuffer (S.Vector a) 1 a where+ withHalideBufferImpl v f =+ S.unsafeWith v $ \dataPtr ->+ bufferFromPtrShape dataPtr [S.length v] f++-- | Storable vectors are one-dimensional buffers. This involves no copying.+instance IsHalideType a => IsHalideBuffer (S.MVector RealWorld a) 1 a where+ withHalideBufferImpl v f =+ SM.unsafeWith v $ \dataPtr ->+ bufferFromPtrShape dataPtr [SM.length v] f++-- | Lists can also act as Halide buffers. __Use for testing only.__+instance IsHalideType a => IsHalideBuffer [a] 1 a where+ withHalideBufferImpl v = withHalideBuffer (S.fromList v)++-- | Lists can also act as Halide buffers. __Use for testing only.__+instance IsHalideType a => IsHalideBuffer [[a]] 2 a where+ withHalideBufferImpl xs f = do+ let d0 = length xs+ d1 = if d0 == 0 then 0 else length (head xs)+ -- we want column-major ordering, so transpose first+ v = S.fromList (List.concat (List.transpose xs))+ when (S.length v /= d0 * d1) $+ error "list doesn't have a regular shape (i.e. rows have varying number of elements)"+ S.unsafeWith v $ \cpuPtr ->+ bufferFromPtrShape cpuPtr [d0, d1] f++-- | Lists can also act as Halide buffers. __Use for testing only.__+instance IsHalideType a => IsHalideBuffer [[[a]]] 3 a where+ withHalideBufferImpl xs f = do+ let d0 = length xs+ d1 = if d0 == 0 then 0 else length (head xs)+ d2 = if d1 == 0 then 0 else length (head (head xs))+ -- we want column-major ordering, so transpose first+ v =+ S.fromList+ . List.concat+ . List.concatMap List.transpose+ . List.transpose+ . fmap List.transpose+ $ xs+ when (S.length v /= d0 * d1 * d2) $+ error "list doesn't have a regular shape (i.e. rows have varying number of elements)"+ S.unsafeWith v $ \cpuPtr ->+ bufferFromPtrShape cpuPtr [d0, d1, d2] f++whenM :: Monad m => m Bool -> m () -> m ()+whenM cond f =+ cond >>= \case+ True -> f+ False -> pure ()++-- | Temporary allocate a CPU buffer.+--+-- This is useful for testing and debugging when you need to allocate an output buffer for your pipeline. E.g.+--+-- @+-- 'allocaCpuBuffer' [3, 3] $ \out -> do+-- myKernel out -- fill the buffer+-- print =<< 'peekToList' out -- print it for debugging+-- @+allocaCpuBuffer+ :: forall n a b+ . (HasCallStack, KnownNat n, IsHalideType a)+ => [Int]+ -> (Ptr (HalideBuffer n a) -> IO b)+ -> IO b+allocaCpuBuffer shape action =+ allocaArray numElements $ \cpuPtr ->+ bufferFromPtrShape cpuPtr shape $ \buf -> do+ r <- action buf+ whenM (isDeviceDirty (castPtr buf)) $+ error $+ "device_dirty is set on a CPU-only buffer; "+ <> "did you forget a copyToHost in your pipeline?"+ pure r+ where+ numElements = foldl' (*) 1 shape++-- | Do we have changes on the device the have not been copied to the host?+isDeviceDirty :: Ptr RawHalideBuffer -> IO Bool+isDeviceDirty p =+ toBool <$> [CU.exp| bool { $(const halide_buffer_t* p)->device_dirty() } |]++-- | Do we have changes on the device the have not been copied to the host?+isHostDirty :: Ptr RawHalideBuffer -> IO Bool+isHostDirty p =+ toBool <$> [CU.exp| bool { $(const halide_buffer_t* p)->host_dirty() } |]++-- | Copy the underlying memory from device to host.+bufferCopyToHost :: Ptr RawHalideBuffer -> IO ()+bufferCopyToHost p =+ [C.throwBlock| void {+ auto& buf = *$(halide_buffer_t* p);+ if (buf.device_dirty()) {+ if (buf.device_interface == nullptr) {+ throw std::runtime_error{"bufferCopyToHost: device_dirty is set, "+ "but device_interface is NULL"};+ }+ if (buf.host == nullptr) {+ throw std::runtime_error{"bufferCopyToHost: host is NULL; "+ "did you forget to allocate memory?"};+ }+ buf.device_interface->copy_to_host(nullptr, &buf);+ }+ } |]++checkNumberOfDimensions :: forall n. (HasCallStack, KnownNat n) => RawHalideBuffer -> IO ()+checkNumberOfDimensions raw = do+ unless (fromIntegral (natVal (Proxy @n)) == raw.halideBufferDimensions) $+ error $+ "type-level and runtime number of dimensions do not match: "+ <> show (natVal (Proxy @n))+ <> " != "+ <> show raw.halideBufferDimensions++-- | Specifies that @a@ can be converted to a list. This is very similar to 'GHC.Exts.IsList' except that+-- we read the list from a @'Ptr'@ rather than converting directly.+class IsListPeek a where+ type ListPeekElem a :: Type+ peekToList :: HasCallStack => Ptr a -> IO [ListPeekElem a]++instance IsHalideType a => IsListPeek (HalideBuffer 0 a) where+ type ListPeekElem (HalideBuffer 0 a) = a+ peekToList p = do+ whenM (isDeviceDirty (castPtr p)) $+ error "cannot peek data from device; call bufferCopyToHost first"+ raw <- peek (castPtr @_ @RawHalideBuffer p)+ checkNumberOfDimensions @0 raw+ fmap pure . peek $ castPtr @_ @a (halideBufferHost raw)++instance IsHalideType a => IsListPeek (HalideBuffer 1 a) where+ type ListPeekElem (HalideBuffer 1 a) = a+ peekToList p = do+ whenM (isDeviceDirty (castPtr p)) $+ error "cannot peek data from device; call bufferCopyToHost first"+ raw <- peek (castPtr @_ @RawHalideBuffer p)+ (HalideDimension min0 extent0 stride0 _) <- peekElemOff (halideBufferDim raw) 0+ let ptr0 = castPtr @_ @a (halideBufferHost raw)+ forM [0 .. extent0 - 1] $ \i0 ->+ peekElemOff ptr0 (fromIntegral (min0 + stride0 * i0))++instance IsHalideType a => IsListPeek (HalideBuffer 2 a) where+ type ListPeekElem (HalideBuffer 2 a) = [a]+ peekToList p = do+ whenM (isDeviceDirty (castPtr p)) $+ error "cannot peek data from device; call bufferCopyToHost first"+ raw <- peek (castPtr @_ @RawHalideBuffer p)+ (HalideDimension min0 extent0 stride0 _) <- peekElemOff (halideBufferDim raw) 0+ (HalideDimension min1 extent1 stride1 _) <- peekElemOff (halideBufferDim raw) 1+ let ptr0 = castPtr @_ @a (halideBufferHost raw)+ forM [0 .. extent0 - 1] $ \i0 -> do+ let ptr1 = ptr0 `advancePtr` fromIntegral (min0 + stride0 * i0)+ forM [0 .. extent1 - 1] $ \i1 ->+ peekElemOff ptr1 (fromIntegral (min1 + stride1 * i1))++instance IsHalideType a => IsListPeek (HalideBuffer 3 a) where+ type ListPeekElem (HalideBuffer 3 a) = [[a]]+ peekToList p = do+ whenM (isDeviceDirty (castPtr p)) $+ error "cannot peek data from device; call bufferCopyToHost first"+ raw <- peek (castPtr @_ @RawHalideBuffer p)+ (HalideDimension min0 extent0 stride0 _) <- peekElemOff (halideBufferDim raw) 0+ (HalideDimension min1 extent1 stride1 _) <- peekElemOff (halideBufferDim raw) 1+ (HalideDimension min2 extent2 stride2 _) <- peekElemOff (halideBufferDim raw) 2+ let ptr0 = castPtr @_ @a (halideBufferHost raw)+ forM [0 .. extent0 - 1] $ \i0 -> do+ let ptr1 = ptr0 `advancePtr` fromIntegral (min0 + stride0 * i0)+ forM [0 .. extent1 - 1] $ \i1 -> do+ let ptr2 = ptr1 `advancePtr` fromIntegral (min1 + stride1 * i1)+ forM [0 .. extent2 - 1] $ \i2 ->+ peekElemOff ptr2 (fromIntegral (min2 + stride2 * i2))
+ src/Language/Halide/Context.hs view
@@ -0,0 +1,139 @@+{-# LANGUAGE TemplateHaskellQuotes #-}++-- |+-- Module : Language.Halide.Context+-- Description : Helpers to setup inline-c for Halide+-- Copyright : (c) Tom Westerhout, 2023+--+-- This module defines a Template Haskell function 'importHalide' that sets up everything you need+-- to call Halide functions from 'Language.C.Inline' and 'Language.C.Inlinde.Cpp' quasiquotes.+--+-- We also define two C++ functions:+--+-- > template <class Func>+-- > auto handle_halide_exceptions(Func&& func);+-- >+-- > template <class T>+-- > auto to_string_via_iostream(T const& x) -> std::string*;+--+-- @handle_halide_exceptions@ can be used to catch various Halide exceptions and convert them to+-- [@std::runtime_error@](https://en.cppreference.com/w/cpp/error/runtime_error). It can be used+-- inside 'C.tryBlock' or 'C.catchBlock' to properly re-throw Halide errors.+--+-- @+-- [C.catchBlock| void {+-- handle_halide_exceptions([=]() {+-- Halide::Func f;+-- Halide::Var i;+-- f(i) = *$(Halide::Expr* e);+-- f.realize(Halide::Pipeline::RealizationArg{$(halide_buffer_t* b)});+-- });+-- } |]+-- @+--+-- @to_string_via_iostream@ is a helper that converts a variable into a string by relying on+-- [iostreams](https://en.cppreference.com/w/cpp/io). It returns a pointer to+-- [@std::string@](https://en.cppreference.com/w/cpp/string/basic_string) that it allocated using the @new@+-- keyword. To convert it to a Haskell string, use the 'Language.Halide.Utils.peekCxxString' and+-- 'Language.Halide.Utils.peekAndDeleteCxxString' functions.+module Language.Halide.Context+ ( importHalide+ )+where++import qualified Language.C.Inline as C+import qualified Language.C.Inline.Cpp as C+import Language.C.Types (CIdentifier)+import Language.Halide.Type+import Language.Haskell.TH (DecsQ, Q, TypeQ, lookupTypeName)+import qualified Language.Haskell.TH as TH++-- | One stop function to include all the neccessary machinery to call Halide functions via inline-c.+--+-- Put @importHalide@ somewhere at the beginning of the file and enjoy using the C++ interface of+-- Halide via inline-c quasiquotes.+importHalide :: DecsQ+importHalide =+ concat+ <$> sequence+ [ C.context =<< halideCxt+ , C.include "<Halide.h>"+ , C.include "<cxxabi.h>"+ , C.include "<dlfcn.h>"+ , defineExceptionHandler+ ]++halideCxt :: Q C.Context+halideCxt = do+ typePairs <- C.cppTypePairs <$> halideTypePairs+ pure (C.cppCtx <> C.fptrCtx <> C.bsCtx <> typePairs)++halideTypePairs :: Q [(CIdentifier, TypeQ)]+halideTypePairs = do+ fmap concat . sequence $ [core, other]+ where+ core =+ pure+ [ ("Halide::Expr", [t|CxxExpr|])+ , ("Halide::Var", [t|CxxVar|])+ , ("Halide::RVar", [t|CxxRVar|])+ , ("Halide::VarOrRVar", [t|CxxVarOrRVar|])+ , ("Halide::Func", [t|CxxFunc|])+ , ("Halide::Internal::Parameter", [t|CxxParameter|])+ , ("Halide::ImageParam", [t|CxxImageParam|])+ , ("Halide::Callable", [t|CxxCallable|])+ , ("Halide::Target", [t|CxxTarget|])+ , ("Halide::JITUserContext", [t|CxxUserContext|])+ , ("Halide::Argument", [t|CxxArgument|])+ , ("std::vector", [t|CxxVector|])+ , ("std::string", [t|CxxString|])+ , ("halide_type_t", [t|HalideType|])+ ]+ other =+ optionals+ [ ("Halide::Internal::StageSchedule", "CxxStageSchedule")+ , ("Halide::Internal::Dim", "Language.Halide.Schedule.Dim")+ , ("Halide::Internal::Split", "Language.Halide.Schedule.Split")+ , ("halide_buffer_t", "Language.Halide.Buffer.RawHalideBuffer")+ , ("Halide::Internal::Dimension", "CxxDimension")+ , ("Halide::LoopLevel", "CxxLoopLevel")+ , ("Halide::Stage", "CxxStage")+ , ("Halide::Buffer", "CxxBuffer")+ , ("Halide::Internal::FusedPair", "FusedPair")+ , ("Halide::Internal::ReductionVariable", "ReductionVariable")+ , ("Halide::Internal::PrefetchDirective", "PrefetchDirective")+ , ("halide_trace_event_t", "TraceEvent")+ ]+ optional :: (CIdentifier, String) -> Q [(CIdentifier, TypeQ)]+ optional (cName, hsName) = do+ hsType <- lookupTypeName hsName+ pure $ maybe [] (\x -> [(cName, pure (TH.ConT x))]) hsType+ optionals :: [(CIdentifier, String)] -> Q [(CIdentifier, TypeQ)]+ optionals pairs = concat <$> mapM optional pairs++defineExceptionHandler :: DecsQ+defineExceptionHandler =+ C.verbatim+ "\+ \template <class Func> \n\+ \auto handle_halide_exceptions(Func&& func) { \n\+ \ try { \n\+ \ return func(); \n\+ \ } catch(Halide::RuntimeError& e) { \n\+ \ throw std::runtime_error{e.what()}; \n\+ \ } catch(Halide::CompileError& e) { \n\+ \ throw std::runtime_error{e.what()}; \n\+ \ } catch(Halide::InternalError& e) { \n\+ \ throw std::runtime_error{e.what()}; \n\+ \ } catch(Halide::Error& e) { \n\+ \ throw std::runtime_error{e.what()}; \n\+ \ } \n\+ \} \n\+ \ \n\+ \template <class T> \n\+ \auto to_string_via_iostream(T const& x) -> std::string* { \n\+ \ std::ostringstream stream; \n\+ \ stream << x; \n\+ \ return new std::string{stream.str()}; \n\+ \} \n\+ \"
+ src/Language/Halide/Dimension.hs view
@@ -0,0 +1,153 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_GHC -Wno-redundant-constraints #-}++-- |+-- Module : Language.Halide.Dimension+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Dimension+ ( Dimension (..)+ , setMin+ , setExtent+ , setStride+ , setEstimate++ -- * Internal+ , CxxDimension+ , wrapCxxDimension+ , withCxxDimension+ )+where++import Foreign.ForeignPtr+import Foreign.Ptr (Ptr)+import GHC.Records (HasField (..))+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Expr+import Language.Halide.Type+import System.IO.Unsafe (unsafePerformIO)+import Prelude hiding (tail)++-- | Haskell counterpart of [@Halide::Internal::Dimension@](https://halide-lang.org/docs/class_halide_1_1_internal_1_1_dimension.html).+data CxxDimension++importHalide++-- | Information about a buffer's dimension, such as the min, extent, and stride.+newtype Dimension = Dimension (ForeignPtr CxxDimension)++instance Show Dimension where+ showsPrec d dim =+ showParen (d > 10) $+ showString "Dimension { min="+ . shows dim.min+ . showString (", extent=" :: String)+ . shows dim.extent+ . showString (", stride=" :: String)+ . shows dim.stride+ . showString " }"++instance HasField "min" Dimension (Expr Int32) where+ getField :: Dimension -> Expr Int32+ getField dim = unsafePerformIO $+ withCxxDimension dim $ \d ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::Dimension* d)->min()} } |]++-- | Set the min in a given dimension to equal the given expression. Setting the mins to+-- zero may simplify some addressing math.+--+-- For more info, see [Halide::Internal::Dimension::set_min](https://halide-lang.org/docs/class_halide_1_1_internal_1_1_dimension.html#a84acaf7733391fdaea4f4cec24a60de2).+setMin :: Expr Int32 -> Dimension -> IO Dimension+setMin expr dim = do+ asExpr expr $ \n ->+ withCxxDimension dim $ \d ->+ [CU.exp| void {+ $(Halide::Internal::Dimension* d)->set_min(*$(const Halide::Expr* n)) } |]+ pure dim++instance HasField "extent" Dimension (Expr Int32) where+ getField :: Dimension -> Expr Int32+ getField dim = unsafePerformIO $+ withCxxDimension dim $ \d ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::Dimension* d)->extent()} } |]++-- | Set the extent in a given dimension to equal the given expression.+--+-- Halide will generate runtime errors for Buffers that fail this check.+--+-- For more info, see [Halide::Internal::Dimension::set_extent](https://halide-lang.org/docs/class_halide_1_1_internal_1_1_dimension.html#a54111d8439a065bdaca5b9ff9bcbd630).+setExtent :: Expr Int32 -> Dimension -> IO Dimension+setExtent expr dim = do+ asExpr expr $ \n ->+ withCxxDimension dim $ \d ->+ [CU.exp| void {+ $(Halide::Internal::Dimension* d)->set_extent(*$(const Halide::Expr* n)) } |]+ pure dim++instance HasField "max" Dimension (Expr Int32) where+ getField :: Dimension -> Expr Int32+ getField dim = unsafePerformIO $+ withCxxDimension dim $ \d ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(Halide::Internal::Dimension* d)->max()} } |]++instance HasField "stride" Dimension (Expr Int32) where+ getField :: Dimension -> Expr Int32+ getField dim = unsafePerformIO $+ withCxxDimension dim $ \d ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(Halide::Internal::Dimension* d)->stride()} } |]++-- | Set the stride in a given dimension to equal the given expression.+--+-- This is particularly useful to set when vectorizing. Known strides for the vectorized+-- dimensions generate better code.+--+-- For more info, see [Halide::Internal::Dimension::set_stride](https://halide-lang.org/docs/class_halide_1_1_internal_1_1_dimension.html#a94f4c432a89907e2cc2aa908b5012cf8).+setStride :: Expr Int32 -> Dimension -> IO Dimension+setStride expr dim = do+ asExpr expr $ \n ->+ withCxxDimension dim $ \d ->+ [CU.exp| void {+ $(Halide::Internal::Dimension* d)->set_stride(*$(const Halide::Expr* n)) } |]+ pure dim++-- | Set estimates for autoschedulers.+setEstimate+ :: Expr Int32+ -- ^ @min@ estimate+ -> Expr Int32+ -- ^ @extent@ estimate+ -> Dimension+ -> IO Dimension+setEstimate minExpr extentExpr dim = do+ asExpr minExpr $ \m ->+ asExpr extentExpr $ \e ->+ withCxxDimension dim $ \d ->+ [CU.exp| void {+ $(Halide::Internal::Dimension* d)->set_estimate(*$(const Halide::Expr* m),+ *$(const Halide::Expr* e)) } |]+ pure dim++wrapCxxDimension :: Ptr CxxDimension -> IO Dimension+wrapCxxDimension = fmap Dimension . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteDimension(Halide::Internal::Dimension* p) { delete p; } |]++withCxxDimension :: Dimension -> (Ptr CxxDimension -> IO a) -> IO a+withCxxDimension (Dimension fp) = withForeignPtr fp
+ src/Language/Halide/Expr.hs view
@@ -0,0 +1,657 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_GHC -Wno-orphans #-}++-- |+-- Module : Language.Halide.Expr+-- Description : Scalar expressions+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Expr+ ( Expr (..)+ , Var+ , RVar+ , VarOrRVar+ , Int32+ , mkExpr+ , mkVar+ , mkRVar+ , cast+ , eq+ , neq+ , lt+ , lte+ , gt+ , gte+ , bool+ , undef+ -- | For debugging, it's often useful to observe the value of an expression when it's evaluated. If you+ -- have a complex expression that does not depend on any buffers or indices, you can 'evaluate' it.+ , evaluate+ -- | However, often an expression is only used within a definition of a pipeline, and it's impossible to+ -- call 'evaluate' on it. In such cases, it can be wrapped with 'printed' to indicate to Halide that the+ -- value of the expression should be dumped to screen when it's computed.+ , printed+ , toIntImm++ -- * Internal+ , exprToForeignPtr+ , cxxConstructExpr+ -- , wrapCxxExpr+ , wrapCxxRVar+ , wrapCxxVarOrRVar+ , wrapCxxParameter+ , asExpr+ , asVar+ , asRVar+ , asVarOrRVar+ , asScalarParam+ , asVectorOf+ , mkScalarParameter+ , withMany+ , binaryOp+ , unaryOp+ , checkType+ )+where++import Control.Exception (bracket)+import Control.Monad (unless)+import Data.IORef+import Data.Int (Int32)+import Data.Proxy+import Data.Ratio (denominator, numerator)+import Data.Text (Text, unpack)+import Data.Text.Encoding qualified as T+import Data.Vector.Storable.Mutable qualified as SM+import Foreign.ForeignPtr+import Foreign.Marshal (alloca, allocaArray, peekArray, toBool, with)+import Foreign.Ptr (Ptr, castPtr, nullPtr)+import Foreign.Storable (peek)+import GHC.Stack (HasCallStack)+import Language.C.Inline qualified as C+import Language.C.Inline.Cpp.Exception qualified as C+import Language.C.Inline.Unsafe qualified as CU+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Type+import Language.Halide.Utils+import System.IO.Unsafe (unsafePerformIO)+import Prelude hiding (min)++importHalide++instanceCxxConstructible "Halide::Expr"+instanceCxxConstructible "Halide::Var"+instanceCxxConstructible "Halide::RVar"+instanceCxxConstructible "Halide::VarOrRVar"++defineIsHalideTypeInstances++instanceHasCxxVector "Halide::Expr"+instanceHasCxxVector "Halide::Var"+instanceHasCxxVector "Halide::RVar"+instanceHasCxxVector "Halide::VarOrRVar"++-- instanceCxxConstructible "Halide::Var"+-- instanceCxxConstructible "Halide::RVar"+-- instanceCxxConstructible "Halide::VarOrRVar"++instance IsHalideType Bool where+ halideTypeFor _ = HalideType HalideTypeUInt 1 1+ toCxxExpr (fromIntegral . fromEnum -> x) =+ cxxConstruct $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{cast(Halide::UInt(1), Halide::Expr{$(int x)})} } |]++type instance FromTuple (Expr a) = Arguments '[Expr a]++-- | A scalar expression in Halide.+--+-- To have a nice experience writing arithmetic expressions in terms of @Expr@s, we want to derive 'Num',+-- 'Floating' etc. instances for @Expr@. Unfortunately, that means that we encode v'Expr', v'Var', v'RVar',+-- and v'ScalarParam' by the same type, and passing an @Expr@ to a function that expects a @Var@ will produce+-- a runtime error.+data Expr a+ = -- | Scalar expression.+ Expr (ForeignPtr CxxExpr)+ | -- | Index variable.+ Var (ForeignPtr CxxVar)+ | -- | Reduction variable.+ RVar (ForeignPtr CxxRVar)+ | -- | Scalar parameter.+ --+ -- The 'IORef' is initialized with 'Nothing' and filled in on the first+ -- call to 'asExpr'.+ ScalarParam (IORef (Maybe (ForeignPtr CxxParameter)))++-- | A v'Var'.+type Var = Expr Int32++-- | An v'RVar'.+type RVar = Expr Int32++-- | Either v'Var' or v'RVar'.+type VarOrRVar = Expr Int32++-- | Create a scalar expression from a Haskell value.+mkExpr :: IsHalideType a => a -> Expr a+mkExpr x = unsafePerformIO $! Expr <$> toCxxExpr x++-- | Create a named index variable.+mkVar :: Text -> IO (Expr Int32)+mkVar (T.encodeUtf8 -> s) = fmap Var . cxxConstruct $ \ptr ->+ [CU.exp| void {+ new ($(Halide::Var* ptr)) Halide::Var{std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}} } |]++-- | Create a named reduction variable.+--+-- For more information about reduction variables, see [@Halide::RDom@](https://halide-lang.org/docs/class_halide_1_1_r_dom.html).+mkRVar+ :: Text+ -- ^ name+ -> Expr Int32+ -- ^ min index+ -> Expr Int32+ -- ^ extent+ -> IO (Expr Int32)+mkRVar name min extent =+ asExpr min $ \min' ->+ asExpr extent $ \extent' ->+ wrapCxxRVar+ =<< [CU.exp| Halide::RVar* {+ new Halide::RVar{static_cast<Halide::RVar>(Halide::RDom{+ *$(const Halide::Expr* min'),+ *$(const Halide::Expr* extent'),+ std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}+ })}+ } |]+ where+ s = T.encodeUtf8 name++-- | Return an undef value of the given type.+--+-- For more information, see [@Halide::undef@](https://halide-lang.org/docs/namespace_halide.html#a9389bcacbed602df70eae94826312e03).+undef :: forall a. IsHalideType a => Expr a+undef = unsafePerformIO $+ with (halideTypeFor (Proxy @a)) $ \tp ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void {+ new ($(Halide::Expr* ptr))+ Halide::Expr{Halide::undef(Halide::Type{*$(const halide_type_t* tp)})} } |]+{-# NOINLINE undef #-}++-- | Cast a scalar expression to a different type.+--+-- Use TypeApplications with this function, e.g. @cast \@Float x@.+cast :: forall to from. (IsHalideType to, IsHalideType from) => Expr from -> Expr to+cast expr = unsafePerformIO $+ asExpr expr $ \e ->+ with (halideTypeFor (Proxy @to)) $ \t ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ Halide::cast(Halide::Type{*$(halide_type_t* t)}, *$(Halide::Expr* e))} } |]++-- | Print the expression to stdout when it's evaluated.+--+-- This is useful for debugging Halide pipelines.+printed :: IsHalideType a => Expr a -> Expr a+printed = unaryOp $ \e ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{print(*$(Halide::Expr* e))} } |]++infix 4 `eq`, `neq`, `lt`, `lte`, `gt`, `gte`++-- | '==' but lifted to return an 'Expr'.+eq :: IsHalideType a => Expr a -> Expr a -> Expr Bool+eq = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) == (*$(Halide::Expr* b))} } |]++-- | '/=' but lifted to return an 'Expr'.+neq :: IsHalideType a => Expr a -> Expr a -> Expr Bool+neq = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) != (*$(Halide::Expr* b))} } |]++-- | '<' but lifted to return an 'Expr'.+lt :: IsHalideType a => Expr a -> Expr a -> Expr Bool+lt = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) < (*$(Halide::Expr* b))} } |]++-- | '<=' but lifted to return an 'Expr'.+lte :: IsHalideType a => Expr a -> Expr a -> Expr Bool+lte = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) <= (*$(Halide::Expr* b))} } |]++-- | '>' but lifted to return an 'Expr'.+gt :: IsHalideType a => Expr a -> Expr a -> Expr Bool+gt = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) > (*$(Halide::Expr* b))} } |]++-- | '>=' but lifted to return an 'Expr'.+gte :: IsHalideType a => Expr a -> Expr a -> Expr Bool+gte = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ (*$(Halide::Expr* a)) >= (*$(Halide::Expr* b))} } |]++-- | Similar to the standard 'Prelude.bool' function from Prelude except that it's+-- lifted to work with 'Expr' types.+bool :: IsHalideType a => Expr Bool -> Expr a -> Expr a -> Expr a+bool condExpr trueExpr falseExpr = unsafePerformIO $+ asExpr condExpr $ \p ->+ asExpr trueExpr $ \t ->+ asExpr falseExpr $ \f ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void {+ new ($(Halide::Expr* ptr)) Halide::Expr{+ Halide::select(*$(Halide::Expr* p),+ *$(Halide::Expr* t), *$(Halide::Expr* f))} } |]++-- | Evaluate a scalar expression.+--+-- It should contain no parameters. If it does contain parameters, an exception will be thrown.+evaluate :: forall a. IsHalideType a => Expr a -> IO a+evaluate expr =+ asExpr expr $ \e -> do+ out <- SM.new 1+ withHalideBuffer out $ \buffer -> do+ let b = castPtr (buffer :: Ptr (HalideBuffer 1 a))+ [C.throwBlock| void {+ handle_halide_exceptions([=]() {+ Halide::Func f;+ Halide::Var i;+ f(i) = *$(Halide::Expr* e);+ f.realize(Halide::Pipeline::RealizationArg{$(halide_buffer_t* b)});+ });+ } |]+ SM.read out 0++-- | Convert expression to integer immediate.+--+-- Tries to extract the value of an expression if it is a compile-time constant. If the expression+-- isn't known at compile-time of the Halide pipeline, returns 'Nothing'.+toIntImm :: IsHalideType a => Expr a -> Maybe Int+toIntImm expr = unsafePerformIO $+ asExpr expr $ \expr' -> do+ intPtr <-+ [CU.block| const int64_t* {+ auto expr = *$(const Halide::Expr* expr');+ Halide::Internal::IntImm const* node = expr.as<Halide::Internal::IntImm>();+ if (node == nullptr) return nullptr;+ return &node->value;+ } |]+ if intPtr == nullPtr+ then pure Nothing+ else Just . fromIntegral <$> peek intPtr++instance IsTuple (Arguments '[Expr a]) (Expr a) where+ toTuple (x ::: Nil) = x+ fromTuple x = x ::: Nil++instance IsHalideType a => Show (Expr a) where+ show (Expr expr) = unpack . unsafePerformIO $ do+ withForeignPtr expr $ \x ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* { to_string_via_iostream(*$(const Halide::Expr* x)) } |]+ show (Var var) = unpack . unsafePerformIO $ do+ withForeignPtr var $ \x ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* { to_string_via_iostream(*$(const Halide::Var* x)) } |]+ show (RVar rvar) = unpack . unsafePerformIO $ do+ withForeignPtr rvar $ \x ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* { to_string_via_iostream(*$(const Halide::RVar* x)) } |]+ show (ScalarParam r) = unpack . unsafePerformIO $ do+ maybeParam <- readIORef r+ case maybeParam of+ Just fp ->+ withForeignPtr fp $ \x ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* {+ new std::string{$(const Halide::Internal::Parameter* x)->name()} } |]+ Nothing -> pure "ScalarParam"++instance (IsHalideType a, Num a) => Num (Expr a) where+ fromInteger :: Integer -> Expr a+ fromInteger x = mkExpr (fromInteger x :: a)+ (+) :: Expr a -> Expr a -> Expr a+ (+) = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{*$(Halide::Expr* a) + *$(Halide::Expr* b)} } |]+ (-) :: Expr a -> Expr a -> Expr a+ (-) = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{*$(Halide::Expr* a) - *$(Halide::Expr* b)} } |]+ (*) :: Expr a -> Expr a -> Expr a+ (*) = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{*$(Halide::Expr* a) * *$(Halide::Expr* b)} } |]++ abs :: Expr a -> Expr a+ abs = unaryOp $ \a ptr ->+ -- If the type is unsigned, then abs does nothing Also note that for signed+ -- integers, in Halide abs returns the unsigned version, so we manually+ -- cast it back.+ [CU.block| void {+ if ($(Halide::Expr* a)->type().is_uint()) {+ new ($(Halide::Expr* ptr)) Halide::Expr{*$(Halide::Expr* a)};+ }+ else {+ new ($(Halide::Expr* ptr)) Halide::Expr{+ Halide::cast($(Halide::Expr* a)->type(), Halide::abs(*$(Halide::Expr* a)))};+ }+ } |]+ negate :: Expr a -> Expr a+ negate = unaryOp $ \a ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{ -(*$(Halide::Expr* a))} } |]+ signum :: Expr a -> Expr a+ signum = error "Num instance of (Expr a) does not implement signum"++instance (IsHalideType a, Fractional a) => Fractional (Expr a) where+ (/) :: Expr a -> Expr a -> Expr a+ (/) = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{*$(Halide::Expr* a) / *$(Halide::Expr* b)} } |]+ fromRational :: Rational -> Expr a+ fromRational r = fromInteger (numerator r) / fromInteger (denominator r)++instance (IsHalideType a, Floating a) => Floating (Expr a) where+ pi :: Expr a+ pi = cast @a @Double $! mkExpr (pi :: Double)+ exp :: Expr a -> Expr a+ exp = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::exp(*$(Halide::Expr* a))} } |]+ log :: Expr a -> Expr a+ log = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::log(*$(Halide::Expr* a))} } |]+ sqrt :: Expr a -> Expr a+ sqrt = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::sqrt(*$(Halide::Expr* a))} } |]+ (**) :: Expr a -> Expr a -> Expr a+ (**) = binaryOp $ \a b ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::pow(*$(Halide::Expr* a), *$(Halide::Expr* b))} } |]+ sin :: Expr a -> Expr a+ sin = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::sin(*$(Halide::Expr* a))} } |]+ cos :: Expr a -> Expr a+ cos = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::cos(*$(Halide::Expr* a))} } |]+ tan :: Expr a -> Expr a+ tan = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::tan(*$(Halide::Expr* a))} } |]+ asin :: Expr a -> Expr a+ asin = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::asin(*$(Halide::Expr* a))} } |]+ acos :: Expr a -> Expr a+ acos = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::acos(*$(Halide::Expr* a))} } |]+ atan :: Expr a -> Expr a+ atan = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::atan(*$(Halide::Expr* a))} } |]+ sinh :: Expr a -> Expr a+ sinh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::sinh(*$(Halide::Expr* a))} } |]+ cosh :: Expr a -> Expr a+ cosh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::cosh(*$(Halide::Expr* a))} } |]+ tanh :: Expr a -> Expr a+ tanh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::tanh(*$(Halide::Expr* a))} } |]+ asinh :: Expr a -> Expr a+ asinh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::asinh(*$(Halide::Expr* a))} } |]+ acosh :: Expr a -> Expr a+ acosh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::acosh(*$(Halide::Expr* a))} } |]+ atanh :: Expr a -> Expr a+ atanh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::atanh(*$(Halide::Expr* a))} } |]++-- | Wrap a raw @Halide::Expr@ pointer in a Haskell value.+--+-- __Note:__ This function checks the runtime type of the expression.+-- wrapCxxExpr :: forall a. (HasCallStack, IsHalideType a) => Ptr CxxExpr -> IO (Expr a)+-- wrapCxxExpr p = do+-- checkType @a p+-- Expr <$> newForeignPtr deleter p+-- where+-- deleter = [C.funPtr| void deleteExpr(Halide::Expr *p) { delete p; } |]+cxxConstructExpr :: forall a. (HasCallStack, IsHalideType a) => (Ptr CxxExpr -> IO ()) -> IO (Expr a)+cxxConstructExpr construct = do+ fp <- cxxConstruct construct+ withForeignPtr fp (checkType @a)+ pure (Expr fp)++-- | Wrap a raw @Halide::RVar@ pointer in a Haskell value.+--+-- __Note:__ v'RVar' objects correspond to expressions of type 'Int32'.+wrapCxxRVar :: Ptr CxxRVar -> IO (Expr Int32)+wrapCxxRVar = fmap RVar . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteExpr(Halide::RVar *p) { delete p; } |]++wrapCxxVarOrRVar :: Ptr CxxVarOrRVar -> IO (Expr Int32)+wrapCxxVarOrRVar p = do+ isRVar <- toBool <$> [CU.exp| bool { $(const Halide::VarOrRVar* p)->is_rvar } |]+ expr <-+ if isRVar+ then wrapCxxRVar =<< [CU.exp| Halide::RVar* { new Halide::RVar{$(const Halide::VarOrRVar* p)->rvar} } |]+ else fmap Var . cxxConstruct $ \ptr ->+ [CU.exp| void { new ($(Halide::Var* ptr)) Halide::Var{$(const Halide::VarOrRVar* p)->var} } |]+ [CU.exp| void { delete $(const Halide::VarOrRVar* p) } |]+ pure expr++class HasHalideType a where+ getHalideType :: a -> IO HalideType++instance HasHalideType (Expr a) where+ getHalideType (Expr fp) =+ withForeignPtr fp $ \e -> alloca $ \t -> do+ [CU.block| void {+ *$(halide_type_t* t) = static_cast<halide_type_t>(+ $(Halide::Expr* e)->type()); } |]+ peek t+ getHalideType (Var fp) =+ withForeignPtr fp $ \e -> alloca $ \t -> do+ [CU.block| void {+ *$(halide_type_t* t) = static_cast<halide_type_t>(+ static_cast<Halide::Expr>(*$(Halide::Var* e)).type()); } |]+ peek t+ getHalideType (RVar fp) =+ withForeignPtr fp $ \e -> alloca $ \t -> do+ [CU.block| void {+ *$(halide_type_t* t) = static_cast<halide_type_t>(+ static_cast<Halide::Expr>(*$(Halide::RVar* e)).type()); } |]+ peek t+ getHalideType _ = error "not implemented"++instance HasHalideType (Ptr CxxExpr) where+ getHalideType e =+ alloca $ \t -> do+ [CU.block| void {+ *$(halide_type_t* t) = static_cast<halide_type_t>($(Halide::Expr* e)->type()); } |]+ peek t++instance HasHalideType (Ptr CxxVar) where+ getHalideType _ = pure $ halideTypeFor (Proxy @Int32)++instance HasHalideType (Ptr CxxRVar) where+ getHalideType _ = pure $ halideTypeFor (Proxy @Int32)++instance HasHalideType (Ptr CxxParameter) where+ getHalideType p =+ alloca $ \t -> do+ [CU.block| void {+ *$(halide_type_t* t) = static_cast<halide_type_t>($(Halide::Internal::Parameter* p)->type()); } |]+ peek t++-- | Wrap a raw @Halide::Internal::Parameter@ pointer in a Haskell value.+--+-- __Note:__ v'Var' objects correspond to expressions of type 'Int32'.+wrapCxxParameter :: Ptr CxxParameter -> IO (ForeignPtr CxxParameter)+wrapCxxParameter = newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteParameter(Halide::Internal::Parameter *p) { delete p; } |]++-- | Helper function to assert that the runtime type of the expression matches it's+-- compile-time type.+--+-- Essentially, given an @(x :: 'Expr' a)@, we check that @x.type()@ in C++ is equal to+-- @'halideTypeFor' (Proxy \@a)@ in Haskell.+checkType :: forall a t. (HasCallStack, IsHalideType a, HasHalideType t) => t -> IO ()+checkType x = do+ let hsType = halideTypeFor (Proxy @a)+ cxxType <- getHalideType x+ unless (cxxType == hsType) . error $+ "Type mismatch: C++ Expr has type "+ <> show cxxType+ <> ", but its Haskell counterpart has type "+ <> show hsType++mkScalarParameter :: forall a. IsHalideType a => Maybe Text -> IO (ForeignPtr CxxParameter)+mkScalarParameter maybeName = do+ with (halideTypeFor (Proxy @a)) $ \t -> do+ let createWithoutName =+ [CU.exp| Halide::Internal::Parameter* {+ new Halide::Internal::Parameter{Halide::Type{*$(halide_type_t* t)}, false, 0} } |]+ createWithName name =+ let s = T.encodeUtf8 name+ in [CU.exp| Halide::Internal::Parameter* {+ new Halide::Internal::Parameter{+ Halide::Type{*$(halide_type_t* t)},+ false,+ 0,+ std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}}+ } |]+ p <- maybe createWithoutName createWithName maybeName+ checkType @a p+ wrapCxxParameter p++getScalarParameter+ :: forall a+ . IsHalideType a+ => Maybe Text+ -> IORef (Maybe (ForeignPtr CxxParameter))+ -> IO (ForeignPtr CxxParameter)+getScalarParameter name r = do+ readIORef r >>= \case+ Just fp -> pure fp+ Nothing -> do+ fp <- mkScalarParameter @a name+ writeIORef r (Just fp)+ pure fp++-- | Make sure that the expression is fully constructed. That means that if we+-- are dealing with a 'ScalarParam' rather than an 'Expr', we force the+-- construction of the underlying @Halide::Internal::Parameter@ and convert it+-- to an 'Expr'.+forceExpr :: forall a. (HasCallStack, IsHalideType a) => Expr a -> IO (Expr a)+forceExpr x@(Expr _) = pure x+forceExpr (Var fp) =+ withForeignPtr fp $ \varPtr ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ static_cast<Halide::Expr>(*$(Halide::Var* varPtr))} } |]+forceExpr (RVar fp) =+ withForeignPtr fp $ \rvarPtr ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ static_cast<Halide::Expr>(*$(Halide::RVar* rvarPtr))} } |]+forceExpr (ScalarParam r) =+ getScalarParameter @a Nothing r >>= \fp -> withForeignPtr fp $ \paramPtr ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ Halide::Internal::Variable::make(+ $(Halide::Internal::Parameter* paramPtr)->type(),+ $(Halide::Internal::Parameter* paramPtr)->name(),+ *$(Halide::Internal::Parameter* paramPtr))} } |]++-- | Use the underlying @Halide::Expr@ in an 'IO' action.+asExpr :: IsHalideType a => Expr a -> (Ptr CxxExpr -> IO b) -> IO b+asExpr x = withForeignPtr (exprToForeignPtr x)++-- | Allows applying 'asExpr', 'asVar', 'asRVar', and 'asVarOrRVar' to multiple arguments.+--+-- Example usage:+--+-- > asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple args) $ \v -> do+-- > withFunc func $ \f ->+-- > [C.throwBlock| void { $(Halide::Func* f)->reorder(+-- > *$(std::vector<Halide::VarOrRVar>* v)); } |]+asVectorOf+ :: forall c k ts a+ . (All c ts, HasCxxVector k)+ => (forall t b. c t => t -> (Ptr k -> IO b) -> IO b)+ -> Arguments ts+ -> (Ptr (CxxVector k) -> IO a)+ -> IO a+asVectorOf asPtr args action =+ bracket (newCxxVector Nothing) deleteCxxVector (go args)+ where+ go+ :: All c ts'+ => Arguments ts'+ -> Ptr (CxxVector k)+ -> IO a+ go Nil v = action v+ go (x ::: xs) v = asPtr x $ \p -> cxxVectorPushBack v p >> go xs v++withMany+ :: forall k t a+ . (HasCxxVector k)+ => (t -> (Ptr k -> IO a) -> IO a)+ -> [t]+ -> (Ptr (CxxVector k) -> IO a)+ -> IO a+withMany asPtr args action =+ bracket (newCxxVector Nothing) deleteCxxVector (go args)+ where+ go [] v = action v+ go (x : xs) v = asPtr x $ \p -> cxxVectorPushBack v p >> go xs v++-- | Use the underlying @Halide::Var@ in an 'IO' action.+asVar :: HasCallStack => Expr Int32 -> (Ptr CxxVar -> IO b) -> IO b+asVar (Var fp) = withForeignPtr fp+asVar _ = error "the expression is not a Var"++-- | Use the underlying @Halide::RVar@ in an 'IO' action.+asRVar :: HasCallStack => Expr Int32 -> (Ptr CxxRVar -> IO b) -> IO b+asRVar (RVar fp) = withForeignPtr fp+asRVar _ = error "the expression is not an RVar"++-- | Use the underlying v'Var' or v'RVar' as @Halide::VarOrRVar@ in an 'IO' action.+asVarOrRVar :: HasCallStack => VarOrRVar -> (Ptr CxxVarOrRVar -> IO b) -> IO b+asVarOrRVar x action = case x of+ Var fp ->+ let allocate p = [CU.exp| Halide::VarOrRVar* { new Halide::VarOrRVar{*$(Halide::Var* p)} } |]+ in withForeignPtr fp (run . allocate)+ RVar fp ->+ let allocate p = [CU.exp| Halide::VarOrRVar* { new Halide::VarOrRVar{*$(Halide::RVar* p)} } |]+ in withForeignPtr fp (run . allocate)+ _ -> error "the expression is not a Var or an RVar"+ where+ destroy p = [CU.exp| void { delete $(Halide::VarOrRVar* p) } |]+ run allocate = bracket allocate destroy action++-- | Use the underlying @Halide::RVar@ in an 'IO' action.+asScalarParam :: forall a b. (HasCallStack, IsHalideType a) => Expr a -> (Ptr CxxParameter -> IO b) -> IO b+asScalarParam (ScalarParam r) action = do+ fp <- getScalarParameter @a Nothing r+ withForeignPtr fp action+asScalarParam _ _ = error "the expression is not a ScalarParam"++-- | Get the underlying 'ForeignPtr CxxExpr'.+exprToForeignPtr :: IsHalideType a => Expr a -> ForeignPtr CxxExpr+exprToForeignPtr x =+ unsafePerformIO $!+ forceExpr x >>= \case+ (Expr fp) -> pure fp+ _ -> error "this cannot happen"++-- | Lift a unary function working with @Halide::Expr@ to work with 'Expr'.+unaryOp :: IsHalideType a => (Ptr CxxExpr -> Ptr CxxExpr -> IO ()) -> Expr a -> Expr a+unaryOp f a = unsafePerformIO $+ asExpr a $ \aPtr ->+ cxxConstructExpr $ \destPtr ->+ f aPtr destPtr++-- | Lift a binary function working with @Halide::Expr@ to work with 'Expr'.+binaryOp+ :: (IsHalideType a, IsHalideType b, IsHalideType c)+ => (Ptr CxxExpr -> Ptr CxxExpr -> Ptr CxxExpr -> IO ())+ -> Expr a+ -> Expr b+ -> Expr c+binaryOp f a b = unsafePerformIO $+ asExpr a $ \aPtr -> asExpr b $ \bPtr ->+ cxxConstructExpr $ \destPtr ->+ f aPtr bPtr destPtr
+ src/Language/Halide/Func.hs view
@@ -0,0 +1,1143 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE PolyKinds #-}+-- {-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_GHC -Wno-redundant-constraints #-}++-- |+-- Module : Language.Halide.Func+-- Description : Functions / Arrays+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Func+ ( -- * Defining pipelines+ Func (..)+ , FuncTy (..)+ , Stage (..)+ , buffer+ , scalar+ , define+ , (!)+ , realize++ -- * Scheduling+ , Schedulable (..)+ , TailStrategy (..)++ -- ** 'Func'-specific+ , computeRoot+ , getStage+ , getLoopLevel+ , getLoopLevelAtStage+ , asUsed+ , asUsedBy+ , copyToDevice+ , copyToHost+ , storeAt+ , computeAt+ , dim+ , estimate+ , bound+ , getArgs+ -- , deepCopy++ -- * Update definitions+ , update+ , hasUpdateDefinitions+ , getUpdateStage++ -- * Debugging+ , prettyLoopNest++ -- * Internal+ , IndexTuple+ , asBufferParam+ , withFunc+ , withBufferParam+ , wrapCxxFunc+ , CxxStage+ , wrapCxxStage+ , withCxxStage+ )+where++import Control.Exception (bracket)+import Control.Monad (forM)+import Data.IORef+import Data.Kind (Type)+import Data.Proxy+import Data.Text (Text)+import Data.Text.Encoding qualified as T+import Foreign.ForeignPtr+import Foreign.Marshal (toBool, with)+import Foreign.Ptr (Ptr, castPtr)+import GHC.Stack (HasCallStack)+import GHC.TypeLits+import Language.C.Inline qualified as C+import Language.C.Inline.Cpp.Exception qualified as C+import Language.C.Inline.Unsafe qualified as CU+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Dimension+import Language.Halide.Expr+import Language.Halide.LoopLevel+import Language.Halide.Target+import Language.Halide.Type+import Language.Halide.Utils+import System.IO.Unsafe (unsafePerformIO)+import Prelude hiding (min, tail)++-- | Haskell counterpart of [Halide::Stage](https://halide-lang.org/docs/class_halide_1_1_stage.html).+data CxxStage++importHalide++-- | A function in Halide. Conceptually, it can be thought of as a lazy+-- @n@-dimensional buffer of type @a@.+--+-- This is a wrapper around the [@Halide::Func@](https://halide-lang.org/docs/class_halide_1_1_func.html)+-- C++ type.+data Func (t :: FuncTy) (n :: Nat) (a :: Type) where+ Func :: {-# UNPACK #-} !(ForeignPtr CxxFunc) -> Func 'FuncTy n a+ Param :: {-# UNPACK #-} !(IORef (Maybe (ForeignPtr CxxImageParam))) -> Func 'ParamTy n a++-- | Function type. It can either be 'FuncTy' which means that we have defined the function ourselves,+-- or 'ParamTy' which means that it's a parameter to our pipeline.+data FuncTy = FuncTy | ParamTy+ deriving stock (Show, Eq, Ord)++-- | A single definition of a t'Func'.+newtype Stage (n :: Nat) (a :: Type) = Stage (ForeignPtr CxxStage)++-- | Different ways to handle a tail case in a split when the split factor does+-- not provably divide the extent.+--+-- This is the Haskell counterpart of [@Halide::TailStrategy@](https://halide-lang.org/docs/namespace_halide.html#a6c6557df562bd7850664e70fdb8fea0f).+data TailStrategy+ = -- | Round up the extent to be a multiple of the split factor.+ --+ -- Not legal for RVars, as it would change the meaning of the algorithm.+ --+ -- * Pros: generates the simplest, fastest code.+ -- * Cons: if used on a stage that reads from the input or writes to the+ -- output, constrains the input or output size to be a multiple of the+ -- split factor.+ TailRoundUp+ | -- | Guard the inner loop with an if statement that prevents evaluation+ -- beyond the original extent.+ --+ -- Always legal. The if statement is treated like a boundary condition, and+ -- factored out into a loop epilogue if possible.+ --+ -- * Pros: no redundant re-evaluation; does not constrain input our output sizes.+ -- * Cons: increases code size due to separate tail-case handling;+ -- vectorization will scalarize in the tail case to handle the if+ -- statement.+ TailGuardWithIf+ | -- | Guard the loads and stores in the loop with an if statement that+ -- prevents evaluation beyond the original extent.+ --+ -- Always legal. The if statement is treated like a boundary condition, and+ -- factored out into a loop epilogue if possible.+ -- * Pros: no redundant re-evaluation; does not constrain input or output+ -- sizes.+ -- * Cons: increases code size due to separate tail-case handling.+ TailPredicate+ | -- | Guard the loads in the loop with an if statement that prevents+ -- evaluation beyond the original extent.+ --+ -- Only legal for innermost splits. Not legal for RVars, as it would change+ -- the meaning of the algorithm. The if statement is treated like a+ -- boundary condition, and factored out into a loop epilogue if possible.+ -- * Pros: does not constrain input sizes, output size constraints are+ -- simpler than full predication.+ -- * Cons: increases code size due to separate tail-case handling,+ -- constrains the output size to be a multiple of the split factor.+ TailPredicateLoads+ | -- | Guard the stores in the loop with an if statement that prevents+ -- evaluation beyond the original extent.+ --+ -- Only legal for innermost splits. Not legal for RVars, as it would change+ -- the meaning of the algorithm. The if statement is treated like a+ -- boundary condition, and factored out into a loop epilogue if possible.+ -- * Pros: does not constrain output sizes, input size constraints are+ -- simpler than full predication.+ -- * Cons: increases code size due to separate tail-case handling,+ -- constraints the input size to be a multiple of the split factor.+ TailPredicateStores+ | -- | Prevent evaluation beyond the original extent by shifting the tail+ -- case inwards, re-evaluating some points near the end.+ --+ -- Only legal for pure variables in pure definitions. If the inner loop is+ -- very simple, the tail case is treated like a boundary condition and+ -- factored out into an epilogue.+ --+ -- This is a good trade-off between several factors. Like 'TailRoundUp', it+ -- supports vectorization well, because the inner loop is always a fixed+ -- size with no data-dependent branching. It increases code size slightly+ -- for inner loops due to the epilogue handling, but not for outer loops+ -- (e.g. loops over tiles). If used on a stage that reads from an input or+ -- writes to an output, this stategy only requires that the input/output+ -- extent be at least the split factor, instead of a multiple of the split+ -- factor as with 'TailRoundUp'.+ TailShiftInwards+ | -- | For pure definitions use 'TailShiftInwards'.+ --+ -- For pure vars in update definitions use 'TailRoundUp'. For RVars in update+ -- definitions use 'TailGuardWithIf'.+ TailAuto+ deriving stock (Eq, Ord, Show)++-- | Specifies that @i@ is a tuple of @'Expr' Int32@.+--+-- @ts@ are deduced from @i@, so you don't have to specify them explicitly.+type IndexTuple i ts = (IsTuple (Arguments ts) i, All ((~) (Expr Int32)) ts)++-- | Common scheduling functions+class (KnownNat n, IsHalideType a) => Schedulable f n a where+ -- | Vectorize the dimension.+ vectorize :: VarOrRVar -> f n a -> IO (f n a)++ -- | Unroll the dimension.+ unroll :: VarOrRVar -> f n a -> IO (f n a)++ -- | Reorder variables to have the given nesting order, from innermost out.+ reorder :: [VarOrRVar] -> f n a -> IO (f n a)++ -- | Split a dimension into inner and outer subdimensions with the given names, where the inner dimension+ -- iterates from @0@ to @factor-1@.+ --+ -- The inner and outer subdimensions can then be dealt with using the other scheduling calls. It's okay+ -- to reuse the old variable name as either the inner or outer variable. The first argument specifies+ -- how the tail should be handled if the split factor does not provably divide the extent.+ split :: TailStrategy -> VarOrRVar -> (VarOrRVar, VarOrRVar) -> Expr Int32 -> f n a -> IO (f n a)++ -- | Join two dimensions into a single fused dimenion.+ --+ -- The fused dimension covers the product of the extents of the inner and outer dimensions given.+ fuse :: (VarOrRVar, VarOrRVar) -> VarOrRVar -> f n a -> IO (f n a)++ -- | Mark the dimension to be traversed serially+ serial :: VarOrRVar -> f n a -> IO (f n a)++ -- | Mark the dimension to be traversed in parallel+ parallel :: VarOrRVar -> f n a -> IO (f n a)++ specialize :: Expr Bool -> f n a -> IO (Stage n a)+ specializeFail :: Text -> f n a -> IO ()+ gpuBlocks :: (IndexTuple i ts, 1 <= Length ts, Length ts <= 3) => DeviceAPI -> i -> f n a -> IO (f n a)+ gpuThreads :: (IndexTuple i ts, 1 <= Length ts, Length ts <= 3) => DeviceAPI -> i -> f n a -> IO (f n a)+ gpuLanes :: DeviceAPI -> VarOrRVar -> f n a -> IO (f n a)++ -- | Schedule the iteration over this stage to be fused with another stage from outermost loop to a+ -- given LoopLevel.+ --+ -- For more info, see [Halide::Stage::compute_with](https://halide-lang.org/docs/class_halide_1_1_stage.html#a82a2ae25a009d6a2d52cb407a25f0a5b).+ computeWith :: LoopAlignStrategy -> f n a -> LoopLevel t -> IO ()++instance (KnownNat n, IsHalideType a) => Schedulable Stage n a where+ vectorize var stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar var $ \var' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->vectorize(*$(const Halide::VarOrRVar* var'));+ });+ } |]+ pure stage+ unroll var stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar var $ \var' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->unroll(*$(const Halide::VarOrRVar* var'));+ });+ } |]+ pure stage+ reorder args stage = do+ withMany asVarOrRVar args $ \args' -> do+ withCxxStage stage $ \stage' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=]() {+ $(Halide::Stage* stage')->reorder(+ *$(const std::vector<Halide::VarOrRVar>* args')); + });+ } |]+ pure stage+ split tail old (outer, inner) factor stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar old $ \old' ->+ asVarOrRVar outer $ \outer' ->+ asVarOrRVar inner $ \inner' ->+ asExpr factor $ \factor' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->split(+ *$(const Halide::VarOrRVar* old'),+ *$(const Halide::VarOrRVar* outer'),+ *$(const Halide::VarOrRVar* inner'),+ *$(const Halide::Expr* factor'),+ static_cast<Halide::TailStrategy>($(int t)));+ });+ } |]+ pure stage+ where+ t = fromIntegral . fromEnum $ tail+ fuse (outer, inner) fused stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar outer $ \outer' ->+ asVarOrRVar inner $ \inner' ->+ asVarOrRVar fused $ \fused' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->fuse(+ *$(const Halide::VarOrRVar* outer'),+ *$(const Halide::VarOrRVar* inner'),+ *$(const Halide::VarOrRVar* fused'));+ });+ } |]+ pure stage+ serial var stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar var $ \var' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->serial(*$(const Halide::VarOrRVar* var'));+ });+ } |]+ pure stage+ parallel var stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar var $ \var' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->parallel(*$(const Halide::VarOrRVar* var'));+ });+ } |]+ pure stage+ specialize cond stage = do+ withCxxStage stage $ \stage' ->+ asExpr cond $ \cond' ->+ wrapCxxStage+ =<< [C.throwBlock| Halide::Stage* {+ return handle_halide_exceptions([=](){+ return new Halide::Stage{$(Halide::Stage* stage')->specialize(+ *$(const Halide::Expr* cond'))};+ });+ } |]+ specializeFail (T.encodeUtf8 -> s) stage =+ withCxxStage stage $ \stage' ->+ [C.throwBlock| void {+ return handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->specialize_fail(+ std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)});+ });+ } |]+ gpuBlocks (fromIntegral . fromEnum -> api) vars stage = do+ withCxxStage stage $ \stage' ->+ asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');+ auto& stage = *$(Halide::Stage* stage');+ auto const device = static_cast<Halide::DeviceAPI>($(int api));+ switch (vars.size()) {+ case 1: stage.gpu_blocks(vars.at(0), device);+ break;+ case 2: stage.gpu_blocks(vars.at(0), vars.at(1), device);+ break;+ case 3: stage.gpu_blocks(vars.at(0), vars.at(1), vars.at(2), device);+ break;+ default: throw std::runtime_error{"unexpected number of arguments in gpuBlocks"};+ }+ });+ } |]+ pure stage+ gpuThreads (fromIntegral . fromEnum -> api) vars stage = do+ withCxxStage stage $ \stage' ->+ asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');+ auto& stage = *$(Halide::Stage* stage');+ auto const device = static_cast<Halide::DeviceAPI>($(int api));+ switch (vars.size()) {+ case 1: stage.gpu_threads(vars.at(0), device);+ break;+ case 2: stage.gpu_threads(vars.at(0), vars.at(1), device);+ break;+ case 3: stage.gpu_threads(vars.at(0), vars.at(1), vars.at(2), device);+ break;+ default: throw std::runtime_error{"unexpected number of arguments in gpuThreads"};+ }+ });+ } |]+ pure stage+ gpuLanes (fromIntegral . fromEnum -> api) var stage = do+ withCxxStage stage $ \stage' ->+ asVarOrRVar var $ \var' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Stage* stage')->gpu_lanes(+ *$(const Halide::VarOrRVar* var'),+ static_cast<Halide::DeviceAPI>($(int api)));+ });+ } |]+ pure stage+ computeWith (fromIntegral . fromEnum -> align) stage level = do+ withCxxStage stage $ \stage' ->+ withCxxLoopLevel level $ \level' ->+ [C.throwBlock| void {+ handle_halide_exceptions([=]() {+ $(Halide::Stage* stage')->compute_with(+ *$(const Halide::LoopLevel* level'),+ static_cast<Halide::LoopAlignStrategy>($(int align)));+ });+ } |]++viaStage1+ :: (KnownNat n, IsHalideType b)+ => (a -> Stage n b -> IO (Stage n b))+ -> a+ -> Func t n b+ -> IO (Func t n b)+viaStage1 f a1 func = do+ _ <- f a1 =<< getStage func+ pure func++viaStage2+ :: (KnownNat n, IsHalideType b)+ => (a1 -> a2 -> Stage n b -> IO (Stage n b))+ -> a1+ -> a2+ -> Func t n b+ -> IO (Func t n b)+viaStage2 f a1 a2 func = do+ _ <- f a1 a2 =<< getStage func+ pure func++{-+viaStage3+ :: (KnownNat n, IsHalideType b)+ => (a1 -> a2 -> a3 -> Stage n b -> IO (Stage n b))+ -> a1+ -> a2+ -> a3+ -> Func t n b+ -> IO (Func t n b)+viaStage3 f a1 a2 a3 func = do+ _ <- f a1 a2 a3 =<< getStage func+ pure func+-}++viaStage4+ :: (KnownNat n, IsHalideType b)+ => (a1 -> a2 -> a3 -> a4 -> Stage n b -> IO (Stage n b))+ -> a1+ -> a2+ -> a3+ -> a4+ -> Func t n b+ -> IO (Func t n b)+viaStage4 f a1 a2 a3 a4 func = do+ _ <- f a1 a2 a3 a4 =<< getStage func+ pure func++instance (KnownNat n, IsHalideType a) => Schedulable (Func t) n a where+ vectorize = viaStage1 vectorize+ unroll = viaStage1 unroll+ reorder = viaStage1 reorder+ split = viaStage4 split+ fuse = viaStage2 fuse+ serial = viaStage1 serial+ parallel = viaStage1 parallel+ specialize cond func = getStage func >>= specialize cond+ specializeFail msg func = getStage func >>= specializeFail msg+ gpuBlocks = viaStage2 gpuBlocks+ gpuThreads = viaStage2 gpuThreads+ gpuLanes = viaStage2 gpuLanes+ computeWith a f l = getStage f >>= \f' -> computeWith a f' l++instance Enum TailStrategy where+ fromEnum =+ fromIntegral . \case+ TailRoundUp -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::RoundUp) } |]+ TailGuardWithIf -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::GuardWithIf) } |]+ TailPredicate -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::Predicate) } |]+ TailPredicateLoads -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::PredicateLoads) } |]+ TailPredicateStores -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::PredicateStores) } |]+ TailShiftInwards -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::ShiftInwards) } |]+ TailAuto -> [CU.pure| int { static_cast<int>(Halide::TailStrategy::Auto) } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::RoundUp) } |] = TailRoundUp+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::GuardWithIf) } |] = TailGuardWithIf+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::Predicate) } |] = TailPredicate+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::PredicateLoads) } |] = TailPredicateLoads+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::PredicateStores) } |] = TailPredicateStores+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::ShiftInwards) } |] = TailShiftInwards+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::Auto) } |] = TailAuto+ | otherwise = error $ "invalid TailStrategy: " <> show k++-- vectorize+-- :: (KnownNat n, IsHalideType a)+-- => TailStrategy+-- -> Func t n a+-- -> Expr Int32+-- -- ^ Variable to vectorize+-- -> Expr Int32+-- -- ^ Split factor+-- -> IO ()+-- vectorize strategy func var factor =+-- withFunc func $ \f ->+-- asVarOrRVar var $ \x ->+-- asExpr factor $ \n ->+-- [C.throwBlock| void {+-- $(Halide::Func* f)->vectorize(*$(Halide::VarOrRVar* x), *$(Halide::Expr* n),+-- static_cast<Halide::TailStrategy>($(int tail)));+-- } |]+-- where+-- tail = fromIntegral (fromEnum strategy)++-- | Split a dimension by the given factor, then unroll the inner dimension.+--+-- This is how you unroll a loop of unknown size by some constant factor. After+-- this call, @var@ refers to the outer dimension of the split.+-- unroll+-- :: (KnownNat n, IsHalideType a)+-- => TailStrategy+-- -> Func t n a+-- -> Expr Int32+-- -- ^ Variable @var@ to vectorize+-- -> Expr Int32+-- -- ^ Split factor+-- -> IO ()+-- unroll strategy func var factor =+-- withFunc func $ \f ->+-- asVarOrRVar var $ \x ->+-- asExpr factor $ \n ->+-- [C.throwBlock| void {+-- $(Halide::Func* f)->unroll(*$(Halide::VarOrRVar* x), *$(Halide::Expr* n),+-- static_cast<Halide::TailStrategy>($(int tail)));+-- } |]+-- where+-- tail = fromIntegral (fromEnum strategy)++-- | Reorder variables to have the given nesting order, from innermost out.+-- reorder+-- :: forall t n a i ts+-- . ( IsTuple (Arguments ts) i+-- , All ((~) (Expr Int32)) ts+-- , Length ts ~ n+-- , KnownNat n+-- , IsHalideType a+-- )+-- => Func t n a+-- -> i+-- -> IO ()+-- reorder func args =+-- asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple args) $ \v -> do+-- withFunc func $ \f ->+-- [C.throwBlock| void { $(Halide::Func* f)->reorder(*$(std::vector<Halide::VarOrRVar>* v)); } |]++-- | Statically declare the range over which the function will be evaluated in the general case.+--+-- This provides a basis for the auto scheduler to make trade-offs and scheduling decisions.+-- The auto generated schedules might break when the sizes of the dimensions are very different from the+-- estimates specified. These estimates are used only by the auto scheduler if the function is a pipeline output.+estimate+ :: (KnownNat n, IsHalideType a)+ => Expr Int32+ -- ^ index variable+ -> Expr Int32+ -- ^ @min@ estimate+ -> Expr Int32+ -- ^ @extent@ estimate+ -> Func t n a+ -> IO ()+estimate var min extent func =+ withFunc func $ \f -> asVar var $ \i -> asExpr min $ \minExpr -> asExpr extent $ \extentExpr ->+ [CU.exp| void {+ $(Halide::Func* f)->set_estimate(+ *$(Halide::Var* i), *$(Halide::Expr* minExpr), *$(Halide::Expr* extentExpr)) } |]++-- | Statically declare the range over which a function should be evaluated.+--+-- This can let Halide perform some optimizations. E.g. if you know there are going to be 4 color channels,+-- you can completely vectorize the color channel dimension without the overhead of splitting it up.+-- If bounds inference decides that it requires more of this function than the bounds you have stated,+-- a runtime error will occur when you try to run your pipeline.+bound+ :: (KnownNat n, IsHalideType a)+ => Expr Int32+ -- ^ index variable+ -> Expr Int32+ -- ^ @min@ estimate+ -> Expr Int32+ -- ^ @extent@ estimate+ -> Func t n a+ -> IO ()+bound var min extent func =+ withFunc func $ \f -> asVar var $ \i -> asExpr min $ \minExpr -> asExpr extent $ \extentExpr ->+ [CU.exp| void {+ $(Halide::Func* f)->bound(+ *$(Halide::Var* i), *$(Halide::Expr* minExpr), *$(Halide::Expr* extentExpr)) } |]++-- | Get the index arguments of the function.+--+-- The returned list contains exactly @n@ elements.+getArgs :: (KnownNat n, IsHalideType a) => Func t n a -> IO [Var]+getArgs func =+ withFunc func $ \func' -> do+ let allocate =+ [CU.exp| std::vector<Halide::Var>* { + new std::vector<Halide::Var>{$(const Halide::Func* func')->args()} } |]+ destroy v = [CU.exp| void { delete $(std::vector<Halide::Var>* v) } |]+ bracket allocate destroy $ \v -> do+ n <- [CU.exp| size_t { $(const std::vector<Halide::Var>* v)->size() } |]+ forM [0 .. n - 1] $ \i ->+ fmap Var . cxxConstruct $ \ptr ->+ [CU.exp| void {+ new ($(Halide::Var* ptr)) Halide::Var{$(const std::vector<Halide::Var>* v)->at($(size_t i))} } |]++-- | Compute all of this function once ahead of time.+--+-- See [Halide::Func::compute_root](https://halide-lang.org/docs/class_halide_1_1_func.html#a29df45a4a16a63eb81407261a9783060) for more info.+computeRoot :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+computeRoot func = do+ withFunc func $ \f ->+ [C.throwBlock| void { handle_halide_exceptions([=](){ $(Halide::Func* f)->compute_root(); }); } |]+ pure func++-- | Creates and returns a new identity Func that wraps this Func.+--+-- During compilation, Halide replaces all calls to this Func done by 'f' with calls to the wrapper.+-- If this Func is already wrapped for use in 'f', will return the existing wrapper.+--+-- For more info, see [Halide::Func::in](https://halide-lang.org/docs/class_halide_1_1_func.html#a9d619f2d0111ea5bf640781d1324d050).+asUsedBy+ :: (KnownNat n, KnownNat m, IsHalideType a, IsHalideType b)+ => Func t1 n a+ -> Func 'FuncTy m b+ -> IO (Func 'FuncTy n a)+asUsedBy g f =+ withFunc g $ \gPtr -> withFunc f $ \fPtr ->+ wrapCxxFunc+ =<< [CU.exp| Halide::Func* {+ new Halide::Func{$(Halide::Func* gPtr)->in(*$(Halide::Func* fPtr))} } |]++-- | Create and return a global identity wrapper, which wraps all calls to this Func by any other Func.+--+-- If a global wrapper already exists, returns it. The global identity wrapper is only used by callers+-- for which no custom wrapper has been specified.+asUsed :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func 'FuncTy n a)+asUsed f =+ withFunc f $ \fPtr ->+ wrapCxxFunc+ =<< [CU.exp| Halide::Func* { new Halide::Func{$(Halide::Func* fPtr)->in()} } |]++-- | Declare that this function should be implemented by a call to @halide_buffer_copy@ with the given+-- target device API.+--+-- Asserts that the @Func@ has a pure definition which is a simple call to a single input, and no update+-- definitions. The wrapper @Func@s returned by 'asUsed' are suitable candidates. Consumes all pure variables,+-- and rewrites the @Func@ to have an extern definition that calls @halide_buffer_copy@.+copyToDevice :: (KnownNat n, IsHalideType a) => DeviceAPI -> Func t n a -> IO (Func t n a)+copyToDevice deviceApi func = do+ withFunc func $ \f ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Func* f)->copy_to_device(static_cast<Halide::DeviceAPI>($(int api)));+ });+ } |]+ pure func+ where+ api = fromIntegral . fromEnum $ deviceApi++-- | Same as @'copyToDevice' 'DeviceHost'@+copyToHost :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+copyToHost = copyToDevice DeviceHost++-- | Split a dimension into inner and outer subdimensions with the given names, where the inner dimension+-- iterates from @0@ to @factor-1@.+--+-- The inner and outer subdimensions can then be dealt with using the other scheduling calls. It's okay+-- to reuse the old variable name as either the inner or outer variable. The first argument specifies+-- how the tail should be handled if the split factor does not provably divide the extent.+-- split+-- :: (KnownNat n, IsHalideType a)+-- => TailStrategy+-- -- ^ how to treat the remainder+-- -> Func t n a+-- -> Expr Int32+-- -- ^ loop variable to split+-- -> Expr Int32+-- -- ^ new outer loop variable+-- -> Expr Int32+-- -- ^ new inner loop variable+-- -> Expr Int32+-- -- ^ split factor+-- -> IO (Func t n a)+-- split tail func old outer inner factor = do+-- withFunc func $ \f ->+-- asVarOrRVar old $ \old' ->+-- asVarOrRVar outer $ \outer' ->+-- asVarOrRVar inner $ \inner' ->+-- asExpr factor $ \factor' ->+-- [C.throwBlock| void {+-- handle_halide_exceptions([=](){+-- $(Halide::Func* f)->split(+-- *$(const Halide::VarOrRVar* old'),+-- *$(const Halide::VarOrRVar* outer'),+-- *$(const Halide::VarOrRVar* inner'),+-- *$(const Halide::Expr* factor'),+-- static_cast<Halide::TailStrategy>($(int t)));+-- }); } |]+-- pure func+-- where+-- t = fromIntegral . fromEnum $ tail++-- | Join two dimensions into a single fused dimenion.+--+-- The fused dimension covers the product of the extents of the inner and outer dimensions given.+-- fuse+-- :: (KnownNat n, IsHalideType a)+-- => Func t n a+-- -> Expr Int32+-- -- ^ inner loop variable+-- -> Expr Int32+-- -- ^ outer loop variable+-- -> Expr Int32+-- -- ^ new fused loop variable+-- -> IO (Func t n a)+-- fuse func outer inner fused = do+-- withFunc func $ \f ->+-- asVarOrRVar outer $ \outer' ->+-- asVarOrRVar inner $ \inner' ->+-- asVarOrRVar fused $ \fused' ->+-- [CU.exp| void {+-- $(Halide::Func* f)->fuse(+-- *$(const Halide::VarOrRVar* outer'),+-- *$(const Halide::VarOrRVar* inner'),+-- *$(const Halide::VarOrRVar* fused')) } |]+-- pure func++-- withVarOrRVarMany :: [Expr Int32] -> (Int -> Ptr (CxxVector CxxVarOrRVar) -> IO a) -> IO a+-- withVarOrRVarMany xs f =+-- bracket allocate destroy $ \v -> do+-- let go !k [] = f k v+-- go !k (y : ys) = withVarOrRVarMany y $ \p -> do+-- [CU.exp| void { $(std::vector<Halide::Expr>* v)->push_back(*$(Halide::VarOrRVar* p)) } |]+-- go (k + 1) ys+-- go 0 xs+-- where+-- count = fromIntegral (length xs)++-- withFunc func $ \f ->+-- withVarOrRVarMany vars $ \count v -> do+-- unless natVal (Proxy @n)+-- handleHalideExceptionsM+-- [C.tryBlock| void {+-- $(Halide::Func* f)->reorder(*$(std::vector<Halide::VarOrRVar>* v));+-- } |]+--+-- class Curry (args :: [Type]) (r :: Type) (f :: Type) | args r -> f where+-- curryG :: (Arguments args -> r) -> f++mkBufferParameter+ :: forall n a. (KnownNat n, IsHalideType a) => Maybe Text -> IO (ForeignPtr CxxImageParam)+mkBufferParameter maybeName = do+ with (halideTypeFor (Proxy @a)) $ \t -> do+ let d = fromIntegral $ natVal (Proxy @n)+ createWithoutName =+ [CU.exp| Halide::ImageParam* {+ new Halide::ImageParam{Halide::Type{*$(halide_type_t* t)}, $(int d)} } |]+ deleter = [C.funPtr| void deleteImageParam(Halide::ImageParam* p) { delete p; } |]+ createWithName name =+ let s = T.encodeUtf8 name+ in [CU.exp| Halide::ImageParam* {+ new Halide::ImageParam{+ Halide::Type{*$(halide_type_t* t)},+ $(int d),+ std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}} } |]+ newForeignPtr deleter =<< maybe createWithoutName createWithName maybeName++getBufferParameter+ :: forall n a+ . (KnownNat n, IsHalideType a)+ => Maybe Text+ -> IORef (Maybe (ForeignPtr CxxImageParam))+ -> IO (ForeignPtr CxxImageParam)+getBufferParameter name r =+ readIORef r >>= \case+ Just fp -> pure fp+ Nothing -> do+ fp <- mkBufferParameter @n @a name+ writeIORef r (Just fp)+ pure fp++-- | Same as 'withFunc', but ensures that we're dealing with 'Param' instead of a 'Func'.+withBufferParam+ :: forall n a b+ . (HasCallStack, KnownNat n, IsHalideType a)+ => Func 'ParamTy n a+ -> (Ptr CxxImageParam -> IO b)+ -> IO b+withBufferParam (Param r) action =+ getBufferParameter @n @a Nothing r >>= flip withForeignPtr action++-- instance (KnownNat n, IsHalideType a) => Named (Func 'ParamTy n a) where+-- setName :: Func 'ParamTy n a -> Text -> IO ()+-- setName (Param r) name = do+-- readIORef r >>= \case+-- Just _ -> error "the name of this Func has already been set"+-- Nothing -> do+-- fp <- mkBufferParameter @n @a (Just name)+-- writeIORef r (Just fp)++-- | Get the underlying pointer to @Halide::Func@ and invoke an 'IO' action with it.+withFunc :: (KnownNat n, IsHalideType a) => Func t n a -> (Ptr CxxFunc -> IO b) -> IO b+withFunc f = withForeignPtr (funcToForeignPtr f)++wrapCxxFunc :: Ptr CxxFunc -> IO (Func 'FuncTy n a)+wrapCxxFunc = fmap Func . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteFunc(Halide::Func *x) { delete x; } |]++forceFunc :: forall t n a. (KnownNat n, IsHalideType a) => Func t n a -> IO (Func 'FuncTy n a)+forceFunc x@(Func _) = pure x+forceFunc (Param r) = do+ fp <- getBufferParameter @n @a Nothing r+ withForeignPtr fp $ \p ->+ wrapCxxFunc+ =<< [CU.exp| Halide::Func* {+ new Halide::Func{static_cast<Halide::Func>(*$(Halide::ImageParam* p))} } |]++funcToForeignPtr :: (KnownNat n, IsHalideType a) => Func t n a -> ForeignPtr CxxFunc+funcToForeignPtr x = unsafePerformIO $! forceFunc x >>= \(Func fp) -> pure fp++-- | Define a Halide function.+--+-- @define "f" i e@ defines a Halide function called "f" such that @f[i] = e@.+--+-- Here, @i@ is an @n@-element tuple of t'Var', i.e. the following are all valid:+--+-- >>> [x, y, z] <- mapM mkVar ["x", "y", "z"]+-- >>> f1 <- define "f1" x (0 :: Expr Float)+-- >>> f2 <- define "f2" (x, y) (0 :: Expr Float)+-- >>> f3 <- define "f3" (x, y, z) (0 :: Expr Float)+define+ :: ( IsTuple (Arguments ts) i+ , All ((~) Var) ts+ , Length ts ~ n+ , KnownNat n+ , IsHalideType a+ )+ => Text+ -> i+ -> Expr a+ -> IO (Func 'FuncTy n a)+define name args expr =+ asVectorOf @((~) (Expr Int32)) asVar (fromTuple args) $ \x -> do+ let s = T.encodeUtf8 name+ asExpr expr $ \y ->+ wrapCxxFunc+ =<< [CU.block| Halide::Func* {+ Halide::Func f{std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}};+ f(*$(std::vector<Halide::Var>* x)) = *$(Halide::Expr* y);+ return new Halide::Func{f};+ } |]++-- | Create an update definition for a Halide function.+--+-- @update f i e@ creates an update definition for @f@ that performs @f[i] = e@.+update+ :: ( IsTuple (Arguments ts) i+ , All ((~) (Expr Int32)) ts+ , Length ts ~ n+ , KnownNat n+ , IsHalideType a+ )+ => Func 'FuncTy n a+ -> i+ -> Expr a+ -> IO ()+update func args expr =+ withFunc func $ \f ->+ asVectorOf @((~) (Expr Int32)) asExpr (fromTuple args) $ \x ->+ asExpr expr $ \y ->+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Func* f)->operator()(*$(std::vector<Halide::Expr>* x)) = *$(Halide::Expr* y);+ });+ } |]++infix 9 !++-- | Apply a Halide function. Conceptually, @f ! i@ is equivalent to @f[i]@, i.e.+-- indexing into a lazy array.+(!)+ :: ( IsTuple (Arguments ts) i+ , All ((~) (Expr Int32)) ts+ , Length ts ~ n+ , KnownNat n+ , IsHalideType a+ )+ => Func t n a+ -> i+ -> Expr a+(!) func args =+ unsafePerformIO $+ withFunc func $ \f ->+ asVectorOf @((~) (Expr Int32)) asExpr (fromTuple args) $ \x ->+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(Halide::Func* f)->operator()(*$(std::vector<Halide::Expr>* x))} } |]++-- | Get a particular dimension of a pipeline parameter.+dim+ :: forall n a+ . (HasCallStack, KnownNat n, IsHalideType a)+ => Int+ -> Func 'ParamTy n a+ -> IO Dimension+dim k func+ | 0 <= k && k < fromIntegral (natVal (Proxy @n)) =+ let n = fromIntegral k+ in withBufferParam func $ \f ->+ wrapCxxDimension+ =<< [CU.exp| Halide::Internal::Dimension* {+ new Halide::Internal::Dimension{$(Halide::ImageParam* f)->dim($(int n))} } |]+ | otherwise =+ error $+ "invalid dimension index: "+ <> show k+ <> "; Func is "+ <> show (natVal (Proxy @n))+ <> "-dimensional"++-- | Write out the loop nests specified by the schedule for this function.+--+-- Helpful for understanding what a schedule is doing.+--+-- For more info, see+-- [@Halide::Func::print_loop_nest@](https://halide-lang.org/docs/class_halide_1_1_func.html#a03f839d9e13cae4b87a540aa618589ae)+-- printLoopNest :: (KnownNat n, IsHalideType r) => Func n r -> IO ()+-- printLoopNest func = withFunc func $ \f ->+-- [C.exp| void { $(Halide::Func* f)->print_loop_nest() } |]++-- | Get the loop nests specified by the schedule for this function.+--+-- Helpful for understanding what a schedule is doing.+--+-- For more info, see+-- [@Halide::Func::print_loop_nest@](https://halide-lang.org/docs/class_halide_1_1_func.html#a03f839d9e13cae4b87a540aa618589ae)+prettyLoopNest :: (KnownNat n, IsHalideType r) => Func t n r -> IO Text+prettyLoopNest func = withFunc func $ \f ->+ peekAndDeleteCxxString+ =<< [C.throwBlock| std::string* {+ return handle_halide_exceptions([=]() {+ return new std::string{Halide::Internal::print_loop_nest(+ std::vector<Halide::Internal::Function>{$(Halide::Func* f)->function()})};+ });+ } |]++-- | Evaluate this function over a rectangular domain.+realize+ :: forall n a t b+ . (KnownNat n, IsHalideType a)+ => Func t n a+ -- ^ Function to evaluate+ -> [Int]+ -- ^ Domain over which to evaluate+ -> (Ptr (HalideBuffer n a) -> IO b)+ -- ^ What to do with the buffer afterwards. Note that the buffer is allocated only temporary,+ -- so do not return it directly.+ -> IO b+realize func shape action =+ withFunc func $ \f ->+ allocaCpuBuffer shape $ \buf -> do+ let raw = castPtr buf+ [C.throwBlock| void {+ handle_halide_exceptions([=](){+ $(Halide::Func* f)->realize(+ Halide::Pipeline::RealizationArg{$(halide_buffer_t* raw)});+ });+ } |]+ action buf++-- \| Evaluate this function over a one-dimensional domain and return the+-- resulting buffer or buffers.+-- realize1D+-- :: forall a t+-- . IsHalideType a+-- => Int+-- -- ^ @size@ of the domain. The function will be evaluated on @[0, ..., size -1]@+-- -> Func t 1 a+-- -- ^ Function to evaluate+-- -> IO (Vector a)+-- realize1D size func = do+-- buf <- SM.new size+-- withHalideBuffer @1 @a buf $ \x -> do+-- let b = castPtr x+-- withFunc func $ \f ->+-- [CU.exp| void {+-- $(Halide::Func* f)->realize(+-- Halide::Pipeline::RealizationArg{$(halide_buffer_t* b)}) } |]+-- S.unsafeFreeze buf++-- | A view pattern to specify the name of a buffer argument.+--+-- Example usage:+--+-- >>> :{+-- _ <- compile $ \(buffer "src" -> src) -> do+-- i <- mkVar "i"+-- define "dest" i $ (src ! i :: Expr Float)+-- :}+--+-- or if we want to specify the dimension and type, we can use type applications:+--+-- >>> :{+-- _ <- compile $ \(buffer @1 @Float "src" -> src) -> do+-- i <- mkVar "i"+-- define "dest" i $ src ! i+-- :}+buffer :: forall n a. (KnownNat n, IsHalideType a) => Text -> Func 'ParamTy n a -> Func 'ParamTy n a+buffer name p@(Param r) = unsafePerformIO $ do+ _ <- getBufferParameter @n @a (Just name) r+ pure p++-- | Similar to 'buffer', but for scalar parameters.+--+-- Example usage:+--+-- >>> :{+-- _ <- compile $ \(scalar @Float "a" -> a) -> do+-- i <- mkVar "i"+-- define "dest" i $ a+-- :}+scalar :: forall a. IsHalideType a => Text -> Expr a -> Expr a+scalar name (ScalarParam r) = unsafePerformIO $ do+ readIORef r >>= \case+ Just _ -> error "the name of this Expr has already been set"+ Nothing -> do+ fp <- mkScalarParameter @a (Just name)+ writeIORef r (Just fp)+ pure (ScalarParam r)+scalar _ _ = error "cannot set the name of an expression that is not a parameter"++wrapCxxStage :: (KnownNat n, IsHalideType a) => Ptr CxxStage -> IO (Stage n a)+wrapCxxStage = fmap Stage . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteStage(Halide::Stage* p) { delete p; } |]++withCxxStage :: (KnownNat n, IsHalideType a) => Stage n a -> (Ptr CxxStage -> IO b) -> IO b+withCxxStage (Stage fp) = withForeignPtr fp++-- | Get the pure stage of a 'Func' for the purposes of scheduling it.+getStage :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Stage n a)+getStage func =+ withFunc func $ \func' ->+ [CU.exp| Halide::Stage* { new Halide::Stage{static_cast<Halide::Stage>(*$(Halide::Func* func'))} } |]+ >>= wrapCxxStage++-- | Return 'True' when the function has update definitions, 'False' otherwise.+hasUpdateDefinitions :: (KnownNat n, IsHalideType a) => Func t n a -> IO Bool+hasUpdateDefinitions func =+ withFunc func $ \func' ->+ toBool <$> [CU.exp| bool { $(const Halide::Func* func')->has_update_definition() } |]++-- | Get a handle to an update step for the purposes of scheduling it.+getUpdateStage :: (KnownNat n, IsHalideType a) => Int -> Func 'FuncTy n a -> IO (Stage n a)+getUpdateStage k func =+ withFunc func $ \func' ->+ let k' = fromIntegral k+ in [CU.exp| Halide::Stage* { new Halide::Stage{$(Halide::Func* func')->update($(int k'))} } |]+ >>= wrapCxxStage++-- | Identify the loop nest corresponding to some dimension of some function.+getLoopLevelAtStage+ :: (KnownNat n, IsHalideType a)+ => Func t n a+ -> Expr Int32+ -> Int+ -- ^ update index+ -> IO (LoopLevel 'LockedTy)+getLoopLevelAtStage func var stageIndex =+ withFunc func $ \f -> asVarOrRVar var $ \i -> do+ (SomeLoopLevel level) <-+ wrapCxxLoopLevel+ =<< [C.throwBlock| Halide::LoopLevel* {+ return handle_halide_exceptions([=](){+ return new Halide::LoopLevel{*$(const Halide::Func* f),+ *$(const Halide::VarOrRVar* i),+ $(int k)};+ });+ } |]+ case level of+ LoopLevel _ -> pure level+ _ -> error $ "getLoopLevelAtStage: got " <> show level <> ", but expected a LoopLevel 'LockedTy"+ where+ k = fromIntegral stageIndex++-- | Same as 'getLoopLevelAtStage' except that the stage is @-1@.+getLoopLevel :: (KnownNat n, IsHalideType a) => Func t n a -> Expr Int32 -> IO (LoopLevel 'LockedTy)+getLoopLevel f i = getLoopLevelAtStage f i (-1)++-- | Allocate storage for this function within a particular loop level.+--+-- Scheduling storage is optional, and can be used to separate the loop level at which storage is allocated+-- from the loop level at which computation occurs to trade off between locality and redundant work.+--+-- For more info, see [Halide::Func::store_at](https://halide-lang.org/docs/class_halide_1_1_func.html#a417c08f8aa3a5cdf9146fba948b65193).+storeAt :: (KnownNat n, IsHalideType a) => Func 'FuncTy n a -> LoopLevel t -> IO (Func 'FuncTy n a)+storeAt func level = do+ withFunc func $ \f ->+ withCxxLoopLevel level $ \l ->+ [CU.exp| void { $(Halide::Func* f)->store_at(*$(const Halide::LoopLevel* l)) } |]+ pure func++-- | Schedule a function to be computed within the iteration over a given loop level.+--+-- For more info, see [Halide::Func::compute_at](https://halide-lang.org/docs/class_halide_1_1_func.html#a800cbcc3ca5e3d3fa1707f6e1990ec83).+computeAt :: (KnownNat n, IsHalideType a) => Func 'FuncTy n a -> LoopLevel t -> IO (Func 'FuncTy n a)+computeAt func level = do+ withFunc func $ \f ->+ withCxxLoopLevel level $ \l ->+ [CU.exp| void { $(Halide::Func* f)->compute_at(*$(const Halide::LoopLevel* l)) } |]+ pure func++-- | Wrap a buffer into a t'Func'.+--+-- Suppose, we are defining a pipeline that adds together two vectors, and we'd like to call 'realize' to+-- evaluate it directly, how do we pass the vectors to the t'Func'? 'asBufferParam' allows to do exactly this.+--+-- > asBuffer [1, 2, 3] $ \a ->+-- > asBuffer [4, 5, 6] $ \b -> do+-- > i <- mkVar "i"+-- > f <- define "vectorAdd" i $ a ! i + b ! i+-- > realize f [3] $ \result ->+-- > print =<< peekToList f+asBufferParam+ :: forall n a t b+ . IsHalideBuffer t n a+ => t+ -- ^ Object to treat as a buffer+ -> (Func 'ParamTy n a -> IO b)+ -- ^ What to do with the __temporary__ buffer+ -> IO b+asBufferParam arr action =+ withHalideBuffer @n @a arr $ \arr' -> do+ param <- mkBufferParameter @n @a Nothing+ withForeignPtr param $ \param' ->+ let buf = (castPtr arr' :: Ptr RawHalideBuffer)+ in [CU.block| void {+ $(Halide::ImageParam* param')->set(Halide::Buffer<>{*$(const halide_buffer_t* buf)});+ } |]+ action . Param =<< newIORef (Just param)
+ src/Language/Halide/Kernel.hs view
@@ -0,0 +1,365 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE UndecidableInstances #-}++-- |+-- Module : Language.Halide.Kernel+-- Description : Compiling functions to kernels+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Kernel+ ( compile+ , compileForTarget+ , compileToCallable+ , compileToLoweredStmt+ , StmtOutputFormat (..)+ , IsFuncBuilder+ , ReturnsFunc+ , Lowered+ )+where++import Control.Exception (bracket)+import Control.Monad.Primitive (touch)+import Control.Monad.ST (RealWorld)+import Data.IORef+import Data.Kind (Type)+import Data.Primitive.PrimArray (MutablePrimArray)+import Data.Primitive.PrimArray qualified as P+import Data.Primitive.Ptr qualified as P+import Data.Proxy+import Data.Text (Text, pack)+import Data.Text.Encoding (encodeUtf8)+import Data.Text.IO qualified as T+import Foreign.C.Types (CUIntPtr (..))+import Foreign.ForeignPtr+import Foreign.ForeignPtr.Unsafe+import Foreign.Ptr (FunPtr, Ptr, castPtr)+import Foreign.Storable+import GHC.TypeNats+import Language.C.Inline qualified as C+import Language.C.Inline.Cpp.Exception qualified as C+import Language.C.Inline.Unsafe qualified as CU+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Expr+import Language.Halide.Func+import Language.Halide.RedundantConstraints+import Language.Halide.Target+import Language.Halide.Type+import System.IO.Temp (withSystemTempDirectory)+import Unsafe.Coerce (unsafeCoerce)++importHalide++data ArgvStorage s+ = ArgvStorage+ {-# UNPACK #-} !(MutablePrimArray s (Ptr ()))+ {-# UNPACK #-} !(MutablePrimArray s CUIntPtr)++newArgvStorage :: Int -> IO (ArgvStorage RealWorld)+newArgvStorage n = ArgvStorage <$> P.newPinnedPrimArray n <*> P.newPinnedPrimArray n++setArgvStorage+ :: (All ValidArgument inputs, All ValidArgument outputs)+ => ArgvStorage RealWorld+ -> Arguments inputs+ -> Arguments outputs+ -> IO ()+setArgvStorage (ArgvStorage argv scalarStorage) inputs outputs = do+ let argvPtr = P.mutablePrimArrayContents argv+ scalarStoragePtr = P.mutablePrimArrayContents scalarStorage+ go :: All ValidArgument ts' => Int -> Arguments ts' -> IO Int+ go !i Nil = pure i+ go !i ((x :: t) ::: xs) = do+ fillSlot+ (castPtr $ argvPtr `P.advancePtr` i)+ (castPtr $ scalarStoragePtr `P.advancePtr` i)+ x+ go (i + 1) xs+ i <- go 0 inputs+ _ <- go i outputs+ touch argv+ touch scalarStorage++-- | Specifies that the type can be used as an argument to a kernel.+class ValidArgument (t :: Type) where+ fillSlot :: Ptr () -> Ptr () -> t -> IO ()++instance IsHalideType t => ValidArgument t where+ fillSlot :: Ptr () -> Ptr () -> t -> IO ()+ fillSlot argv scalarStorage x = do+ poke (castPtr scalarStorage :: Ptr t) x+ poke (castPtr argv :: Ptr (Ptr ())) scalarStorage+ {-# INLINE fillSlot #-}++instance {-# OVERLAPPING #-} ValidArgument (Ptr CxxUserContext) where+ fillSlot :: Ptr () -> Ptr () -> Ptr CxxUserContext -> IO ()+ fillSlot argv scalarStorage x = do+ poke (castPtr scalarStorage :: Ptr (Ptr CxxUserContext)) x+ poke (castPtr argv :: Ptr (Ptr ())) scalarStorage+ {-# INLINE fillSlot #-}++instance {-# OVERLAPPING #-} ValidArgument (Ptr (HalideBuffer n a)) where+ fillSlot :: Ptr () -> Ptr () -> Ptr (HalideBuffer n a) -> IO ()+ fillSlot argv _ x = do+ poke (castPtr argv :: Ptr (Ptr (HalideBuffer n a))) x+ {-# INLINE fillSlot #-}++class ValidArgument (Lowered t) => ValidParameter (t :: Type) where+ appendToArgList :: Ptr (CxxVector CxxArgument) -> t -> IO ()+ prepareParameter :: IO t++instance IsHalideType a => ValidParameter (Expr a) where+ appendToArgList :: Ptr (CxxVector CxxArgument) -> Expr a -> IO ()+ appendToArgList v expr =+ asScalarParam expr $ \p ->+ [CU.exp| void { $(std::vector<Halide::Argument>* v)->emplace_back(+ $(Halide::Internal::Parameter const* p)->name(),+ Halide::Argument::InputScalar,+ $(Halide::Internal::Parameter const* p)->type(),+ $(Halide::Internal::Parameter const* p)->dimensions(),+ $(Halide::Internal::Parameter const* p)->get_argument_estimates()) } |]+ prepareParameter :: IO (Expr a)+ prepareParameter = ScalarParam <$> newIORef Nothing++instance (KnownNat n, IsHalideType a) => ValidParameter (Func t n a) where+ appendToArgList :: Ptr (CxxVector CxxArgument) -> Func t n a -> IO ()+ appendToArgList v func@(Param _) =+ withBufferParam func $ \p ->+ [CU.exp| void { $(std::vector<Halide::Argument>* v)->push_back(+ *$(Halide::ImageParam const* p)) } |]+ appendToArgList _ _ = error "appendToArgList called on Func; this should never happen"+ prepareParameter :: IO (Func t n a)+ prepareParameter = unsafeCoerce $ Param <$> newIORef Nothing++class PrepareParameters ts where+ prepareParameters :: IO (Arguments ts)++instance PrepareParameters '[] where+ prepareParameters :: IO (Arguments '[])+ prepareParameters = pure Nil++instance (ValidParameter t, PrepareParameters ts) => PrepareParameters (t ': ts) where+ prepareParameters :: IO (Arguments (t : ts))+ prepareParameters = do+ t <- prepareParameter @t+ ts <- prepareParameters @ts+ pure $ t ::: ts++prepareCxxArguments+ :: forall ts b+ . (All ValidParameter ts, KnownNat (Length ts))+ => Arguments ts+ -> (Ptr (CxxVector CxxArgument) -> IO b)+ -> IO b+prepareCxxArguments args action = do+ let count = fromIntegral (natVal (Proxy @(Length ts)))+ allocate =+ [CU.block| std::vector<Halide::Argument>* {+ auto p = new std::vector<Halide::Argument>{};+ p->reserve($(size_t count));+ return p;+ } |]+ destroy p = [CU.exp| void { delete $(std::vector<Halide::Argument>* p) } |]+ bracket allocate destroy $ \v -> do+ let go :: All ValidParameter ts' => Arguments ts' -> IO ()+ go Nil = pure ()+ go (x ::: xs) = appendToArgList v x >> go xs+ go args+ action v++deleteCxxUserContext :: FunPtr (Ptr CxxUserContext -> IO ())+deleteCxxUserContext = [C.funPtr| void deleteUserContext(Halide::JITUserContext* p) { delete p; } |]++wrapCxxUserContext :: Ptr CxxUserContext -> IO (ForeignPtr CxxUserContext)+wrapCxxUserContext = newForeignPtr deleteCxxUserContext++newEmptyCxxUserContext :: IO (ForeignPtr CxxUserContext)+newEmptyCxxUserContext =+ wrapCxxUserContext =<< [CU.exp| Halide::JITUserContext* { new Halide::JITUserContext{} } |]++wrapCxxCallable :: Ptr CxxCallable -> IO (Callable inputs outputs)+wrapCxxCallable = fmap Callable . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteCallable(Halide::Callable* p) { delete p; } |]++type Lowered :: forall k. k -> k++-- | Specifies how t'Expr' and t'Func' parameters become scalar and buffer arguments in compiled kernels.+type family Lowered (t :: k) :: k where+ Lowered (Expr a) = a+ Lowered (Func t n a) = Ptr (HalideBuffer n a)+ Lowered '[] = '[]+ Lowered (Expr a ': ts) = (a ': Lowered ts)+ Lowered (Func t n a ': ts) = (Ptr (HalideBuffer n a) ': Lowered ts)++-- | A constraint that specifies that the function @f@ returns @'IO' ('Func' t n a)@.+class (FunctionReturn f ~ IO (Func t n a), IsHalideType a, KnownNat n) => ReturnsFunc f t n a | f -> t n a++instance (FunctionReturn f ~ IO (Func t n a), IsHalideType a, KnownNat n) => ReturnsFunc f t n a++type IsFuncBuilder f t n a =+ ( All ValidParameter (FunctionArguments f)+ , All ValidArgument (Lowered (FunctionArguments f))+ , UnCurry f (FunctionArguments f) (FunctionReturn f)+ , PrepareParameters (FunctionArguments f)+ , ReturnsFunc f t n a+ , KnownNat (Length (FunctionArguments f))+ , KnownNat (Length (Lowered (FunctionArguments f)))+ )++buildFunc :: (IsFuncBuilder f t n a) => f -> IO (Arguments (FunctionArguments f), Func t n a)+buildFunc builder = do+ parameters <- prepareParameters+ func <- uncurryG builder parameters+ pure (parameters, func)++newtype Callable (inputs :: [Type]) (output :: Type) = Callable (ForeignPtr CxxCallable)++compileToCallable+ :: forall n a t f inputs output+ . ( IsFuncBuilder f t n a+ , Lowered (FunctionArguments f) ~ inputs+ , Ptr (HalideBuffer n a) ~ output+ )+ => Target+ -> f+ -> IO (Callable inputs output)+compileToCallable target builder = do+ (args, func) <- buildFunc builder+ prepareCxxArguments args $ \args' ->+ withFunc func $ \func' ->+ withCxxTarget target $ \target' ->+ wrapCxxCallable+ =<< [C.throwBlock| Halide::Callable* {+ return handle_halide_exceptions([=]() {+ return new Halide::Callable{+ $(Halide::Func* func')->compile_to_callable(+ *$(const std::vector<Halide::Argument>* args'),+ *$(const Halide::Target* target'))};+ });+ } |]+ where+ _ = keepRedundantConstraint (Proxy @(Ptr (HalideBuffer n a) ~ output))++callableToFunction+ :: forall inputs output kernel+ . ( Curry inputs (output -> IO ()) kernel+ , KnownNat (Length inputs)+ , All ValidArgument inputs+ , ValidArgument output+ )+ => Callable inputs output+ -> IO kernel+callableToFunction (Callable callable) = do+ context <- newEmptyCxxUserContext+ -- +1 comes from CxxUserContext and another +1 comes from output+ let argc = 2 + fromIntegral (natVal (Proxy @(Length inputs)))+ storage@(ArgvStorage argv scalarStorage) <- newArgvStorage (fromIntegral argc)+ let argvPtr = P.mutablePrimArrayContents argv+ contextPtr = unsafeForeignPtrToPtr context+ callablePtr = unsafeForeignPtrToPtr callable+ kernel args out = do+ setArgvStorage storage (contextPtr ::: args) (out ::: Nil)+ [CU.exp| void {+ handle_halide_exceptions([=]() {+ return $(Halide::Callable* callablePtr)->call_argv_fast(+ $(int argc), $(const void* const* argvPtr));+ })+ } |]+ touch argv+ touch scalarStorage+ touch context+ touch callable+ pure $ curryG @inputs @(output -> IO ()) kernel++-- | Convert a function that builds a Halide 'Func' into a normal Haskell function acccepting scalars and+-- 'HalideBuffer's.+--+-- For example:+--+-- @+-- builder :: Expr Float -> Func 'ParamTy 1 Float -> IO (Func 'FuncTy 1 Float)+-- builder scale inputVector = do+-- i <- 'mkVar' "i"+-- scaledVector <- 'define' "scaledVector" i $ scale * inputVector '!' i+-- pure scaledVector+-- @+--+-- The @builder@ function accepts a scalar parameter and a vector and scales the vector by the given factor.+-- We can now pass @builder@ to 'compile':+--+-- @+-- scaler <- 'compile' builder+-- 'withHalideBuffer' @1 @Float [1, 1, 1] $ \inputVector ->+-- 'allocaCpuBuffer' [3] $ \outputVector -> do+-- -- invoke the kernel+-- scaler 2.0 inputVector outputVector+-- -- print the result+-- print =<< 'peekToList' outputVector+-- @+compile+ :: forall n a t f kernel+ . ( IsFuncBuilder f t n a+ , Curry (Lowered (FunctionArguments f)) (Ptr (HalideBuffer n a) -> IO ()) kernel+ )+ => f+ -- ^ Function to compile+ -> IO kernel+ -- ^ Compiled kernel+compile = compileForTarget hostTarget++-- | Similar to 'compile', but the first argument lets you explicitly specify the compilation target.+compileForTarget+ :: forall n a t f kernel+ . ( IsFuncBuilder f t n a+ , Curry (Lowered (FunctionArguments f)) (Ptr (HalideBuffer n a) -> IO ()) kernel+ )+ => Target+ -> f+ -> IO kernel+compileForTarget target builder = compileToCallable target builder >>= callableToFunction++-- | Format in which to return the lowered code.+data StmtOutputFormat+ = -- | plain text+ StmtText+ | -- | HTML+ StmtHTML+ deriving stock (Show, Eq)++instance Enum StmtOutputFormat where+ fromEnum =+ fromIntegral . \case+ StmtText -> [CU.pure| int { static_cast<int>(Halide::StmtOutputFormat::Text) } |]+ StmtHTML -> [CU.pure| int { static_cast<int>(Halide::StmtOutputFormat::HTML) } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::StmtOutputFormat::Text) } |] = StmtText+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::StmtOutputFormat::HTML) } |] = StmtHTML+ | otherwise = error $ "invalid StmtOutputFormat " <> show k++-- | Get the internal representation of lowered code.+--+-- Useful for analyzing and debugging scheduling. Can emit HTML or plain text.+compileToLoweredStmt+ :: forall n a t f. (IsFuncBuilder f t n a) => StmtOutputFormat -> Target -> f -> IO Text+compileToLoweredStmt format target builder = do+ withSystemTempDirectory "halide-haskell" $ \dir -> do+ let s = encodeUtf8 (pack (dir <> "/code.stmt"))+ o = fromIntegral (fromEnum format)+ (parameters, func) <- buildFunc builder+ prepareCxxArguments parameters $ \v ->+ withFunc func $ \f ->+ withCxxTarget target $ \t ->+ [C.throwBlock| void {+ handle_halide_exceptions([=]() {+ $(Halide::Func* f)->compile_to_lowered_stmt(+ std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)},+ *$(const std::vector<Halide::Argument>* v),+ static_cast<Halide::StmtOutputFormat>($(int o)),+ *$(Halide::Target* t));+ });+ } |]+ T.readFile (dir <> "/code.stmt")
+ src/Language/Halide/LoopLevel.hs view
@@ -0,0 +1,177 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}++-- |+-- Module : Language.Halide.LoopLevel+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.LoopLevel+ ( LoopLevel (..)+ , LoopLevelTy (..)+ , SomeLoopLevel (..)+ , LoopAlignStrategy (..)++ -- * Internal+ , CxxLoopLevel+ , withCxxLoopLevel+ , wrapCxxLoopLevel+ )+where++import Control.Exception (bracket)+import Data.Text (Text)+import Foreign.ForeignPtr+import Foreign.Marshal (toBool)+import Foreign.Ptr (Ptr)+import GHC.Records (HasField (..))+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Cpp.Exception as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Context+import Language.Halide.Expr+import Language.Halide.Type+import Language.Halide.Utils+import System.IO.Unsafe (unsafePerformIO)+import Prelude hiding (min, tail)++-- | Haskell counterpart of @Halide::LoopLevel@+data CxxLoopLevel++importHalide++data LoopLevelTy = InlinedTy | RootTy | LockedTy++-- | A reference to a site in a Halide statement at the top of the body of a particular for loop.+data LoopLevel (t :: LoopLevelTy) where+ InlinedLoopLevel :: LoopLevel 'InlinedTy+ RootLoopLevel :: LoopLevel 'RootTy+ LoopLevel :: !(ForeignPtr CxxLoopLevel) -> LoopLevel 'LockedTy++data SomeLoopLevel where+ SomeLoopLevel :: LoopLevel t -> SomeLoopLevel++deriving stock instance Show SomeLoopLevel++instance Eq SomeLoopLevel where+ (SomeLoopLevel InlinedLoopLevel) == (SomeLoopLevel InlinedLoopLevel) = True+ (SomeLoopLevel RootLoopLevel) == (SomeLoopLevel RootLoopLevel) = True+ (SomeLoopLevel a@(LoopLevel _)) == (SomeLoopLevel b@(LoopLevel _)) = a == b+ _ == _ = False++instance Eq (LoopLevel t) where+ level1 == level2 =+ toBool . unsafePerformIO $+ withCxxLoopLevel level1 $ \l1 ->+ withCxxLoopLevel level2 $ \l2 ->+ [CU.exp| bool { *$(const Halide::LoopLevel* l1) == *$(const Halide::LoopLevel* l2) } |]++instance Show (LoopLevel t) where+ showsPrec _ InlinedLoopLevel = showString "InlinedLoopLevel"+ showsPrec _ RootLoopLevel = showString "RootLoopLevel"+ showsPrec d level@(LoopLevel _) =+ showParen (d > 10) $+ showString "LoopLevel {func = "+ . shows (level.func :: Text)+ . showString ", var = "+ . shows (level.var :: Expr Int32)+ . showString "}"++-- desc+-- where+-- desc = unpack . unsafePerformIO $+-- withCxxLoopLevel level $ \l ->+-- peekAndDeleteCxxString+-- =<< [C.throwBlock| std::string* {+-- return handle_halide_exceptions([=](){+-- return new std::string{$(const Halide::LoopLevel* l)->to_string()};+-- });+-- } |]++-- | Different ways to handle the case when the start/end of the loops of stages computed with (fused)+-- are not aligned.+data LoopAlignStrategy+ = -- | Shift the start of the fused loops to align.+ LoopAlignStart+ | -- | Shift the end of the fused loops to align.+ LoopAlignEnd+ | -- | 'computeWith' will make no attempt to align the start/end of the fused loops.+ LoopNoAlign+ | -- | By default, LoopAlignStrategy is set to 'LoopNoAlign'.+ LoopAlignAuto+ deriving stock (Show, Eq, Ord)++instance Enum LoopAlignStrategy where+ fromEnum =+ fromIntegral . \case+ LoopAlignStart -> [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::AlignStart) } |]+ LoopAlignEnd -> [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::AlignEnd) } |]+ LoopNoAlign -> [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::NoAlign) } |]+ LoopAlignAuto -> [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::Auto) } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::AlignStart) } |] = LoopAlignStart+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::AlignEnd) } |] = LoopAlignEnd+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::NoAlign) } |] = LoopNoAlign+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::LoopAlignStrategy::Auto) } |] = LoopAlignAuto+ | otherwise = error $ "invalid LoopAlignStrategy: " <> show k++isInlined :: LoopLevel t -> Bool+isInlined InlinedLoopLevel = True+isInlined _ = False++isRoot :: LoopLevel t -> Bool+isRoot RootLoopLevel = True+isRoot _ = False++instance HasField "func" (LoopLevel 'LockedTy) Text where+ getField level = unsafePerformIO $+ withCxxLoopLevel level $ \level' ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* {+ new std::string{$(const Halide::LoopLevel* level')->func()} } |]++instance HasField "var" (LoopLevel 'LockedTy) (Expr Int32) where+ getField level = unsafePerformIO $+ withCxxLoopLevel level $ \level' ->+ wrapCxxVarOrRVar+ =<< [CU.exp| Halide::VarOrRVar* {+ new Halide::VarOrRVar{$(const Halide::LoopLevel* level')->var()} } |]++wrapCxxLoopLevel :: Ptr CxxLoopLevel -> IO SomeLoopLevel+wrapCxxLoopLevel p = do+ [C.throwBlock| void { handle_halide_exceptions([=]() { $(Halide::LoopLevel* p)->lock(); }); } |]+ inlined <-+ toBool+ <$> [C.throwBlock| bool {+ return handle_halide_exceptions([=](){+ return $(const Halide::LoopLevel* p)->is_inlined(); });+ } |]+ root <-+ toBool+ <$> [C.throwBlock| bool {+ return handle_halide_exceptions([=](){+ return $(const Halide::LoopLevel* p)->is_root(); });+ } |]+ let level+ | inlined = [CU.exp| void { delete $(Halide::LoopLevel *p) } |] >> pure (SomeLoopLevel InlinedLoopLevel)+ | root = [CU.exp| void { delete $(Halide::LoopLevel *p) } |] >> pure (SomeLoopLevel RootLoopLevel)+ | otherwise = do+ let deleter = [C.funPtr| void deleteLoopLevel(Halide::LoopLevel* p) { delete p; } |]+ SomeLoopLevel . LoopLevel <$> newForeignPtr deleter p+ level++withCxxLoopLevel :: LoopLevel t -> (Ptr CxxLoopLevel -> IO a) -> IO a+withCxxLoopLevel (LoopLevel fp) action = withForeignPtr fp action+withCxxLoopLevel level action = do+ let allocate+ | isInlined level = [CU.exp| Halide::LoopLevel* { new Halide::LoopLevel{Halide::LoopLevel::inlined()} } |]+ | isRoot level = [CU.exp| Halide::LoopLevel* { new Halide::LoopLevel{Halide::LoopLevel::root()} } |]+ | otherwise = error "this should never happen"+ destroy p = [CU.exp| void { delete $(Halide::LoopLevel *p) } |]+ bracket allocate destroy action
+ src/Language/Halide/Prelude.hs view
@@ -0,0 +1,41 @@+module Language.Halide.Prelude+ ( (==)+ , (/=)+ , (+)+ , (-)+ )+where++import Data.Kind (Type)+import Language.Halide.Expr+import Language.Halide.Type+import Prelude (Bool, undefined)++type family Promoted a b :: Type++infix 4 ==, /=++(==) :: Expr a -> Expr b -> Expr Bool+(==) = undefined++(/=) :: Expr a -> Expr b -> Expr Bool+(/=) = undefined++infix 6 +, -++(+) :: Expr a -> Expr b -> Expr (Promoted a b)+(+) = undefined++(-) :: Expr a -> Expr b -> Expr (Promoted a b)+(-) = undefined++infix 7 *, /++(*) :: Expr a -> Expr b -> Expr (Promoted a b)+(*) = undefined++(/) :: Expr a -> Expr b -> Expr (Promoted a b)+(/) = undefined++mkExpr :: IsHalideType a => a -> Expr a+mkExpr = undefined
+ src/Language/Halide/RedundantConstraints.hs view
@@ -0,0 +1,18 @@+{-# LANGUAGE ConstraintKinds #-}+{-# OPTIONS_GHC -Wno-redundant-constraints #-}++module Language.Halide.RedundantConstraints+ ( keepRedundantConstraint+ ) where++-- | Can be used to silence individual "redundant constraint" warnings+--+-- > foo :: ConstraintUsefulForDebugging => ...+-- > foo =+-- > ..+-- > where+-- > _ = keepRedundantConstraint (Proxy @ConstraintUsefulForDebugging))+--+-- __Note:__ this function is taken from [input-output-hk/ouroboros-network](https://github.com/input-output-hk/ouroboros-network/blob/master/ouroboros-consensus/src/Ouroboros/Consensus/Util/RedundantConstraints.hs).+keepRedundantConstraint :: c => proxy c -> ()+keepRedundantConstraint _ = ()
+ src/Language/Halide/Schedule.hs view
@@ -0,0 +1,562 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE DuplicateRecordFields #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}++module Language.Halide.Schedule+ ( Dim (..)+ , DimType (..)+ , ForType (..)+ , SplitContents (..)+ , FuseContents (..)+ , Split (..)+ , Bound (..)+ , StorageDim (..)+ , FusedPair (..)+ , FuseLoopLevel (..)+ , StageSchedule (..)+ , ReductionVariable (..)+ , PrefetchDirective (..)+ , getStageSchedule+ -- , getStageSchedule+ -- , getFusedPairs+ -- , getReductionVariables+ -- , getFuseLoopLevel+ -- , getDims+ -- , getSplits+ , AutoScheduler (..)+ , loadAutoScheduler+ , applyAutoScheduler+ , getHalideLibraryPath+ , applySplits+ , applyDims+ , applySchedule+ )+where++import Control.Monad (void)+import Data.Text (Text, pack, unpack)+import qualified Data.Text as T+import Data.Text.Encoding (encodeUtf8)+import Foreign.C.Types (CInt (..))+import Foreign.ForeignPtr+import Foreign.Marshal (allocaArray, peekArray, toBool)+import Foreign.Ptr (Ptr, nullPtr)+import Foreign.Storable+import GHC.TypeLits+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Cpp.Exception as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Context+import Language.Halide.Expr+import Language.Halide.Func+import Language.Halide.LoopLevel+import Language.Halide.Target+import Language.Halide.Type+import Language.Halide.Utils+import System.FilePath (takeDirectory)+import Prelude hiding (tail)++#if USE_DLOPEN+import qualified System.Posix.DynamicLinker as DL++loadLibrary :: Text -> IO ()+loadLibrary path = do+ _ <- DL.dlopen (unpack path) [DL.RTLD_LAZY]+ pure ()++#else+import qualified System.Win32.DLL as Win32++loadLibrary :: Text -> IO ()+loadLibrary path = do+ _ <- Win32.loadLibrary (unpack path)+ pure ()+#endif++-- | Type of dimension that tells which transformations are legal on it.+data DimType = DimPureVar | DimPureRVar | DimImpureRVar+ deriving stock (Show, Eq)++-- | Specifies how loop values are traversed.+data ForType+ = ForSerial+ | ForParallel+ | ForVectorized+ | ForUnrolled+ | ForExtern+ | ForGPUBlock+ | ForGPUThread+ | ForGPULane+ deriving stock (Show, Eq)++data Dim = Dim {var :: !Text, forType :: !ForType, deviceApi :: !DeviceAPI, dimType :: !DimType}+ deriving stock (Show, Eq)++data FuseContents = FuseContents+ { fuseOuter :: !Text+ , fuseInner :: !Text+ , fuseNew :: !Text+ }+ deriving stock (Show, Eq)++data SplitContents = SplitContents+ { splitOld :: !Text+ , splitOuter :: !Text+ , splitInner :: !Text+ , splitFactor :: !(Expr Int32)+ , splitExact :: !Bool+ , splitTail :: !TailStrategy+ }+ deriving stock (Show)++data Split+ = SplitVar !SplitContents+ | FuseVars !FuseContents+ deriving stock (Show)++data Bound = Bound+ { boundVar :: !Text+ , boundMin :: !(Maybe (Expr Int32))+ , boundExtent :: !(Expr Int32)+ , boundModulus :: !(Maybe (Expr Int32))+ , boundRemainder :: !(Maybe (Expr Int32))+ }+ deriving stock (Show)++data StorageDim = StorageDim+ { storageVar :: !Text+ , storageAlignment :: !(Maybe (Expr Int32))+ , storageBound :: !(Maybe (Expr Int32))+ , storageFold :: !(Maybe (Expr Int32, Bool))+ }+ deriving stock (Show)++data FusedPair = FusedPair !Text !(Text, Int) !(Text, Int)+ deriving stock (Show, Eq)++data FuseLoopLevel = FuseLoopLevel !SomeLoopLevel+ deriving stock (Show, Eq)++data ReductionVariable = ReductionVariable {varName :: !Text, minExpr :: !(Expr Int32), extentExpr :: !(Expr Int32)}+ deriving stock (Show)++data PrefetchBoundStrategy+ = PrefetchClamp+ | PrefetchGuardWithIf+ | PrefetchNonFaulting+ deriving stock (Show, Eq)++data PrefetchDirective = PrefetchDirective+ { prefetchFunc :: !Text+ , prefetchAt :: !Text+ , prefetchFrom :: !Text+ , prefetchOffset :: !(Expr Int32)+ , prefetchStrategy :: !PrefetchBoundStrategy+ , prefetchParameter :: !(Maybe (ForeignPtr CxxParameter))+ }+ deriving stock (Show)++data StageSchedule = StageSchedule+ { rvars :: ![ReductionVariable]+ , splits :: ![Split]+ , dims :: ![Dim]+ , prefetches :: ![PrefetchDirective]+ , fuseLevel :: !FuseLoopLevel+ , fusedPairs :: ![FusedPair]+ , allowRaceConditions :: !Bool+ , atomic :: !Bool+ , overrideAtomicAssociativityTest :: !Bool+ }+ deriving stock (Show)++importHalide++instanceHasCxxVector "Halide::Internal::Dim"+instanceHasCxxVector "Halide::Internal::Split"+instanceHasCxxVector "Halide::Internal::FusedPair"+instanceHasCxxVector "Halide::Internal::ReductionVariable"++instance Enum ForType where+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::Serial) } |] =+ ForSerial+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::Parallel) } |] =+ ForParallel+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::Vectorized) } |] =+ ForVectorized+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::Unrolled) } |] =+ ForUnrolled+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::Extern) } |] =+ ForExtern+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::GPUBlock) } |] =+ ForGPUBlock+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::GPUThread) } |] =+ ForGPUThread+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::ForType::GPULane) } |] =+ ForGPULane+ | otherwise = error $ "invalid ForType: " <> show k+ fromEnum = error "Enum instance for ForType does not implement fromEnum"++instance Enum DimType where+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::DimType::PureVar) } |] =+ DimPureVar+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::DimType::PureRVar) } |] =+ DimPureRVar+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::Internal::DimType::ImpureRVar) } |] =+ DimImpureRVar+ | otherwise = error $ "invalid DimType: " <> show k+ fromEnum = error "Enum instance for DimType does not implement fromEnum"++instance Enum PrefetchBoundStrategy where+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::PrefetchBoundStrategy::Clamp) } |] =+ PrefetchClamp+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::PrefetchBoundStrategy::GuardWithIf) } |] =+ PrefetchGuardWithIf+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::PrefetchBoundStrategy::NonFaulting) } |] =+ PrefetchNonFaulting+ | otherwise = error $ "invalid PrefetchBoundStrategy: " <> show k+ fromEnum = error "Enum instance for ForType does not implement fromEnum"++instance Storable FusedPair where+ sizeOf _ = fromIntegral [CU.pure| size_t { sizeof(Halide::Internal::FusedPair) } |]+ alignment _ = fromIntegral [CU.pure| size_t { alignof(Halide::Internal::FusedPair) } |]+ peek p = do+ func1 <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::FusedPair* p)->func_1 } |]+ func2 <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::FusedPair* p)->func_2 } |]+ stage1 <-+ fromIntegral+ <$> [CU.exp| size_t { $(const Halide::Internal::FusedPair* p)->stage_1 } |]+ stage2 <-+ fromIntegral+ <$> [CU.exp| size_t { $(const Halide::Internal::FusedPair* p)->stage_2 } |]+ varName <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::FusedPair* p)->var_name } |]+ pure $ FusedPair varName (func1, stage1) (func2, stage2)+ poke _ _ = error "Storable instance of FusedPair does not implement poke"++instance Storable ReductionVariable where+ sizeOf _ = fromIntegral [CU.pure| size_t { sizeof(Halide::Internal::ReductionVariable) } |]+ alignment _ = fromIntegral [CU.pure| size_t { alignof(Halide::Internal::ReductionVariable) } |]+ peek p = do+ varName <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::ReductionVariable* p)->var } |]+ minExpr <-+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::ReductionVariable* p)->min} } |]+ extentExpr <-+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::ReductionVariable* p)->extent} } |]+ pure $ ReductionVariable varName minExpr extentExpr+ poke _ _ = error "Storable instance of ReductionVariable does not implement poke"++instance Storable PrefetchDirective where+ sizeOf _ = fromIntegral [CU.pure| size_t { sizeof(Halide::Internal::PrefetchDirective) } |]+ alignment _ = fromIntegral [CU.pure| size_t { alignof(Halide::Internal::PrefetchDirective) } |]+ peek p = do+ funcName' <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::PrefetchDirective* p)->name } |]+ atVar' <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::PrefetchDirective* p)->at } |]+ fromVar' <-+ peekCxxString+ =<< [CU.exp| const std::string* { &$(const Halide::Internal::PrefetchDirective* p)->from } |]+ offset' <-+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::PrefetchDirective* p)->offset} } |]+ strategy' <-+ toEnum . fromIntegral+ <$> [CU.exp| int { static_cast<int>($(const Halide::Internal::PrefetchDirective* p)->strategy) } |]+ -- isDefined <-+ -- toBool+ -- <$> [CU.exp| bool { $(const Halide::Internal::PrefetchDirective* p)->param.defined() } |]+ -- param' <-+ -- if isDefined+ -- then+ -- fmap Just $+ -- wrapCxxParameter+ -- =<< [CU.exp| Halide::Internal::Parameter* {+ -- new Halide::Internal::Parameter{$(const Halide::Internal::PrefetchDirective* p)->param} } |]+ -- else pure Nothing+ pure $ PrefetchDirective funcName' atVar' fromVar' offset' strategy' Nothing+ poke _ _ = error "Storable instance for PrefetchDirective does not implement poke"++getReductionVariables :: Ptr CxxStageSchedule -> IO [ReductionVariable]+getReductionVariables schedule =+ peekCxxVector+ =<< [CU.exp| const std::vector<Halide::Internal::ReductionVariable>* {+ &$(const Halide::Internal::StageSchedule* schedule)->rvars() } |]++getSplits :: Ptr CxxStageSchedule -> IO [Split]+getSplits schedule =+ peekCxxVector+ =<< [CU.exp| const std::vector<Halide::Internal::Split>* {+ &$(const Halide::Internal::StageSchedule* schedule)->splits() } |]++getDims :: Ptr CxxStageSchedule -> IO [Dim]+getDims schedule =+ peekCxxVector+ =<< [CU.exp| const std::vector<Halide::Internal::Dim>* {+ &$(const Halide::Internal::StageSchedule* schedule)->dims() } |]++getFuseLoopLevel :: Ptr CxxStageSchedule -> IO FuseLoopLevel+getFuseLoopLevel schedule =+ fmap FuseLoopLevel $+ wrapCxxLoopLevel+ =<< [CU.exp| Halide::LoopLevel* {+ new Halide::LoopLevel{$(const Halide::Internal::StageSchedule* schedule)->fuse_level().level}+ } |]++getFusedPairs :: Ptr CxxStageSchedule -> IO [FusedPair]+getFusedPairs schedule = do+ peekCxxVector+ =<< [CU.exp| const std::vector<Halide::Internal::FusedPair>* {+ &$(const Halide::Internal::StageSchedule* schedule)->fused_pairs() } |]++peekStageSchedule :: Ptr CxxStageSchedule -> IO StageSchedule+peekStageSchedule schedule = do+ rvars' <- getReductionVariables schedule+ splits' <- getSplits schedule+ dims' <- getDims schedule+ let prefetches' = []+ fuseLevel' <- getFuseLoopLevel schedule+ fusedPairs' <- getFusedPairs schedule+ allowRaceConditions' <-+ toBool+ <$> [CU.exp| bool { $(const Halide::Internal::StageSchedule* schedule)->allow_race_conditions() } |]+ atomic' <-+ toBool+ <$> [CU.exp| bool { $(const Halide::Internal::StageSchedule* schedule)->atomic() } |]+ overrideAtomicAssociativityTest' <-+ toBool+ <$> [CU.exp| bool { $(const Halide::Internal::StageSchedule* schedule)->override_atomic_associativity_test() } |]+ pure $+ StageSchedule+ { rvars = rvars'+ , splits = splits'+ , dims = dims'+ , prefetches = prefetches'+ , fuseLevel = fuseLevel'+ , fusedPairs = fusedPairs'+ , allowRaceConditions = allowRaceConditions'+ , atomic = atomic'+ , overrideAtomicAssociativityTest = overrideAtomicAssociativityTest'+ }++instance Storable Dim where+ sizeOf _ = fromIntegral [CU.pure| size_t { sizeof(Halide::Internal::Dim) } |]+ alignment _ = fromIntegral [CU.pure| size_t { alignof(Halide::Internal::Dim) } |]+ peek p = do+ name <- peekCxxString =<< [CU.exp| const std::string* { &$(Halide::Internal::Dim* p)->var } |]+ forType' <-+ toEnum . fromIntegral+ <$> [CU.exp| int { static_cast<int>($(Halide::Internal::Dim* p)->for_type) } |]+ device <-+ toEnum . fromIntegral+ <$> [CU.exp| int { static_cast<int>($(Halide::Internal::Dim* p)->device_api) } |]+ dimType' <-+ toEnum . fromIntegral+ <$> [CU.exp| int { static_cast<int>($(Halide::Internal::Dim* p)->dim_type) } |]+ pure $ Dim name forType' device dimType'+ poke _ = error "Storable instance for Dim does not implement poke"++peekOld :: Ptr Split -> IO Text+peekOld p = peekCxxString =<< [CU.exp| const std::string* { &$(const Halide::Internal::Split* p)->old_var } |]++peekOuter :: Ptr Split -> IO Text+peekOuter p = peekCxxString =<< [CU.exp| const std::string* { &$(const Halide::Internal::Split* p)->outer } |]++peekInner :: Ptr Split -> IO Text+peekInner p = peekCxxString =<< [CU.exp| const std::string* { &$(const Halide::Internal::Split* p)->inner } |]++peekFactor :: Ptr Split -> IO (Expr Int32)+peekFactor p =+ cxxConstructExpr $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+ $(const Halide::Internal::Split* p)->factor} } |]++instance Storable Split where+ sizeOf _ = fromIntegral [CU.pure| size_t { sizeof(Halide::Internal::Split) } |]+ alignment _ = fromIntegral [CU.pure| size_t { alignof(Halide::Internal::Split) } |]+ peek p = do+ isRename <- toBool <$> [CU.exp| bool { $(const Halide::Internal::Split* p)->is_rename() } |]+ isSplit <- toBool <$> [CU.exp| bool { $(const Halide::Internal::Split* p)->is_split() } |]+ isFuse <- toBool <$> [CU.exp| bool { $(const Halide::Internal::Split* p)->is_fuse() } |]+ isPurify <- toBool <$> [CU.exp| bool { $(const Halide::Internal::Split* p)->is_purify() } |]+ let r+ | isSplit =+ fmap SplitVar $+ SplitContents+ <$> peekOld p+ <*> peekOuter p+ <*> peekInner p+ <*> peekFactor p+ <*> (toBool <$> [CU.exp| bool { $(const Halide::Internal::Split* p)->exact } |])+ <*> fmap+ (toEnum . fromIntegral)+ [CU.exp| int { static_cast<int>($(const Halide::Internal::Split* p)->tail) } |]+ | isFuse =+ fmap FuseVars $+ FuseContents+ <$> peekOuter p+ <*> peekInner p+ <*> peekOld p+ | isRename = error "renames are not yet implemented"+ | isPurify = error "purify is not yet implemented"+ | otherwise = error "invalid split type"+ r+ poke _ = error "Storable instance for Split does not implement poke"++-- wrapCxxStageSchedule :: Ptr CxxStageSchedule -> IO StageSchedule+-- wrapCxxStageSchedule = fmap StageSchedule . newForeignPtr deleter+-- where+-- deleter =+-- [C.funPtr| void deleteSchedule(Halide::Internal::StageSchedule* p) {+-- std::cout << "deleting ..." << std::endl;+-- delete p; } |]++getStageSchedule :: (KnownNat n, IsHalideType a) => Stage n a -> IO StageSchedule+getStageSchedule stage =+ withCxxStage stage $ \stage' ->+ peekStageSchedule+ =<< [CU.exp| const Halide::Internal::StageSchedule* {+ &$(const Halide::Stage* stage')->get_schedule() } |]++#if USE_DLOPEN++getHalideLibraryPath :: IO (Maybe Text)+getHalideLibraryPath = do+ ptr <-+ [CU.block| std::string* {+ Dl_info info;+ if (dladdr((void const*)&Halide::load_plugin, &info) != 0 && info.dli_sname != nullptr) {+ auto symbol = dlsym(RTLD_NEXT, info.dli_sname);+ if (dladdr(symbol, &info) != 0 && info.dli_fname != nullptr) {+ return new std::string{info.dli_fname};+ }+ }+ return nullptr;+ } |]+ if ptr == nullPtr+ then pure Nothing+ else Just . pack . takeDirectory . unpack <$> peekAndDeleteCxxString ptr++#else++getHalideLibraryPath :: IO (Maybe Text)+getHalideLibraryPath = pure Nothing++#endif++data AutoScheduler+ = Adams2019+ | Li2018+ | Mullapudi2016+ deriving stock (Eq, Show)++loadAutoScheduler :: AutoScheduler -> IO ()+loadAutoScheduler scheduler = do+ lib <- getHalideLibraryPath+ let prepare s+ | Just dir <- lib = dir <> "/lib" <> s <> ".so"+ | Nothing <- lib = "lib" <> s <> ".so"+ path = prepare $+ case scheduler of+ Adams2019 -> "autoschedule_adams2019"+ Li2018 -> "autoschedule_li2018"+ Mullapudi2016 -> "autoschedule_mullapudi2016"+ loadLibrary path++applyAutoScheduler :: (KnownNat n, IsHalideType a) => AutoScheduler -> Target -> Func t n a -> IO Text+applyAutoScheduler scheduler target func = do+ let s = encodeUtf8 . pack . show $ scheduler+ withFunc func $ \f ->+ withCxxTarget target $ \t -> do+ peekAndDeleteCxxString+ =<< [C.throwBlock| std::string* {+ return handle_halide_exceptions([=](){+ auto name = std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)};+ auto pipeline = Halide::Pipeline{*$(Halide::Func* f)};+ auto params = Halide::AutoschedulerParams{name};+ auto results = pipeline.apply_autoscheduler(*$(Halide::Target* t), params);+ return new std::string{std::move(results.schedule_source)};+ });+ } |]++makeUnqualified :: Text -> Text+makeUnqualified = snd . T.breakOnEnd "."++applySplit :: (KnownNat n, IsHalideType a) => Split -> Stage n a -> IO ()+applySplit (SplitVar x) stage = do+ oldVar <- mkVar (makeUnqualified x.splitOld)+ outerVar <- mkVar (makeUnqualified x.splitOuter)+ innerVar <- mkVar (makeUnqualified x.splitInner)+ void $ Language.Halide.Func.split x.splitTail oldVar (outerVar, innerVar) x.splitFactor stage+applySplit (FuseVars x) stage = do+ newVar <- mkVar (makeUnqualified x.fuseNew)+ innerVar <- mkVar (makeUnqualified x.fuseInner)+ outerVar <- mkVar (makeUnqualified x.fuseOuter)+ void $ Language.Halide.Func.fuse (innerVar, outerVar) newVar stage++applySplits :: (KnownNat n, IsHalideType a) => [Split] -> Stage n a -> IO ()+applySplits xs stage = mapM_ (`applySplit` stage) xs++applyDim :: (KnownNat n, IsHalideType a) => Dim -> Stage n a -> IO ()+applyDim x stage = do+ var' <- mkVar (makeUnqualified x.var)+ void $+ case x.forType of+ ForSerial -> pure stage+ ForParallel -> parallel var' stage+ ForVectorized -> vectorize var' stage+ ForUnrolled -> unroll var' stage+ ForExtern -> error "extern ForType is not yet supported by applyDim"+ ForGPUBlock -> gpuBlocks x.deviceApi var' stage+ ForGPUThread -> gpuThreads x.deviceApi var' stage+ ForGPULane -> gpuLanes x.deviceApi var' stage++applyDims :: (KnownNat n, IsHalideType a) => [Dim] -> Stage n a -> IO ()+applyDims xs stage = do+ mapM_ (`applyDim` stage) xs+ vars <- mapM (mkVar . makeUnqualified . (.var)) xs+ void $ reorder vars stage++applySchedule :: (KnownNat n, IsHalideType a) => StageSchedule -> Stage n a -> IO ()+applySchedule schedule stage = do+ applySplits schedule.splits stage+ applyDims schedule.dims stage++-- data SplitContents = SplitContents+-- { old :: !Text+-- , outer :: !Text+-- , inner :: !Text+-- , factor :: !(Maybe Int)+-- , exact :: !Bool+-- , tail :: !TailStrategy+-- }+-- deriving stock (Show, Eq)+--+--+-- applySplits :: [Split] -> Stage n a -> IO ()+-- applySplits splits stage =++-- v <-+-- [CU.exp| Halide::Internal::Dim* {+-- $(Halide::Internal::StageSchedule* schedule)->dims().data() } |]+-- putStrLn $ "n = " <> show n+-- mapM (\i -> print i >> peekElemOff v i) [0 .. n - 1]
+ src/Language/Halide/Target.hs view
@@ -0,0 +1,518 @@+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeApplications #-}++-- |+-- Module : Language.Halide.Target+-- Description : Compilation targets and their features+-- Copyright : (c) Tom Westerhout, 2023+--+-- This module defines a data type 'Target' that represents the compilation target.+-- This could be the host target, or a CUDA device with CUDA capability of at least 7,+-- or an OpenCL device, etc.+--+-- You would typically start with 'hostTarget' and then extend it using various 'TargetFeature's.+-- This can be done with the 'setFeature' function.+module Language.Halide.Target+ ( Target (..)+ , hostTarget+ , gpuTarget+ , hostSupportsTargetDevice+ , setFeature+ , hasGpuFeature+ , TargetFeature (..)+ , DeviceAPI (..)+ -- , targetFeatureForDeviceAPI++ -- * Internal+ , withCxxTarget+ , testCUDA+ , testOpenCL+ )+where++import Data.Text (unpack)+import Foreign.ForeignPtr+import Foreign.Ptr (Ptr)+import GHC.IO (unsafePerformIO)+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Cpp.Exception as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Context+import Language.Halide.Type+import Language.Halide.Utils+import Prelude hiding (tail)++importHalide++-- | The compilation target.+--+-- This is the Haskell counterpart of [@Halide::Target@](https://halide-lang.org/docs/struct_halide_1_1_target.html).+newtype Target = Target (ForeignPtr CxxTarget)++instance Eq Target where+ (==) target1 target2 =+ toEnum . fromIntegral . unsafePerformIO $+ withCxxTarget target1 $ \t1 ->+ withCxxTarget target2 $ \t2 ->+ [CU.exp| bool { *$(const Halide::Target* t1) == *$(const Halide::Target* t2) } |]++instance Show Target where+ show target =+ unpack . unsafePerformIO $ withCxxTarget target $ \t ->+ peekAndDeleteCxxString+ =<< [CU.exp| std::string* {+ new std::string{$(const Halide::Target* t)->to_string()} } |]++-- | Return the target that Halide will use by default.+--+-- If the @HL_TARGET@ environment variable is set, it uses that. Otherwise, it+-- returns the target corresponding to the host machine.+hostTarget :: Target+hostTarget =+ unsafePerformIO $+ wrapCxxTarget+ =<< [CU.exp| Halide::Target* { new Halide::Target{Halide::get_target_from_environment()} } |]+{-# NOINLINE hostTarget #-}++-- | Get the default GPU target. We first check for CUDA and then for OpenCL.+-- If neither of the two is usable, 'Nothing' is returned.+gpuTarget :: Maybe Target+gpuTarget+ | hostSupportsTargetDevice cudaTarget = Just cudaTarget+ | hostSupportsTargetDevice openCLTarget = Just openCLTarget+ | otherwise = Nothing+ where+ openCLTarget = setFeature FeatureOpenCL hostTarget+ cudaTarget = setFeature FeatureCUDA hostTarget++-- | Attempt to sniff whether a given 'Target' (and its implied 'DeviceAPI') is usable on the+-- current host.+--+-- __Note__ that a return value of @True@ does not guarantee that future usage of that device will+-- succeed; it is intended mainly as a simple diagnostic to allow early-exit when a desired device+-- is definitely not usable.+--+-- Also note that this call is __NOT threadsafe__, as it temporarily redirects various global+-- error-handling hooks in Halide.+hostSupportsTargetDevice+ :: Target+ -> Bool+ -- ^ Whether the target appears to be usable+hostSupportsTargetDevice target =+ unsafePerformIO . fmap (toEnum . fromIntegral) $+ withCxxTarget target $ \t ->+ [CU.exp| bool { Halide::host_supports_target_device(*$(Halide::Target* t)) } |]++-- | Add a feature to target.+setFeature+ :: TargetFeature+ -- ^ Feature to add+ -> Target+ -- ^ Initial target+ -> Target+ -- ^ New target+setFeature feature target = unsafePerformIO $+ withCxxTarget target $ \t ->+ wrapCxxTarget+ =<< [CU.exp| Halide::Target* {+ new Halide::Target{$(Halide::Target* t)->with_feature(+ static_cast<Halide::Target::Feature>($(int f)))} + } |]+ where+ f = fromIntegral . fromEnum $ feature++-- | Return whether a GPU compute runtime is enabled.+--+-- Checks whether 'Language.Halide.Func.gpuBlocks' and similar are going to work.+--+-- For more info, see [@Target::has_gpu_feature@](https://halide-lang.org/docs/struct_halide_1_1_target.html#a22bf80aa6dc3a700c9732050d2341a80).+hasGpuFeature :: Target -> Bool+hasGpuFeature target =+ unsafePerformIO . fmap (toEnum . fromIntegral) $+ withCxxTarget target $ \t ->+ [CU.exp| bool { $(Halide::Target* t)->has_gpu_feature() } |]++-- | An enum describing the type of device API.+--+-- This is the Haskell counterpart of [@Halide::DeviceAPI@](https://halide-lang.org/docs/namespace_halide.html#aa26c7f430d2b1c44ba3e1d3f6df2ba6e).+data DeviceAPI+ = DeviceNone+ | DeviceHost+ | DeviceDefaultGPU+ | DeviceCUDA+ | DeviceOpenCL+ | DeviceOpenGLCompute+ | DeviceMetal+ | DeviceHexagon+ | DeviceHexagonDma+ | DeviceD3D12Compute+ deriving stock (Show, Eq, Ord)++instance Enum DeviceAPI where+ fromEnum =+ fromIntegral . \case+ DeviceNone -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::None) } |]+ DeviceHost -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Host) } |]+ DeviceDefaultGPU -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Default_GPU) } |]+ DeviceCUDA -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::CUDA) } |]+ DeviceOpenCL -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::OpenCL) } |]+ DeviceOpenGLCompute -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::OpenGLCompute) } |]+ DeviceMetal -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Metal) } |]+ DeviceHexagon -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Hexagon) } |]+ DeviceHexagonDma -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::HexagonDma) } |]+ DeviceD3D12Compute -> [CU.pure| int { static_cast<int>(Halide::DeviceAPI::D3D12Compute) } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::None) } |] = DeviceNone+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Host) } |] = DeviceHost+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Default_GPU) } |] = DeviceDefaultGPU+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::CUDA) } |] = DeviceCUDA+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::OpenCL) } |] = DeviceOpenCL+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::OpenGLCompute) } |] = DeviceOpenGLCompute+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Metal) } |] = DeviceMetal+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::Hexagon) } |] = DeviceHexagon+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::HexagonDma) } |] = DeviceHexagonDma+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::D3D12Compute) } |] = DeviceD3D12Compute+ | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::DeviceAPI::D3D12Compute) } |] = DeviceD3D12Compute+ | otherwise = error $ "invalid DeviceAPI: " <> show k++wrapCxxTarget :: Ptr CxxTarget -> IO Target+wrapCxxTarget = fmap Target . newForeignPtr deleter+ where+ deleter = [C.funPtr| void deleteTarget(Halide::Target* p) { delete p; } |]++-- | Convert 'Target' into @Halide::Target*@ and use it in an 'IO' action.+withCxxTarget :: Target -> (Ptr CxxTarget -> IO a) -> IO a+withCxxTarget (Target fp) = withForeignPtr fp++-- targetFeatureForDeviceAPI :: DeviceAPI -> Maybe TargetFeature+-- targetFeatureForDeviceAPI deviceAPI =+-- toFeature . unsafePerformIO $+-- [CU.block| int {+-- auto feature = Halide::target_feature_for_device_api(+-- static_cast<Halide::DeviceAPI>($(int api)));+-- return (feature == Halide::Target::FeatureEnd) ? (-1) : static_cast<int>(feature);+-- } |]+-- where+-- api = fromIntegral . fromEnum $ deviceAPI+-- toFeature n+-- | n > 0 = Just . toEnum . fromIntegral $ n+-- | otherwise = Nothing++-- | A test that tries to compile and run a Halide pipeline using 'FeatureCUDA'.+--+-- This is implemented fully in C++ to make sure that we test the installation+-- rather than our Haskell code.+--+-- On non-NixOS systems one should do the following:+--+-- > nixGLNvidia cabal repl --ghc-options='-fobject-code -O0'+-- > ghci> testCUDA+testCUDA :: IO ()+testCUDA = do+ [C.throwBlock| void {+ handle_halide_exceptions([](){+ // Define a gradient function.+ Halide::Func f;+ Halide::Var x, y, xo, xi, yo, yi;+ f(x, y) = x + y;+ // Schedule f on the GPU in 16x16 tiles.+ f.gpu_tile(x, y, xo, yo, xi, yi, 16, 16);+ // Construct a target that uses the GPU.+ Halide::Target target = Halide::get_host_target();+ // Set CUDA as the GPU backend.+ target.set_feature(Halide::Target::CUDA);+ // Enable debugging so that you can see what CUDA API calls we do.+ target.set_feature(Halide::Target::Debug);+ // JIT-compile the pipeline.+ f.compile_jit(target);+ // Run it.+ Halide::Buffer<int> result = f.realize({32, 32});+ // Check correctness+ for (int y = 0; y < result.height(); y++) {+ for (int x = 0; x < result.width(); x++) {+ if (result(x, y) != x + y) {+ printf("result(%d, %d) = %d instead of %d\n",+ x, y, result(x, y), x + y);+ }+ }+ }+ });+ } |]++-- | Similar to 'testCUDA' but for 'FeatureOpenCL'.+testOpenCL :: IO ()+testOpenCL = do+ [C.throwBlock| void {+ handle_halide_exceptions([](){+ Halide::Func f;+ Halide::Var x, y, xo, xi, yo, yi;+ f(x, y) = x + y;+ f.gpu_tile(x, y, xo, yo, xi, yi, 4, 4);+ Halide::Target target = Halide::get_host_target();+ target.set_feature(Halide::Target::OpenCL);+ target.set_feature(Halide::Target::Debug);+ fprintf(stderr, "Compiling ...\n");+ f.compile_jit(target);+ fprintf(stderr, "Running on OpenCL ...\n");+ Halide::Buffer<int> result = f.realize({32, 32});+ for (int y = 0; y < result.height(); y++) {+ for (int x = 0; x < result.width(); x++) {+ if (result(x, y) != x + y) {+ printf("result(%d, %d) = %d instead of %d\n",+ x, y, result(x, y), x + y);+ }+ }+ }+ });+ } |]++-- |+--+-- Note: generated automatically using+--+-- > cat $HALIDE_PATH/include/Halide.h | \+-- > grep -E '.* = halide_target_feature_.*' | \+-- > sed -E 's/^\s*(.*) = .*$/ | \1/g' | \+-- > grep -v FeatureEnd+data TargetFeature+ = FeatureJIT+ | FeatureDebug+ | FeatureNoAsserts+ | FeatureNoBoundsQuery+ | FeatureSSE41+ | FeatureAVX+ | FeatureAVX2+ | FeatureFMA+ | FeatureFMA4+ | FeatureF16C+ | FeatureARMv7s+ | FeatureNoNEON+ | FeatureVSX+ | FeaturePOWER_ARCH_2_07+ | FeatureCUDA+ | FeatureCUDACapability30+ | FeatureCUDACapability32+ | FeatureCUDACapability35+ | FeatureCUDACapability50+ | FeatureCUDACapability61+ | FeatureCUDACapability70+ | FeatureCUDACapability75+ | FeatureCUDACapability80+ | FeatureCUDACapability86+ | FeatureOpenCL+ | FeatureCLDoubles+ | FeatureCLHalf+ | FeatureCLAtomics64+ | FeatureOpenGLCompute+ | FeatureEGL+ | FeatureUserContext+ | FeatureProfile+ | FeatureNoRuntime+ | FeatureMetal+ | FeatureCPlusPlusMangling+ | FeatureLargeBuffers+ | FeatureHexagonDma+ | FeatureHVX_128+ | FeatureHVX_v62+ | FeatureHVX_v65+ | FeatureHVX_v66+ | -- Removed in upstream Halide FeatureHVX_shared_object+ FeatureFuzzFloatStores+ | FeatureSoftFloatABI+ | FeatureMSAN+ | FeatureAVX512+ | FeatureAVX512_KNL+ | FeatureAVX512_Skylake+ | FeatureAVX512_Cannonlake+ | FeatureAVX512_SapphireRapids+ | FeatureTraceLoads+ | FeatureTraceStores+ | FeatureTraceRealizations+ | FeatureTracePipeline+ | FeatureD3D12Compute+ | FeatureStrictFloat+ | FeatureTSAN+ | FeatureASAN+ | FeatureCheckUnsafePromises+ | FeatureEmbedBitcode+ | FeatureEnableLLVMLoopOpt+ | FeatureWasmSimd128+ | FeatureWasmSignExt+ | FeatureWasmSatFloatToInt+ | FeatureWasmThreads+ | FeatureWasmBulkMemory+ | FeatureSVE+ | FeatureSVE2+ | FeatureARMDotProd+ | FeatureARMFp16+ | FeatureRVV+ | FeatureARMv81a+ | FeatureSanitizerCoverage+ | FeatureProfileByTimer+ | FeatureSPIRV+ | FeatureSemihosting+ deriving stock (Eq, Show, Ord)++instance Enum TargetFeature where+ -- Generated using+ -- cat $HALIDE_PATH/include/Halide.h | grep -E '.* = halide_target_feature_.*' | grep -v 'FeatureEnd' | sed -E 's/[ \t]+(.*) = ([^,]*).*/ Feature\1 -> [CU.pure| int { \2 } |]/'+ fromEnum =+ fromIntegral . \case+ FeatureJIT -> [CU.pure| int { halide_target_feature_jit } |]+ FeatureDebug -> [CU.pure| int { halide_target_feature_debug } |]+ FeatureNoAsserts -> [CU.pure| int { halide_target_feature_no_asserts } |]+ FeatureNoBoundsQuery -> [CU.pure| int { halide_target_feature_no_bounds_query } |]+ FeatureSSE41 -> [CU.pure| int { halide_target_feature_sse41 } |]+ FeatureAVX -> [CU.pure| int { halide_target_feature_avx } |]+ FeatureAVX2 -> [CU.pure| int { halide_target_feature_avx2 } |]+ FeatureFMA -> [CU.pure| int { halide_target_feature_fma } |]+ FeatureFMA4 -> [CU.pure| int { halide_target_feature_fma4 } |]+ FeatureF16C -> [CU.pure| int { halide_target_feature_f16c } |]+ FeatureARMv7s -> [CU.pure| int { halide_target_feature_armv7s } |]+ FeatureNoNEON -> [CU.pure| int { halide_target_feature_no_neon } |]+ FeatureVSX -> [CU.pure| int { halide_target_feature_vsx } |]+ FeaturePOWER_ARCH_2_07 -> [CU.pure| int { halide_target_feature_power_arch_2_07 } |]+ FeatureCUDA -> [CU.pure| int { halide_target_feature_cuda } |]+ FeatureCUDACapability30 -> [CU.pure| int { halide_target_feature_cuda_capability30 } |]+ FeatureCUDACapability32 -> [CU.pure| int { halide_target_feature_cuda_capability32 } |]+ FeatureCUDACapability35 -> [CU.pure| int { halide_target_feature_cuda_capability35 } |]+ FeatureCUDACapability50 -> [CU.pure| int { halide_target_feature_cuda_capability50 } |]+ FeatureCUDACapability61 -> [CU.pure| int { halide_target_feature_cuda_capability61 } |]+ FeatureCUDACapability70 -> [CU.pure| int { halide_target_feature_cuda_capability70 } |]+ FeatureCUDACapability75 -> [CU.pure| int { halide_target_feature_cuda_capability75 } |]+ FeatureCUDACapability80 -> [CU.pure| int { halide_target_feature_cuda_capability80 } |]+ FeatureCUDACapability86 -> [CU.pure| int { halide_target_feature_cuda_capability86 } |]+ FeatureOpenCL -> [CU.pure| int { halide_target_feature_opencl } |]+ FeatureCLDoubles -> [CU.pure| int { halide_target_feature_cl_doubles } |]+ FeatureCLHalf -> [CU.pure| int { halide_target_feature_cl_half } |]+ FeatureCLAtomics64 -> [CU.pure| int { halide_target_feature_cl_atomic64 } |]+ FeatureOpenGLCompute -> [CU.pure| int { halide_target_feature_openglcompute } |]+ FeatureEGL -> [CU.pure| int { halide_target_feature_egl } |]+ FeatureUserContext -> [CU.pure| int { halide_target_feature_user_context } |]+ FeatureProfile -> [CU.pure| int { halide_target_feature_profile } |]+ FeatureNoRuntime -> [CU.pure| int { halide_target_feature_no_runtime } |]+ FeatureMetal -> [CU.pure| int { halide_target_feature_metal } |]+ FeatureCPlusPlusMangling -> [CU.pure| int { halide_target_feature_c_plus_plus_mangling } |]+ FeatureLargeBuffers -> [CU.pure| int { halide_target_feature_large_buffers } |]+ FeatureHexagonDma -> [CU.pure| int { halide_target_feature_hexagon_dma } |]+ FeatureHVX_128 -> [CU.pure| int { halide_target_feature_hvx_128 } |]+ FeatureHVX_v62 -> [CU.pure| int { halide_target_feature_hvx_v62 } |]+ FeatureHVX_v65 -> [CU.pure| int { halide_target_feature_hvx_v65 } |]+ FeatureHVX_v66 -> [CU.pure| int { halide_target_feature_hvx_v66 } |]+ -- FeatureHVX_shared_object -> [CU.pure| int { halide_target_feature_hvx_use_shared_object } |]+ FeatureFuzzFloatStores -> [CU.pure| int { halide_target_feature_fuzz_float_stores } |]+ FeatureSoftFloatABI -> [CU.pure| int { halide_target_feature_soft_float_abi } |]+ FeatureMSAN -> [CU.pure| int { halide_target_feature_msan } |]+ FeatureAVX512 -> [CU.pure| int { halide_target_feature_avx512 } |]+ FeatureAVX512_KNL -> [CU.pure| int { halide_target_feature_avx512_knl } |]+ FeatureAVX512_Skylake -> [CU.pure| int { halide_target_feature_avx512_skylake } |]+ FeatureAVX512_Cannonlake -> [CU.pure| int { halide_target_feature_avx512_cannonlake } |]+ FeatureAVX512_SapphireRapids -> [CU.pure| int { halide_target_feature_avx512_sapphirerapids } |]+ FeatureTraceLoads -> [CU.pure| int { halide_target_feature_trace_loads } |]+ FeatureTraceStores -> [CU.pure| int { halide_target_feature_trace_stores } |]+ FeatureTraceRealizations -> [CU.pure| int { halide_target_feature_trace_realizations } |]+ FeatureTracePipeline -> [CU.pure| int { halide_target_feature_trace_pipeline } |]+ FeatureD3D12Compute -> [CU.pure| int { halide_target_feature_d3d12compute } |]+ FeatureStrictFloat -> [CU.pure| int { halide_target_feature_strict_float } |]+ FeatureTSAN -> [CU.pure| int { halide_target_feature_tsan } |]+ FeatureASAN -> [CU.pure| int { halide_target_feature_asan } |]+ FeatureCheckUnsafePromises -> [CU.pure| int { halide_target_feature_check_unsafe_promises } |]+ FeatureEmbedBitcode -> [CU.pure| int { halide_target_feature_embed_bitcode } |]+ FeatureEnableLLVMLoopOpt -> [CU.pure| int { halide_target_feature_enable_llvm_loop_opt } |]+ FeatureWasmSimd128 -> [CU.pure| int { halide_target_feature_wasm_simd128 } |]+ FeatureWasmSignExt -> [CU.pure| int { halide_target_feature_wasm_signext } |]+ FeatureWasmSatFloatToInt -> [CU.pure| int { halide_target_feature_wasm_sat_float_to_int } |]+ FeatureWasmThreads -> [CU.pure| int { halide_target_feature_wasm_threads } |]+ FeatureWasmBulkMemory -> [CU.pure| int { halide_target_feature_wasm_bulk_memory } |]+ FeatureSVE -> [CU.pure| int { halide_target_feature_sve } |]+ FeatureSVE2 -> [CU.pure| int { halide_target_feature_sve2 } |]+ FeatureARMDotProd -> [CU.pure| int { halide_target_feature_arm_dot_prod } |]+ FeatureARMFp16 -> [CU.pure| int { halide_target_feature_arm_fp16 } |]+ FeatureRVV -> [CU.pure| int { halide_target_feature_rvv } |]+ FeatureARMv81a -> [CU.pure| int { halide_target_feature_armv81a } |]+ FeatureSanitizerCoverage -> [CU.pure| int { halide_target_feature_sanitizer_coverage } |]+ FeatureProfileByTimer -> [CU.pure| int { halide_target_feature_profile_by_timer } |]+ FeatureSPIRV -> [CU.pure| int { halide_target_feature_spirv } |]+ FeatureSemihosting -> [CU.pure| int { halide_target_feature_semihosting } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { halide_target_feature_jit } |] = FeatureJIT+ | fromIntegral k == [CU.pure| int { halide_target_feature_debug } |] = FeatureDebug+ | fromIntegral k == [CU.pure| int { halide_target_feature_no_asserts } |] = FeatureNoAsserts+ | fromIntegral k == [CU.pure| int { halide_target_feature_no_bounds_query } |] = FeatureNoBoundsQuery+ | fromIntegral k == [CU.pure| int { halide_target_feature_sse41 } |] = FeatureSSE41+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx } |] = FeatureAVX+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx2 } |] = FeatureAVX2+ | fromIntegral k == [CU.pure| int { halide_target_feature_fma } |] = FeatureFMA+ | fromIntegral k == [CU.pure| int { halide_target_feature_fma4 } |] = FeatureFMA4+ | fromIntegral k == [CU.pure| int { halide_target_feature_f16c } |] = FeatureF16C+ | fromIntegral k == [CU.pure| int { halide_target_feature_armv7s } |] = FeatureARMv7s+ | fromIntegral k == [CU.pure| int { halide_target_feature_no_neon } |] = FeatureNoNEON+ | fromIntegral k == [CU.pure| int { halide_target_feature_vsx } |] = FeatureVSX+ | fromIntegral k == [CU.pure| int { halide_target_feature_power_arch_2_07 } |] = FeaturePOWER_ARCH_2_07+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda } |] = FeatureCUDA+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability30 } |] = FeatureCUDACapability30+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability32 } |] = FeatureCUDACapability32+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability35 } |] = FeatureCUDACapability35+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability50 } |] = FeatureCUDACapability50+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability61 } |] = FeatureCUDACapability61+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability70 } |] = FeatureCUDACapability70+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability75 } |] = FeatureCUDACapability75+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability80 } |] = FeatureCUDACapability80+ | fromIntegral k == [CU.pure| int { halide_target_feature_cuda_capability86 } |] = FeatureCUDACapability86+ | fromIntegral k == [CU.pure| int { halide_target_feature_opencl } |] = FeatureOpenCL+ | fromIntegral k == [CU.pure| int { halide_target_feature_cl_doubles } |] = FeatureCLDoubles+ | fromIntegral k == [CU.pure| int { halide_target_feature_cl_half } |] = FeatureCLHalf+ | fromIntegral k == [CU.pure| int { halide_target_feature_cl_atomic64 } |] = FeatureCLAtomics64+ | fromIntegral k == [CU.pure| int { halide_target_feature_openglcompute } |] = FeatureOpenGLCompute+ | fromIntegral k == [CU.pure| int { halide_target_feature_egl } |] = FeatureEGL+ | fromIntegral k == [CU.pure| int { halide_target_feature_user_context } |] = FeatureUserContext+ | fromIntegral k == [CU.pure| int { halide_target_feature_profile } |] = FeatureProfile+ | fromIntegral k == [CU.pure| int { halide_target_feature_no_runtime } |] = FeatureNoRuntime+ | fromIntegral k == [CU.pure| int { halide_target_feature_metal } |] = FeatureMetal+ | fromIntegral k == [CU.pure| int { halide_target_feature_c_plus_plus_mangling } |] = FeatureCPlusPlusMangling+ | fromIntegral k == [CU.pure| int { halide_target_feature_large_buffers } |] = FeatureLargeBuffers+ | fromIntegral k == [CU.pure| int { halide_target_feature_hexagon_dma } |] = FeatureHexagonDma+ | fromIntegral k == [CU.pure| int { halide_target_feature_hvx_128 } |] = FeatureHVX_128+ | fromIntegral k == [CU.pure| int { halide_target_feature_hvx_v62 } |] = FeatureHVX_v62+ | fromIntegral k == [CU.pure| int { halide_target_feature_hvx_v65 } |] = FeatureHVX_v65+ | fromIntegral k == [CU.pure| int { halide_target_feature_hvx_v66 } |] = FeatureHVX_v66+ -- \| fromIntegral k == [CU.pure| int { halide_target_feature_hvx_use_shared_object } |] = FeatureHVX_shared_object+ | fromIntegral k == [CU.pure| int { halide_target_feature_fuzz_float_stores } |] = FeatureFuzzFloatStores+ | fromIntegral k == [CU.pure| int { halide_target_feature_soft_float_abi } |] = FeatureSoftFloatABI+ | fromIntegral k == [CU.pure| int { halide_target_feature_msan } |] = FeatureMSAN+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx512 } |] = FeatureAVX512+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx512_knl } |] = FeatureAVX512_KNL+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx512_skylake } |] = FeatureAVX512_Skylake+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx512_cannonlake } |] = FeatureAVX512_Cannonlake+ | fromIntegral k == [CU.pure| int { halide_target_feature_avx512_sapphirerapids } |] = FeatureAVX512_SapphireRapids+ | fromIntegral k == [CU.pure| int { halide_target_feature_trace_loads } |] = FeatureTraceLoads+ | fromIntegral k == [CU.pure| int { halide_target_feature_trace_stores } |] = FeatureTraceStores+ | fromIntegral k == [CU.pure| int { halide_target_feature_trace_realizations } |] = FeatureTraceRealizations+ | fromIntegral k == [CU.pure| int { halide_target_feature_trace_pipeline } |] = FeatureTracePipeline+ | fromIntegral k == [CU.pure| int { halide_target_feature_d3d12compute } |] = FeatureD3D12Compute+ | fromIntegral k == [CU.pure| int { halide_target_feature_strict_float } |] = FeatureStrictFloat+ | fromIntegral k == [CU.pure| int { halide_target_feature_tsan } |] = FeatureTSAN+ | fromIntegral k == [CU.pure| int { halide_target_feature_asan } |] = FeatureASAN+ | fromIntegral k == [CU.pure| int { halide_target_feature_check_unsafe_promises } |] = FeatureCheckUnsafePromises+ | fromIntegral k == [CU.pure| int { halide_target_feature_embed_bitcode } |] = FeatureEmbedBitcode+ | fromIntegral k == [CU.pure| int { halide_target_feature_enable_llvm_loop_opt } |] = FeatureEnableLLVMLoopOpt+ | fromIntegral k == [CU.pure| int { halide_target_feature_wasm_simd128 } |] = FeatureWasmSimd128+ | fromIntegral k == [CU.pure| int { halide_target_feature_wasm_signext } |] = FeatureWasmSignExt+ | fromIntegral k == [CU.pure| int { halide_target_feature_wasm_sat_float_to_int } |] = FeatureWasmSatFloatToInt+ | fromIntegral k == [CU.pure| int { halide_target_feature_wasm_threads } |] = FeatureWasmThreads+ | fromIntegral k == [CU.pure| int { halide_target_feature_wasm_bulk_memory } |] = FeatureWasmBulkMemory+ | fromIntegral k == [CU.pure| int { halide_target_feature_sve } |] = FeatureSVE+ | fromIntegral k == [CU.pure| int { halide_target_feature_sve2 } |] = FeatureSVE2+ | fromIntegral k == [CU.pure| int { halide_target_feature_arm_dot_prod } |] = FeatureARMDotProd+ | fromIntegral k == [CU.pure| int { halide_target_feature_arm_fp16 } |] = FeatureARMFp16+ | fromIntegral k == [CU.pure| int { halide_target_feature_rvv } |] = FeatureRVV+ | fromIntegral k == [CU.pure| int { halide_target_feature_armv81a } |] = FeatureARMv81a+ | fromIntegral k == [CU.pure| int { halide_target_feature_sanitizer_coverage } |] = FeatureSanitizerCoverage+ | fromIntegral k == [CU.pure| int { halide_target_feature_profile_by_timer } |] = FeatureProfileByTimer+ | fromIntegral k == [CU.pure| int { halide_target_feature_spirv } |] = FeatureSPIRV+ | fromIntegral k == [CU.pure| int { halide_target_feature_semihosting } |] = FeatureSemihosting+ | otherwise = error $ "unknown Target feature: " <> show k
+ src/Language/Halide/Trace.hs view
@@ -0,0 +1,181 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeApplications #-}++-- |+-- Module : Language.Halide.Trace+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Trace+ ( TraceEvent (..)+ , TraceEventCode (..)+ , TraceLoadStoreContents (..)+ , setCustomTrace+ , traceStores+ , traceLoads+ , collectIterationOrder+ )+where++import Control.Concurrent.MVar+import Control.Exception (bracket, bracket_)+import Data.ByteString (packCString)+import Data.Int (Int32)+import Data.Text (Text)+import Data.Text.Encoding (decodeUtf8)+import Foreign.Marshal (peekArray)+import Foreign.Ptr (FunPtr, Ptr, freeHaskellFunPtr)+import Foreign.Storable+import GHC.TypeLits+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Buffer+import Language.Halide.Context+import Language.Halide.Dimension+import Language.Halide.Func+import Language.Halide.LoopLevel+import Language.Halide.Type+import Prelude hiding (min, tail)++-- | Haskell counterpart of [@halide_trace_event_code_t@](https://halide-lang.org/docs/_halide_runtime_8h.html#a485130f12eb8bb5fa5a9478eeb6b0dfa).+data TraceEventCode+ = TraceLoad+ | TraceStore+ | TraceBeginRealization+ | TraceEndRealization+ | TraceProduce+ | TraceEndProduce+ | TraceConsume+ | TraceEndConsume+ | TraceBeginPipeline+ | TraceEndPipeline+ | TraceTag+ deriving stock (Show, Eq, Ord)++data TraceLoadStoreContents = TraceLoadStoreContents+ { valuePtr :: !(Ptr ())+ , valueType :: !HalideType+ , coordinates :: ![Int]+ }+ deriving stock (Show)++data TraceEvent = TraceEvent+ { funcName :: !Text+ , eventCode :: !TraceEventCode+ , loadStoreContents :: !(Maybe TraceLoadStoreContents)+ }+ deriving stock (Show)++importHalide++instance Enum TraceEventCode where+ fromEnum =+ fromIntegral . \case+ TraceLoad -> [CU.pure| int { halide_trace_load } |]+ TraceStore -> [CU.pure| int { halide_trace_store } |]+ TraceBeginRealization -> [CU.pure| int { halide_trace_begin_realization } |]+ TraceEndRealization -> [CU.pure| int { halide_trace_end_realization } |]+ TraceProduce -> [CU.pure| int { halide_trace_produce } |]+ TraceEndProduce -> [CU.pure| int { halide_trace_end_produce } |]+ TraceConsume -> [CU.pure| int { halide_trace_consume } |]+ TraceEndConsume -> [CU.pure| int { halide_trace_end_consume } |]+ TraceBeginPipeline -> [CU.pure| int { halide_trace_begin_pipeline } |]+ TraceEndPipeline -> [CU.pure| int { halide_trace_end_pipeline } |]+ TraceTag -> [CU.pure| int { halide_trace_tag } |]+ toEnum k+ | fromIntegral k == [CU.pure| int { halide_trace_load } |] = TraceLoad+ | fromIntegral k == [CU.pure| int { halide_trace_store } |] = TraceStore+ | fromIntegral k == [CU.pure| int { halide_trace_begin_realization } |] = TraceBeginRealization+ | fromIntegral k == [CU.pure| int { halide_trace_end_realization } |] = TraceEndRealization+ | fromIntegral k == [CU.pure| int { halide_trace_produce } |] = TraceProduce+ | fromIntegral k == [CU.pure| int { halide_trace_end_produce } |] = TraceEndProduce+ | fromIntegral k == [CU.pure| int { halide_trace_consume } |] = TraceConsume+ | fromIntegral k == [CU.pure| int { halide_trace_end_consume } |] = TraceEndConsume+ | fromIntegral k == [CU.pure| int { halide_trace_begin_pipeline } |] = TraceBeginPipeline+ | fromIntegral k == [CU.pure| int { halide_trace_end_pipeline } |] = TraceEndPipeline+ | fromIntegral k == [CU.pure| int { halide_trace_tag } |] = TraceTag+ | otherwise = error $ "invalid TraceEventCode: " <> show k++peekTraceLoadStoreContents :: Ptr TraceEvent -> IO TraceLoadStoreContents+peekTraceLoadStoreContents p = do+ v <- [CU.exp| void* { $(const halide_trace_event_t* p)->value } |]+ tp <- peek =<< [CU.exp| const halide_type_t* { &$(const halide_trace_event_t* p)->type } |]+ n <- fromIntegral <$> [CU.exp| int { $(const halide_trace_event_t* p)->dimensions } |]+ cs <- peekArray n =<< [CU.exp| const int32_t* { $(const halide_trace_event_t* p)->coordinates } |]+ pure $ TraceLoadStoreContents v tp (fromIntegral <$> cs)++peekTraceEvent :: Ptr TraceEvent -> IO TraceEvent+peekTraceEvent p = do+ f <-+ fmap decodeUtf8 $+ packCString+ =<< [CU.exp| const char* { $(const halide_trace_event_t* p)->func } |]+ c <- toEnum . fromIntegral <$> [CU.exp| int { $(const halide_trace_event_t* p)->event } |]+ contents <-+ case c of+ TraceLoad -> Just <$> peekTraceLoadStoreContents p+ TraceStore -> Just <$> peekTraceLoadStoreContents p+ _ -> pure Nothing+ pure $ TraceEvent f c contents++withTrace+ :: (TraceEvent -> IO ()) -> (FunPtr (Ptr CxxUserContext -> Ptr TraceEvent -> IO Int32) -> IO a) -> IO a+withTrace customTrace = bracket allocate destroy+ where+ allocate = do+ $(C.mkFunPtr [t|Ptr CxxUserContext -> Ptr TraceEvent -> IO Int32|]) $ \_ p ->+ peekTraceEvent p >>= customTrace >> pure 0+ destroy = freeHaskellFunPtr++setCustomTrace+ :: (KnownNat n, IsHalideType a)+ => (TraceEvent -> IO ()) -- ^ Custom trace function+ -> Func t n a -- ^ For which func to enable it+ -> IO b -- ^ For the duration of which computation to enable it+ -> IO b+setCustomTrace customTrace f action =+ withTrace customTrace $ \tracePtr ->+ bracket_ (set tracePtr) unset action+ where+ set tracePtr =+ withFunc f $ \f' ->+ [CU.block| void {+ auto& func = *$(Halide::Func* f');+ func.jit_handlers().custom_trace = $(int32_t (*tracePtr)(Halide::JITUserContext*, const halide_trace_event_t*));+ } |]+ unset =+ withFunc f $ \f' ->+ [CU.block| void {+ auto& func = *$(Halide::Func* f');+ func.jit_handlers().custom_trace = nullptr;+ } |]++traceStores :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+traceStores f = do+ withFunc f $ \f' ->+ [CU.exp| void { $(Halide::Func* f')->trace_stores() } |]+ pure f++traceLoads :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+traceLoads f = do+ withFunc f $ \f' ->+ [CU.exp| void { $(Halide::Func* f')->trace_loads() } |]+ pure f++collectIterationOrder+ :: (KnownNat n, IsHalideType a)+ => (TraceEventCode -> Bool)+ -> Func t n a+ -> IO b+ -> IO ([[Int]], b)+collectIterationOrder cond f action = do+ m <- newMVar []+ let tracer (TraceEvent _ c' (Just payload))+ | cond c' = modifyMVar_ m $ pure . (payload.coordinates :)+ tracer _ = pure ()+ setCustomTrace tracer f $ do+ r <- action+ cs <- readMVar m+ pure (reverse cs, r)
+ src/Language/Halide/Type.hs view
@@ -0,0 +1,414 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DeriveLift #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-unused-local-binds -Wno-unused-matches #-}++-- |+-- Module : Language.Halide.Type+-- Description : Low-level types+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Type+ ( HalideTypeCode (..)+ , HalideType (..)+ , IsHalideType (..)+ , CxxExpr+ , CxxVar+ , CxxRVar+ , CxxVarOrRVar+ , CxxFunc+ , CxxParameter+ , CxxArgument+ , CxxImageParam+ , CxxVector+ , CxxUserContext+ , CxxCallable+ , CxxTarget+ , CxxStageSchedule+ , CxxString+ , Arguments (..)+ , Length+ , Append+ , argumentsAppend+ , FunctionArguments+ , FunctionReturn+ , All+ , UnCurry (..)+ , Curry (..)+ , defineIsHalideTypeInstances+ , instanceHasCxxVector+ , HasCxxVector (..)+ , IsTuple (..)+ , FromTuple+ , ToTuple+ , instanceCxxConstructible+ , CxxConstructible (..)+ -- defineCastableInstances,+ -- defineCurriedTypeFamily,+ -- defineUnCurriedTypeFamily,+ -- defineCurryInstances,+ -- defineUnCurryInstances,+ )+where++import Data.Coerce+import Data.Constraint+import Data.Int+import Data.Kind (Type)+import qualified Data.Text as T+import Data.Word+import Foreign.C.Types+import Foreign.ForeignPtr+import Foreign.Ptr+import Foreign.Storable+import GHC.TypeLits+import qualified Language.C.Inline as C+import qualified Language.C.Inline.Unsafe as CU+import qualified Language.Haskell.TH as TH+import Language.Haskell.TH.Syntax (Lift)++-- | Haskell counterpart of @Halide::Expr@.+data CxxExpr++-- | Haskell counterpart of @Halide::Var@.+data CxxVar++-- | Haskell counterpart of @Halide::RVar@.+data CxxRVar++-- | Haskell counterpart of @Halide::VarOrRVar@.+data CxxVarOrRVar++-- | Haskell counterpart of @Halide::Internal::Parameter@.+data CxxParameter++-- | Haskell counterpart of @Halide::Argument@.+data CxxArgument++-- | Haskell counterpart of @Halide::ImageParam@.+data CxxImageParam++-- | Haskell counterpart of @Halide::Func@.+data CxxFunc++-- | Haskell counterpart of @Halide::JITUserContext@.+data CxxUserContext++-- | Haskell counterpart of @Halide::Callable@.+data CxxCallable++-- | Haskell counterpart of @Halide::Target@.+data CxxTarget++-- | Haskell counterpart of @std::vector@.+data CxxVector a++-- | Haskell counterpart of @Halide::Internal::StageSchedule@.+data CxxStageSchedule++-- | Haskell counterpart of @std::string@+data CxxString++class CxxConstructible a where+ cxxSizeOf :: Int+ cxxConstruct :: (Ptr a -> IO ()) -> IO (ForeignPtr a)++cxxConstructWithDeleter :: Int -> FinalizerPtr a -> (Ptr a -> IO ()) -> IO (ForeignPtr a)+cxxConstructWithDeleter size deleter constructor = do+ fp <- mallocForeignPtrBytes size+ withForeignPtr fp constructor+ addForeignPtrFinalizer deleter fp+ pure fp++-- data Split =+-- SplitVar !Text !Text !Text !(Expr Int32) !++-- | Haskell counterpart of @halide_type_code_t@.+data HalideTypeCode+ = HalideTypeInt+ | HalideTypeUInt+ | HalideTypeFloat+ | HalideTypeHandle+ | HalideTypeBfloat+ deriving stock (Read, Show, Eq, Lift)++instance Enum HalideTypeCode where+ fromEnum :: HalideTypeCode -> Int+ fromEnum x = case x of+ HalideTypeInt -> 0+ HalideTypeUInt -> 1+ HalideTypeFloat -> 2+ HalideTypeHandle -> 3+ HalideTypeBfloat -> 4+ toEnum :: Int -> HalideTypeCode+ toEnum x = case x of+ 0 -> HalideTypeInt+ 1 -> HalideTypeUInt+ 2 -> HalideTypeFloat+ 3 -> HalideTypeHandle+ 4 -> HalideTypeBfloat+ _ -> error $ "invalid HalideTypeCode: " <> show x++-- | Haskell counterpart of @halide_type_t@.+data HalideType = HalideType+ { halideTypeCode :: !HalideTypeCode+ , halideTypeBits :: {-# UNPACK #-} !Word8+ , halideTypeLanes :: {-# UNPACK #-} !Word16+ }+ deriving stock (Read, Show, Eq)++instance Storable HalideType where+ sizeOf :: HalideType -> Int+ sizeOf _ = 4+ alignment :: HalideType -> Int+ alignment _ = 4+ peek :: Ptr HalideType -> IO HalideType+ peek p =+ HalideType+ <$> (toEnum . (fromIntegral :: Word8 -> Int) <$> peekByteOff p 0)+ <*> peekByteOff p 1+ <*> peekByteOff p 2+ poke :: Ptr HalideType -> HalideType -> IO ()+ poke p (HalideType code bits lanes) = do+ pokeByteOff p 0 . (fromIntegral :: Int -> Word8) . fromEnum $ code+ pokeByteOff p 1 bits+ pokeByteOff p 2 lanes++-- | Specifies that a type is supported by Halide.+class Storable a => IsHalideType a where+ halideTypeFor :: proxy a -> HalideType+ toCxxExpr :: a -> IO (ForeignPtr CxxExpr)++-- | Helper function to coerce 'Float' to 'CFloat' and 'Double' to 'CDouble'+-- before passing them to inline-c quasiquotes. This is needed because inline-c+-- assumes that @float@ in C corresponds to 'CFloat' in Haskell.+optionallyCast :: String -> TH.TypeQ -> TH.ExpQ+optionallyCast cType hsType' = do+ hsType <- hsType'+ hsTargetType <- C.getHaskellType False cType+ if hsType == hsTargetType then [e|id|] else [e|coerce|]++-- | Template Haskell splice that defines instances of 'IsHalideType' for a+-- given Haskell type.+instanceIsHalideType :: (String, TH.TypeQ, HalideTypeCode) -> TH.DecsQ+instanceIsHalideType (cType, hsType, typeCode) =+ C.substitute+ [("T", \x -> "$(" <> cType <> " " <> x <> ")")]+ [d|+ instance IsHalideType $hsType where+ halideTypeFor _ = HalideType typeCode bits 1+ where+ bits = fromIntegral $ 8 * sizeOf (undefined :: $hsType)+ toCxxExpr y =+ cxxConstruct $ \ptr ->+ [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{@T(x)} } |]+ where+ x = $(optionallyCast cType hsType) y+ |]++-- | Derive 'IsHalideType' instances for all supported types.+defineIsHalideTypeInstances :: TH.DecsQ+defineIsHalideTypeInstances = concat <$> mapM instanceIsHalideType halideTypes++instanceCxxConstructible :: String -> TH.DecsQ+instanceCxxConstructible cType =+ C.substitute+ [ ("T", const cType)+ , ("Deleter", const $ "deleter(" <> cType <> "* p)")+ , ("Class", const . T.unpack . snd $ T.breakOnEnd "::" (T.pack cType))+ ]+ [d|+ instance CxxConstructible $(C.getHaskellType False cType) where+ cxxSizeOf = fromIntegral [CU.pure| size_t { sizeof(@T()) } |]+ cxxConstruct = cxxConstructWithDeleter size deleter+ where+ size = fromIntegral [CU.pure| size_t { sizeof(@T()) } |]+ deleter = [C.funPtr| void @Deleter() { p->~@Class()(); } |]+ |]++-- | Specifies that a given Haskell type can be used with @std::vector@.+--+-- E.g. if we have @HasCxxVector Int16@, then using @std::vector<int16_t>*@+-- in inline-c quotes will work.+class HasCxxVector a where+ newCxxVector :: Maybe Int -> IO (Ptr (CxxVector a))+ deleteCxxVector :: Ptr (CxxVector a) -> IO ()+ cxxVectorSize :: Ptr (CxxVector a) -> IO Int+ cxxVectorPushBack :: Ptr (CxxVector a) -> Ptr a -> IO ()+ cxxVectorData :: Ptr (CxxVector a) -> IO (Ptr a)+ peekCxxVector :: Storable a => Ptr (CxxVector a) -> IO [a]++-- | Template Haskell splice that defines an instance of 'HasCxxVector' for a given C type name.+instanceHasCxxVector :: String -> TH.DecsQ+instanceHasCxxVector cType =+ C.substitute+ [ ("T", const cType)+ , ("VEC", \var -> "$(std::vector<" ++ cType ++ ">* " ++ var ++ ")")+ ]+ [d|+ instance HasCxxVector $(C.getHaskellType False cType) where+ newCxxVector maybeSize = do+ v <- [CU.exp| std::vector<@T()>* { new std::vector<@T()>() } |]+ case maybeSize of+ Just size ->+ let n = fromIntegral size+ in [CU.exp| void { @VEC(v)->reserve($(size_t n)) } |]+ Nothing -> pure ()+ pure v+ deleteCxxVector vec = [CU.exp| void { delete @VEC(vec) } |]+ cxxVectorSize vec = fromIntegral <$> [CU.exp| size_t { @VEC(vec)->size() } |]+ cxxVectorPushBack vec x = [CU.exp| void { @VEC(vec)->push_back(*$(@T()* x)) } |]+ cxxVectorData vec = [CU.exp| @T()* { @VEC(vec)->data() } |]+ peekCxxVector vec = do+ n <- cxxVectorSize vec+ allocaArray n $ \out -> do+ [CU.block| void {+ auto const& vec = *@VEC(vec);+ auto* out = $(@T()* out);+ std::uninitialized_copy(std::begin(vec), std::end(vec), out);+ } |]+ peekArray n out+ |]++-- | List of all supported types.+halideTypes :: [(String, TH.TypeQ, HalideTypeCode)]+halideTypes =+ [ ("float", [t|Float|], HalideTypeFloat)+ , ("float", [t|CFloat|], HalideTypeFloat)+ , ("double", [t|Double|], HalideTypeFloat)+ , ("double", [t|CDouble|], HalideTypeFloat)+ , ("int8_t", [t|Int8|], HalideTypeInt)+ , ("int16_t", [t|Int16|], HalideTypeInt)+ , ("int32_t", [t|Int32|], HalideTypeInt)+ , ("int64_t", [t|Int64|], HalideTypeInt)+ , ("uint8_t", [t|Word8|], HalideTypeUInt)+ , ("uint16_t", [t|Word16|], HalideTypeUInt)+ , ("uint32_t", [t|Word32|], HalideTypeUInt)+ , ("uint64_t", [t|Word64|], HalideTypeUInt)+ ]++infixr 5 :::++-- | A heterogeneous list.+data Arguments (k :: [Type]) where+ Nil :: Arguments '[]+ (:::) :: !t -> !(Arguments ts) -> Arguments (t ': ts)++-- | A type family that returns the length of a type-level list.+type family Length (xs :: [k]) :: Nat where+ Length '[] = 0+ Length (x ': xs) = 1 + Length xs++-- | Append to a type-level list.+type family Append (xs :: [k]) (y :: k) :: [k] where+ Append '[] y = '[y]+ Append (x ': xs) y = x ': Append xs y++-- | Append a value to 'Arguments'+argumentsAppend :: Arguments xs -> t -> Arguments (Append xs t)+argumentsAppend = go+ where+ go :: forall xs t. Arguments xs -> t -> Arguments (Append xs t)+ go Nil y = y ::: Nil+ go (x ::: xs) y = x ::: go xs y++-- | Return the list of arguments to of a function type.+type family FunctionArguments (f :: Type) :: [Type] where+ FunctionArguments (a -> b) = a ': FunctionArguments b+ FunctionArguments a = '[]++-- | Get the return type of a function.+type family FunctionReturn (f :: Type) :: Type where+ FunctionReturn (a -> b) = FunctionReturn b+ FunctionReturn a = a++-- | Apply constraint to all types in a list.+type family All (c :: Type -> Constraint) (ts :: [Type]) :: Constraint where+ All c '[] = ()+ All c (t ': ts) = (c t, All c ts)++-- | A helper typeclass to convert a normal curried function to a function that+-- takes 'Arguments' as input.+--+-- For instance, if we have a function @f :: Int -> Float -> Double@, then it+-- will be converted to @f' :: Arguments '[Int, Float] -> Double@.+class UnCurry (f :: Type) (args :: [Type]) (r :: Type) | args r -> f where+ uncurryG :: f -> Arguments args -> r++instance (FunctionArguments f ~ '[], FunctionReturn f ~ r, f ~ r) => UnCurry f '[] r where+ uncurryG f Nil = f+ {-# INLINE uncurryG #-}++instance (UnCurry f args r) => UnCurry (a -> f) (a ': args) r where+ uncurryG f (a ::: args) = uncurryG (f a) args+ {-# INLINE uncurryG #-}++-- | A helper typeclass to convert a function that takes 'Arguments' as input+-- into a normal curried function. This is the inverse of 'UnCurry'.+--+-- For instance, if we have a function @f :: Arguments '[Int, Float] -> Double@, then+-- it will be converted to @f' :: Int -> Float -> Double@.+class Curry (args :: [Type]) (r :: Type) (f :: Type) | args r -> f where+ curryG :: (Arguments args -> r) -> f++instance Curry '[] r r where+ curryG f = f Nil+ {-# INLINE curryG #-}++instance Curry args r f => Curry (a ': args) r (a -> f) where+ curryG f a = curryG (\args -> f (a ::: args))++-- | Type family that maps @'Arguments' ts@ to the corresponding tuple type.+type family ToTuple t where+ ToTuple (Arguments '[]) = ()+ ToTuple (Arguments '[a1]) = a1+ ToTuple (Arguments '[a1, a2]) = (a1, a2)+ ToTuple (Arguments '[a1, a2, a3]) = (a1, a2, a3)+ ToTuple (Arguments '[a1, a2, a3, a4]) = (a1, a2, a3, a4)+ ToTuple (Arguments '[a1, a2, a3, a4, a5]) = (a1, a2, a3, a4, a5)++-- | Type family that maps tuples to the corresponding @'Arguments' ts@ type. This is essentially the inverse+-- of 'ToTuple'.+type family FromTuple t++type instance FromTuple () = Arguments '[]+type instance FromTuple (a1, a2) = Arguments '[a1, a2]+type instance FromTuple (a1, a2, a3) = Arguments '[a1, a2, a3]+type instance FromTuple (a1, a2, a3, a4) = Arguments '[a1, a2, a3, a4]+type instance FromTuple (a1, a2, a3, a4, a5) = Arguments '[a1, a2, a3, a4, a5]++-- | Specifies that there is an isomorphism between a type @a@ and a tuple @t@.+--+-- We use this class to convert between 'Arguments' and normal tuples.+class (ToTuple a ~ t, FromTuple t ~ a) => IsTuple a t | a -> t, t -> a where+ toTuple :: a -> t+ fromTuple :: t -> a++instance IsTuple (Arguments '[]) () where+ toTuple Nil = ()+ fromTuple () = Nil++instance IsTuple (Arguments '[a1, a2]) (a1, a2) where+ toTuple (a1 ::: a2 ::: Nil) = (a1, a2)+ fromTuple (a1, a2) = a1 ::: a2 ::: Nil++instance IsTuple (Arguments '[a1, a2, a3]) (a1, a2, a3) where+ toTuple (a1 ::: a2 ::: a3 ::: Nil) = (a1, a2, a3)+ fromTuple (a1, a2, a3) = a1 ::: a2 ::: a3 ::: Nil++instance IsTuple (Arguments '[a1, a2, a3, a4]) (a1, a2, a3, a4) where+ toTuple (a1 ::: a2 ::: a3 ::: a4 ::: Nil) = (a1, a2, a3, a4)+ fromTuple (a1, a2, a3, a4) = a1 ::: a2 ::: a3 ::: a4 ::: Nil++instance IsTuple (Arguments '[a1, a2, a3, a4, a5]) (a1, a2, a3, a4, a5) where+ toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: Nil) = (a1, a2, a3, a4, a5)+ fromTuple (a1, a2, a3, a4, a5) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: Nil++-- instance IsTuple (Arguments '[Expr a]) (Expr a) where+-- toTuple (x ::: Nil) = x+-- fromTuple () = Nil
+ src/Language/Halide/Utils.hs view
@@ -0,0 +1,37 @@+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}++-- |+-- Module : Language.Halide.Utils+-- Description : Utilities for writing FFI code+-- Copyright : (c) Tom Westerhout, 2023+module Language.Halide.Utils+ ( peekCxxString+ , peekAndDeleteCxxString+ ) where++import Data.ByteString (packCString)+import Data.Text (Text)+import qualified Data.Text.Encoding as T+import Foreign.Ptr (Ptr)+import qualified Language.C.Inline.Unsafe as CU+import Language.Halide.Context+import Language.Halide.Type++importHalide++-- | Convert a pointer to @std::string@ into a string.+--+-- It properly handles unicode characters.+peekCxxString :: Ptr CxxString -> IO Text+peekCxxString p =+ fmap T.decodeUtf8 $+ packCString+ =<< [CU.exp| char const* { $(const std::string* p)->c_str() } |]++-- | Call 'peekCxxString' and @delete@ the pointer.+peekAndDeleteCxxString :: Ptr CxxString -> IO Text+peekAndDeleteCxxString p = do+ s <- peekCxxString p+ [CU.exp| void { delete $(const std::string* p) } |]+ pure s
+ test/Language/Halide/BufferSpec.hs view
@@ -0,0 +1,54 @@+{-# LANGUAGE OverloadedRecordDot #-}++module Language.Halide.BufferSpec (spec) where++import Data.Int (Int64)+import Foreign.Ptr (Ptr, nullPtr)+import Language.Halide+import Test.Hspec+import Test.Hspec.QuickCheck+import Test.QuickCheck++newtype ListVector a = ListVector [a]+ deriving stock (Show)++newtype ListMatrix a = ListMatrix [[a]]+ deriving stock (Show)++newtype ListTensor3D a = ListTensor3D [[[a]]]+ deriving stock (Show)++instance Arbitrary a => Arbitrary (ListVector a) where+ arbitrary = ListVector <$> listOf arbitrary++instance Arbitrary a => Arbitrary (ListMatrix a) where+ arbitrary = do+ d0 <- chooseInt (0, 50)+ d1 <- chooseInt (0, 50)+ ListMatrix <$> vectorOf d0 (vector d1)++instance Arbitrary a => Arbitrary (ListTensor3D a) where+ arbitrary = do+ d0 <- chooseInt (0, 30)+ d1 <- chooseInt (0, 30)+ d2 <- chooseInt (0, 30)+ ListTensor3D <$> vectorOf d0 (vectorOf d1 (vector d2))++spec :: Spec+spec = do+ it "rowMajorStrides" $ do+ rowMajorStrides [1, 1, 1] `shouldBe` ([1, 1, 1] :: [Int])+ rowMajorStrides [2, 1, 3] `shouldBe` ([3, 3, 1] :: [Int])+ rowMajorStrides [3, 2] `shouldBe` ([2, 1] :: [Int])+ rowMajorStrides [] `shouldBe` ([] :: [Int])+ it "bufferFromPtrShapeStrides" $ do+ bufferFromPtrShapeStrides nullPtr [3, 2, 1] [1, 1, 1] (\(_ :: Ptr (HalideBuffer 2 Int32)) -> pure ())+ `shouldThrow` anyErrorCall+ bufferFromPtrShapeStrides nullPtr [3] [1] (\(_ :: Ptr (HalideBuffer 2 Int32)) -> pure ())+ `shouldThrow` anyErrorCall+ prop "works with [a]" $ \(ListVector xs :: ListVector Float) ->+ withHalideBuffer @1 @Float xs peekToList `shouldReturn` xs+ prop "works with [[a]]" $ \(ListMatrix xs :: ListMatrix Int64) ->+ withHalideBuffer @2 @Int64 xs peekToList `shouldReturn` xs+ prop "works with [[[a]]]" $ \(ListTensor3D xs :: ListTensor3D Double) ->+ withHalideBuffer @3 @Double xs peekToList `shouldReturn` xs
+ test/Language/Halide/ExprSpec.hs view
@@ -0,0 +1,112 @@+module Language.Halide.ExprSpec (spec) where++import Control.Monad (unless, when)+import Data.Int+import Data.Word+import Language.Halide+import Test.Hspec+import Test.Hspec.QuickCheck+import Test.QuickCheck (Property)+import Test.QuickCheck.Monadic (PropertyM, assert, monadicIO, run)+import Type.Reflection+import Utils++isOverflowing :: Typeable a => (Integer -> Integer -> Integer) -> a -> a -> Bool+isOverflowing op x y+ | Just HRefl <- eqTypeRep (typeOf x) (typeRep @Int32) =+ op (toInteger x) (toInteger y) > toInteger (maxBound @Int32)+ || op (toInteger x) (toInteger y) < toInteger (minBound @Int32)+ | Just HRefl <- eqTypeRep (typeOf x) (typeRep @Int64) =+ op (toInteger x) (toInteger y) > toInteger (maxBound @Int64)+ || op (toInteger x) (toInteger y) < toInteger (minBound @Int64)+ | otherwise = False++infix 1 `evaluatesTo`++evaluatesTo :: (Eq a, IsHalideType a) => Expr a -> a -> PropertyM IO ()+evaluatesTo expr expected =+ assert . (expected ==) =<< (run . evaluate) expr++infix 1 `evaluatesToApprox`++evaluatesToApprox :: (Ord a, IsHalideType a, HasEpsilon a) => Expr a -> a -> PropertyM IO ()+evaluatesToApprox expr expected =+ assert . approx' expected =<< (run . evaluate) expr++infix 1 `shouldEvaluateTo`++shouldEvaluateTo :: (Eq a, IsHalideType a, Show a) => Expr a -> a -> Expectation+shouldEvaluateTo expr expected =+ evaluate expr `shouldReturn` expected++spec :: Spec+spec = do+ describe "mkExpr" $ modifyMaxSuccess (const 10) $ do+ let p :: forall a. (IsHalideType a, Eq a) => a -> Property+ p x = monadicIO $ mkExpr x `evaluatesTo` x+ prop "Bool" $ p @Bool++ describe "Num Expr" $ modifyMaxSuccess (const 10) $ do+ let whenNotOverflowing op x y check+ | isOverflowing op x y = pure ()+ | otherwise = check+ p :: forall a. (IsHalideType a, Eq a, Num a, Typeable a) => a -> a -> Property+ p x y =+ monadicIO $ do+ whenNotOverflowing (+) x y $+ mkExpr x + mkExpr y `evaluatesTo` x + y+ whenNotOverflowing (-) x y $+ mkExpr x - mkExpr y `evaluatesTo` x - y+ whenNotOverflowing (*) x y $+ mkExpr x * mkExpr y `evaluatesTo` x * y+ -- Temporary disable: see https://github.com/halide/Halide/issues/7365+ when (x /= -128) $+ abs (mkExpr x) `evaluatesTo` abs x+ negate (mkExpr x) `evaluatesTo` negate x+ prop "Int8" $ p @Int8+ prop "Int16" $ p @Int16+ prop "Int32" $ p @Int32+ prop "Int64" $ p @Int64+ prop "Word8" $ p @Word8+ prop "Word16" $ p @Word16+ prop "Word32" $ p @Word32+ prop "Word64" $ p @Word64+ prop "Float" $ p @Float+ prop "Double" $ p @Double++ describe "Fractional Expr" $ modifyMaxSuccess (const 10) $ do+ let p :: forall a. (IsHalideType a, Eq a, Fractional a) => a -> a -> Property+ p x y =+ monadicIO $+ unless (x == 0 && y == 0) $+ mkExpr x / mkExpr y `evaluatesTo` x / y+ prop "Float" $ p @Float+ prop "Double" $ p @Double++ describe "Floating Expr" $ modifyMaxSuccess (const 10) $ do+ let p :: forall a. (IsHalideType a, Ord a, Floating a, HasEpsilon a) => a -> Property+ p x = monadicIO $ do+ when (x > 0) $ do+ log (mkExpr x) `evaluatesToApprox` log x+ sqrt (mkExpr x) `evaluatesToApprox` sqrt x+ exp (mkExpr x) `evaluatesToApprox` exp x+ sin (mkExpr x) `evaluatesToApprox` sin x+ cos (mkExpr x) `evaluatesToApprox` cos x+ tan (mkExpr x) `evaluatesToApprox` tan x+ when (-1 <= x && x <= 1) $ do+ asin (mkExpr x) `evaluatesToApprox` asin x+ acos (mkExpr x) `evaluatesToApprox` acos x+ atan (mkExpr x) `evaluatesToApprox` atan x+ sinh (mkExpr x) `evaluatesToApprox` sinh x+ cosh (mkExpr x) `evaluatesToApprox` cosh x+ tanh (mkExpr x) `evaluatesToApprox` tanh x+ asinh (mkExpr x) `evaluatesToApprox` asinh x+ when (x >= 1) $+ acosh (mkExpr x) `evaluatesToApprox` acosh x+ when (-1 <= x && x <= 1) $+ atanh (mkExpr x) `evaluatesToApprox` atanh x+ prop "Float" $ p @Float+ prop "Double" $ p @Double+ it "defines pi" $ do+ (pi :: Expr Float) `shouldEvaluateTo` pi+ (pi :: Expr Double) `shouldEvaluateTo` pi
+ test/Language/Halide/FuncSpec.hs view
@@ -0,0 +1,252 @@+{-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE ViewPatterns #-}++module Language.Halide.FuncSpec (spec) where++import Control.Monad.ST (RealWorld)+import qualified Data.Vector.Storable.Mutable as SM+import Language.Halide+import Test.Hspec hiding (parallel)+import Utils++importHalide++data Matrix v a = Matrix+ { matrixRows :: !Int+ , matrixCols :: !Int+ , matrixData :: !(v a)+ }+ deriving stock (Show, Eq)++instance IsHalideType a => IsHalideBuffer (Matrix (SM.MVector RealWorld) a) 2 a where+ withHalideBufferImpl (Matrix n m v) f =+ SM.unsafeWith v $ \dataPtr ->+ bufferFromPtrShapeStrides dataPtr [n, m] [1, n] f++spec :: Spec+spec = do+ describe "indexing" $ do+ it "supports empty tuples" $ do+ let x = mkExpr (5 :: Double)+ f <- define "f" () $ x * x - 2 * x + 5 + 3 / x+ g <- define "g" () $ f ! ()+ realize g [] peekToList `shouldReturn` [20.6]++ describe "vectorize" $ do+ it "vectorizes loops" $ do+ i <- mkVar "i"+ ii <- mkVar "inner"+ func <- define "func" i $ (3 * i + 1) * (i - 5)+ -- by default, nothing is vectorized+ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "vectorized"+ s `shouldNotContainText` "[0, 3]"++ void $+ split TailShiftInwards i (i, ii) 4 func+ >>= vectorize ii++ -- now, the inner loop is vectorized+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "vectorized i.inner"+ s `shouldContainText` "in [0, 3]"++ let n = 10+ realize func [n] peekToList+ `shouldReturn` [(3 * k + 1) * (k - 5) | k <- [0 .. fromIntegral n - 1]]++ describe "unroll" $ do+ it "unrolls loops" $ do+ i <- mkVar "i"+ ii <- mkVar "inner"+ func <- define "func" i $ (3 * i + 1) * (i - 5)+ -- by default, nothing is unrolled+ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "unrolled"+ s `shouldNotContainText` "[0, 2]"++ void $+ split TailGuardWithIf i (i, ii) 3 func+ >>= unroll ii++ -- now, the inner loop is unrolled+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "unrolled i.inner"+ s `shouldContainText` "in [0, 2]"++ let n = 17+ realize func [n] peekToList+ `shouldReturn` [(3 * k + 1) * (k - 5) | k <- [0 .. fromIntegral n - 1]]++ describe "reorder" $ do+ it "reorders loops" $ do+ x <- mkVar "x"+ y <- mkVar "y"+ z <- mkVar "z"+ func <- define "func" (x, y, z) $ x * (x + y) - 3 * z++ -- we have+ --+ -- for z+ -- for y+ -- for x+ prettyLoopNest func >>= \s -> do+ s & "for z" `appearsBeforeText` "for y"+ s & "for y" `appearsBeforeText` "for x"++ void $ reorder [z, x, y] func++ -- now we expect+ --+ -- for y+ -- for x+ -- for z+ prettyLoopNest func >>= \s -> do+ s & "for y" `appearsBeforeText` "for x"+ s & "for x" `appearsBeforeText` "for z"++ describe "split" $ do+ it "splits loops into sub-loops" $ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ -- we have+ --+ -- for y+ -- for x+ prettyLoopNest func >>= \s -> do+ s & "for y" `appearsBeforeText` "for x"+ s `shouldNotContainText` "outer"+ s `shouldNotContainText` "inner"++ outer <- mkVar "outer"+ inner <- mkVar "inner"+ void $ split TailGuardWithIf x (outer, inner) 7 func++ -- now we expect+ --+ -- for y+ -- for x.outer+ -- for x.inner+ prettyLoopNest func >>= \s -> do+ s & "for y" `appearsBeforeText` "for x.outer"+ s & "for x.outer" `appearsBeforeText` "for x.inner"++ describe "fuse" $ do+ it "merges sub-loops into one" $ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ -- we have+ --+ -- for y+ -- for x+ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "common"++ common <- mkVar "common"+ void $ fuse (x, y) common func++ -- now we expect+ --+ -- for common+ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "for x:"+ s `shouldNotContainText` "for y"+ s `shouldContainText` "for x.common"++ describe "parallel" $ do+ it "marks dimensions as parallel" $ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ prettyLoopNest func >>= \s ->+ s `shouldNotContainText` "parallel"++ void $+ parallel x func+ >>= serial y++ prettyLoopNest func >>= \s ->+ s `shouldContainText` "parallel x"++ describe "gpuBlocks" $ do+ it "marks dimensions as corresponding to GPU blocks" $ do+ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "gpu_block"+ s `shouldNotContainText` "Default_GPU"+ void $ gpuBlocks DeviceDefaultGPU (x, y) func+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "gpu_block y<Default_GPU>"+ s `shouldContainText` "gpu_block x<Default_GPU>"++ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "gpu_block"+ s `shouldNotContainText` "CUDA"+ void $ gpuBlocks DeviceCUDA y func+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "gpu_block y<CUDA>"++ describe "gpuThreads" $ do+ it "marks dimensions as corresponding to GPU threads" $ do+ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "gpu_thread"+ s `shouldNotContainText` "Default_GPU"+ void $ gpuThreads DeviceDefaultGPU (x, y) func+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "gpu_thread y<Default_GPU>"+ s `shouldContainText` "gpu_thread x<Default_GPU>"++ do+ x <- mkVar "x"+ y <- mkVar "y"+ func <- define "func" (x, y) $ x * y++ prettyLoopNest func >>= \s -> do+ s `shouldNotContainText` "gpu_block"+ s `shouldNotContainText` "gpu_thread"+ s `shouldNotContainText` "CUDA"+ void $+ gpuBlocks DeviceCUDA y func+ >>= gpuThreads DeviceCUDA x+ prettyLoopNest func >>= \s -> do+ s `shouldContainText` "gpu_block y<CUDA>"+ s `shouldContainText` "gpu_thread x<CUDA>"++ describe "reductions" $ do+ it "computes reductions" $ do+ asBufferParam @1 @Int32 ([1, 2, 3, 4, 5] :: [Int32]) $ \src -> do+ n <- (.extent) <$> dim 0 src+ r <- mkRVar "r" 0 n+ i <- mkVar "i"+ f <- define "sum" i 0+ update f (0 :: Expr Int32) $ f ! (0 :: Expr Int32) + src ! r+ realize f [1] peekToList `shouldReturn` ([15] :: [Int32])++ describe "undef" $ do+ it "allows to skip stores" $ do+ i <- mkVar "i"+ f <- define "f" i $ bool (i `gt` 5) i 0+ update f i $ bool ((f ! i) `eq` 0) (2 * i) undef+ realize f [10] peekToList `shouldReturn` ([0, 2, 4, 6, 8, 10] <> [6 .. 9] :: [Int32])
+ test/Language/Halide/KernelSpec.hs view
@@ -0,0 +1,66 @@+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE ViewPatterns #-}++module Language.Halide.KernelSpec (spec) where++import Language.Halide+import Test.Hspec+import Utils++spec :: Spec+spec = do+ describe "compile" $ do+ it "compiles a kernel that adds two vectors together" $ do+ vectorPlus <- compile $ \a b -> do+ i <- mkVar "i"+ define "out" i $ a ! i + b ! i+ let n = 10+ a = replicate 10 (1 :: Float)+ b = replicate 10 (2 :: Float)+ withHalideBuffer a $ \a' ->+ withHalideBuffer b $ \b' ->+ allocaCpuBuffer [n] $ \out' -> do+ vectorPlus a' b' out'+ peekToList out' `shouldReturn` zipWith (+) a b++ it "compiles a kernel that generates a scaled diagonal matrix declaratively" $ do+ scaledDiagonal <- compile $ \(scale :: Expr Double) v -> do+ i <- mkVar "i"+ j <- mkVar "j"+ define "out" (i, j) $+ bool+ (i `eq` j)+ (v ! i / scale)+ 0+ let a :: [Double]+ a = [1.0, 2.0, 3.0]+ withHalideBuffer a $ \a' ->+ allocaCpuBuffer [3, 3] $ \out' -> do+ scaledDiagonal 2 a' out'+ peekToList out' `shouldReturn` [[0.5, 0, 0], [0, 1, 0], [0, 0, 1.5]]++ it "compiles a kernel that generates a scaled diagonal matrix statefully" $ do+ scaledDiagonal <- compile $ \(scale :: Expr Double) v -> do+ i <- mkVar "i"+ j <- mkVar "j"+ out <- define "out" (i, j) 0+ update out (i, i) (v ! i / scale)+ pure out+ let a :: [Double]+ a = [1.0, 2.0, 3.0]+ withHalideBuffer a $ \a' ->+ allocaCpuBuffer [3, 3] $ \out' -> do+ scaledDiagonal 2 a' out'+ peekToList out' `shouldReturn` [[0.5, 0, 0], [0, 1, 0], [0, 0, 1.5]]++ describe "compileToLoweredStmt" $ do+ it "compiles to lowered stmt file" $ do+ let builder (buffer "src" -> src) (scalar @Float "c" -> c) = do+ i <- mkVar "i"+ define "dest1234" i $ c * src ! i+ target =+ setFeature FeatureNoAsserts . setFeature FeatureNoBoundsQuery $+ hostTarget+ s <- compileToLoweredStmt StmtText target builder+ s `shouldContainText` "func dest1234 (src, c, dest1234) {"+ s `shouldContainText` "produce dest1234 {"
+ test/Language/Halide/LoopLevelSpec.hs view
@@ -0,0 +1,120 @@+-- {-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE ViewPatterns #-}++module Language.Halide.LoopLevelSpec (spec) where++-- import Control.Exception (catch)+-- import Control.Monad (void)+-- import Control.Monad.ST (RealWorld)+-- import Data.Function ((&))+-- import Data.Int+-- import Data.Text (Text)+-- import qualified Data.Text as T+-- import qualified Data.Text.Encoding as T+-- import qualified Data.Text.IO as T+-- import qualified Data.Vector.Storable as S+-- import qualified Data.Vector.Storable.Mutable as SM+-- import qualified Language.C.Inline.Cpp.Exception as C+-- import qualified Language.C.Inline.Unsafe as CU+-- import Language.Halide.Context+-- import Language.Halide.LoopLevel+import Language.Halide+import Test.Hspec++importHalide++spec :: Spec+spec = do+ pure ()++-- describe "computeAt" $ do+-- it "schedules the computation to happen at a particular loop level" $ do+-- let innerLoop = do+-- x <- mkVar "x"+-- y <- mkVar "y"+-- g <- define "g" (x, y) $ x * y+-- -- f <- define "f" (x, y) $ g ! (x, y) + g ! (x, y + 1) + g ! (x + 1, y) + g ! (x + 1, y + 1)+-- f <-+-- define "f" (x, y) $+-- sum $+-- (g !) <$> [(x, y), (x, y + 1), (x + 1, y), (x + 1, y + 1)]+-- -- T.putStrLn =<< prettyLoopNest f+-- computeAt g =<< getLoopLevel f x+-- s <- prettyLoopNest f+-- s `shouldContainText` "produce g"+-- s `shouldContainText` "consume g"+-- -- T.putStrLn s+-- -- Both loops should appear before the produce statement+-- s & "for y" `appearsBeforeText` "produce g"+-- s & "for x" `appearsBeforeText` "produce g"+-- outerLoop = do+-- x <- mkVar "x"+-- y <- mkVar "y"+-- g <- define "g" (x, y) $ x * y+-- -- f <- define "f" (x, y) $ g ! (x, y) + g ! (x, y + 1) + g ! (x + 1, y) + g ! (x + 1, y + 1)+-- f <-+-- define "f" (x, y) $+-- sum $+-- (g !) <$> [(x, y), (x, y + 1), (x + 1, y), (x + 1, y + 1)]+-- computeAt g =<< getLoopLevel f y+-- s <- prettyLoopNest f+-- -- The produce statement should appear between for y and for x+-- s & "for y" `appearsBeforeText` "produce g"+-- s & "produce g" `appearsBeforeText` "for x"+-- innerLoop+-- outerLoop++-- describe "computeWith" $ do+-- it "schedules outer loops to be fused with another computation" $ do+-- x <- mkVar "x"+-- y <- mkVar "y"+-- f <- define "f" (x, y) $ x + y+-- g <- define "g" (x, y) $ x - y+-- h <- define "h" (x, y) $ f ! (x, y) + g ! (x, y)+-- computeRoot f+-- computeRoot g+-- xi <- mkVar "xi"+-- xo <- mkVar "xo"+-- split TailAuto f x xo xi 8+-- split TailAuto g x xo xi 8++-- prettyLoopNest h >>= \s -> do+-- s `shouldContainText` "for x.xo"+-- s `shouldContainText` "for x.xi"+-- s `shouldNotContainText` "fused"++-- computeWith LoopAlignAuto g =<< getLoopLevelAtStage f xo 0++-- prettyLoopNest h >>= \s -> do+-- s `shouldContainText` "for x.xi"+-- s `shouldContainText` "for fused.y"+-- s `shouldContainText` "for x.fused.xo"++-- describe "storeAt" $ do+-- it "allocates storage at a particular loop level" $ do+-- -- [C.throwBlock| void {+-- -- using namespace Halide;+-- -- Func f, g;+-- -- Var x, y;+-- -- g(x, y) = x*y;+-- -- f(x, y) = g(x, y) + g(x, y+1) + g(x+1, y) + g(x+1, y+1);+-- -- g.compute_at(f, x);++-- -- f.print_loop_nest();+-- -- } |]++-- x <- mkVar "x"+-- y <- mkVar "y"+-- g <- define "g" (x, y) $ x * y+-- f <- define "f" (x, y) $ g ! (x, y) + g ! (x, y + 1) + g ! (x + 1, y) + g ! (x + 1, y + 1)+-- computeAt g =<< getLoopLevel f x+-- T.putStrLn =<< prettyLoopNest f+-- storeAt g =<< getLoopLevel f y+-- T.putStrLn =<< prettyLoopNest f+-- s <- prettyLoopNest f+-- (pure (<) <*> (startIdxOf s "for y") <*> (startIdxOf s "store g"))+-- `shouldBe` Just True+-- (pure (>) <*> (startIdxOf s "for x") <*> (startIdxOf s "store g"))+-- `shouldBe` Just True
+ test/Language/Halide/ScheduleSpec.hs view
@@ -0,0 +1,229 @@+{-# LANGUAGE OverloadedRecordDot #-}++module Language.Halide.ScheduleSpec (spec) where++import Control.Monad (forM_)+import qualified Data.Text.IO as T+import GHC.TypeLits+import Language.Halide+import Test.Hspec+import Test.Hspec.QuickCheck+import Utils++checkScheduleRoundTrip :: (KnownNat n, IsHalideType a) => IO (Func t n a) -> (Func t n a -> IO ()) -> Expectation+checkScheduleRoundTrip prepare schedule = do+ f1 <- prepare+ f2 <- prepare++ schedule f1+ s1 <- getStageSchedule =<< getStage f1+ l1 <- prettyLoopNest f1++ applySchedule s1 =<< getStage f2+ s2 <- getStageSchedule =<< getStage f2+ l2 <- prettyLoopNest f2++ l1 `shouldBe` l2+ s1 `shouldBeEqForTesting` s2++spec :: Spec+spec = do+ describe "Extracts schedules" $ do+ it "supports vectorize" $ do+ [x, y, z] <- mapM mkVar ["x", "y", "z"]+ xInner <- mkVar "xInner"+ f <- define "f" (x, y, z) $ sin (cast @Float (x * y * z))+ void $+ split TailAuto x (x, xInner) 2 f+ >>= vectorize xInner+ schedule <- getStageSchedule =<< getStage f+ head schedule.dims `shouldBe` Dim "x.xInner" ForVectorized DeviceNone DimPureVar++ it "supports fuse" $ do+ [x, y, z] <- mapM mkVar ["x", "y", "z"]+ k <- mkVar "k"+ f <- define "f" (x, y, z) $ sin (cast @Float (x * y * z))+ void $ fuse (y, z) k f+ schedule <- getStageSchedule =<< getStage f+ print schedule++ describe "Applies schedules" $ do+ it "supports split" $ do+ let prepare = do+ [x, y, z] <- mapM mkVar ["x", "y", "z"]+ define "f" (x, y, z) $ sin (cast @Float (x * y * z))+ schedule f = do+ [x, _, _] <- getArgs f+ xInner <- mkVar "xInner"+ void $ split TailAuto x (x, xInner) 2 f+ checkScheduleRoundTrip prepare schedule+ it "supports fuse" $ do+ -- pendingWith "fails for unknown reason"+ let prepare = do+ [x, y, z] <- mapM mkVar ["x", "y", "z"]+ define "f" (x, y, z) $ sin (cast @Float (x * y * z))+ schedule f = do+ [_, y, z] <- getArgs f+ k <- mkVar "k"+ void $ fuse (y, z) k f+ checkScheduleRoundTrip prepare schedule+ it "supports vectorize" $ do+ -- pendingWith "fails for unknown reason"+ let prepare = do+ [x, y, z] <- mapM mkVar ["x", "y", "z"]+ define "f" (x, y, z) $ sin (cast @Float (x * y * z))+ schedule f = do+ [x, _, _] <- getArgs f+ xOuter <- mkVar "xOuter"+ xInnerOuter <- mkVar "xInnerOuter"+ xInnerInner <- mkVar "xInnerInner"+ void $+ split TailAuto x (x, xOuter) 4 f+ >>= split TailAuto xOuter (xInnerOuter, xInnerInner) 2+ >>= vectorize xInnerInner+ checkScheduleRoundTrip prepare schedule+ it "supports computeWith" $ do+ -- pendingWith "fails for unknown reason"+ let prepare = do+ x <- mkVar "x"+ y <- mkVar "y"+ f <- define "f" (x, y) $ x + y+ g <- define "g" (x, y) $ x - y+ h <- define "h" (x, y) $ f ! (x, y) + g ! (x, y)+ estimate x 0 200 h+ estimate y 0 200 h+ pure h+ let schedule h = do+ loadAutoScheduler Adams2019+ T.putStrLn =<< applyAutoScheduler Adams2019 hostTarget h+ checkScheduleRoundTrip prepare schedule++ -- [x, _, _] <- getArgs f+ -- xOuter <- mkVar "xOuter"+ -- xInnerOuter <- mkVar "xInnerOuter"+ -- xInnerInner <- mkVar "xInnerInner"+ -- void $+ -- split TailAuto x (x, xOuter) 4 f+ -- >>= split TailAuto xOuter (xInnerOuter, xInnerInner) 2+ -- >>= vectorize xInnerInner+ -- computeRoot f+ -- computeRoot g+ -- computeRoot k+ -- xi <- mkVar "xi"+ -- xo <- mkVar "xo"+ -- split TailAuto x (xo, xi) 8 f+ -- split TailAuto x (xo, xi) 8 g+ -- split TailAuto x (xo, xi) 8 k+ -- l <- getLoopLevelAtStage f xo 0+ -- print l+ -- computeWith LoopAlignAuto g l+ -- computeWith LoopAlignAuto f =<< getLoopLevelAtStage k xo 0++ prop "supports autoschedulers" $ do+ -- pendingWith "fails for unknown reason"+ let prepare1 = do+ [x, y] <- mapM mkVar ["x", "y"]+ f <- define "f" (x, y) $ x * y+ estimate x 0 100 f+ estimate y 0 100 f+ pure f+ let prepare2 = do+ x <- mkVar "x"+ y <- mkVar "y"+ f <- define "f" (x, y) $ x + y+ g <- define "g" (x, y) $ x - y+ h <- define "h" (x, y) $ f ! (x, y) + g ! (x, y)+ estimate x 0 200 h+ estimate y 0 200 h+ pure h+ let schedule target (Just scheduler) f = do+ loadAutoScheduler scheduler+ void $ applyAutoScheduler scheduler target f+ schedule _ _ _ = pure ()++ forM_ [prepare1, prepare2] $ \prepare -> do+ checkScheduleRoundTrip prepare (schedule hostTarget Nothing)+ checkScheduleRoundTrip prepare (schedule hostTarget (Just Adams2019))+ checkScheduleRoundTrip prepare (schedule hostTarget (Just Li2018))+ checkScheduleRoundTrip prepare (schedule hostTarget (Just Mullapudi2016))++-- (x, y, z, xInner, f1) <- prepare+-- split TailAuto x (x, xInner) 2 f1+-- nest1 <- prettyLoopNest f1+-- schedule1 <- getStageSchedule =<< getStage f1+-- print schedule1+-- (_, _, _, _, f2) <- prepare+-- applySchedule schedule1 =<< getStage f2+-- schedule2 <- getStageSchedule =<< getStage f2+-- nest2 <- prettyLoopNest f2+-- nest1 `shouldBe` nest2+-- T.putStrLn nest2+-- print schedule2++{-+describe "prints schedules" $ do+ it "of auto-scheduled pipelines" $ do+ let builder :: Bool -> Target -> Func 'ParamTy 1 Int64 -> IO (Func 'FuncTy 1 Float)+ builder useAutoScheduler target src = do+ i <- mkVar "i"+ dest <- define "dest1" i $ sin (cast @Float (src ! i))+ -- dim 0 src >>= setEstimate 0 1000+ -- dim 0 src >>= setMin 0 >>= setStride 1 >>= print+ -- schedule <- do+ estimate i 0 1000 dest++ when useAutoScheduler $ do+ loadAutoScheduler Adams2019+ T.putStrLn =<< applyAutoScheduler Adams2019 target dest+ print =<< getStageSchedule =<< getStage dest+ -- print =<< getStageSchedule =<< getStage dest++ -- T.putStrLn =<< prettyLoopNest dest+ -- T.putStrLn =<< prettyLoopNest clone+ -- schedule <- getStageSchedule dest+ -- print schedule.dims+ -- print =<< (getSplits <$> getStageSchedule dest)+ pure dest+ let target = hostTarget -- setFeature FeatureOpenCL hostTarget+ copy <- compileForTarget target (builder True target)+ -- let src :: S.Vector Int64+ -- src = S.generate 100 fromIntegral+ pure ()+ -}+{-+it "of computeWith" $ do+ x <- mkVar "x"+ y <- mkVar "y"+ f <- define "f" (x, y) $ x + y+ g <- define "g" (x, y) $ x - y+ k <- define "k" (x, y) $ x * y+ h <- define "h" (x, y) $ f ! (x, y) + g ! (x, y) + k ! (x, y)+ computeRoot f+ computeRoot g+ computeRoot k+ xi <- mkVar "xi"+ xo <- mkVar "xo"+ split TailAuto x (xo, xi) 8 f+ split TailAuto x (xo, xi) 8 g+ split TailAuto x (xo, xi) 8 k+ l <- getLoopLevelAtStage f xo 0+ print l+ computeWith LoopAlignAuto g l+ computeWith LoopAlignAuto f =<< getLoopLevelAtStage k xo 0++ hPutStrLn stderr =<< prettyLoopNest h++ schedule <- getStageSchedule =<< getStage g+ print schedule+-}++-- prettyLoopNest h >>= \s -> do+-- s `shouldContainText` "for x.xi"+-- s `shouldContainText` "for fused.y"+-- s `shouldContainText` "for x.fused.xo"++-- dest <- SM.new (S.length src)+-- withHalideBuffer src $ \srcPtr ->+-- withHalideBuffer dest $ \destPtr ->+-- copy srcPtr destPtr+-- S.unsafeFreeze dest `shouldReturn` src
+ test/Language/Halide/TargetSpec.hs view
@@ -0,0 +1,10 @@+module Language.Halide.TargetSpec (spec) where++import Language.Halide+import Test.Hspec++spec :: Spec+spec = do+ describe "setFeature" $ do+ it "adds features to JIT targets" $ do+ setFeature FeatureCUDA hostTarget `shouldSatisfy` hasGpuFeature
+ test/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}
+ test/Utils.hs view
@@ -0,0 +1,141 @@+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE OverloadedRecordDot #-}+{-# LANGUAGE UndecidableInstances #-}++module Utils+ ( shouldContainText+ , shouldNotContainText+ , appearsBeforeText+ , shouldApproxBe+ , testOnGpu+ , approx+ , approx'+ , (&)+ , void+ , T.hPutStrLn+ , stderr+ , HasEpsilon+ , eps+ , showInCodeLenses+ , EqForTesting (..)+ , shouldBeEqForTesting+ )+where++import Control.Exception (throwIO)+import Control.Monad (unless, void)+import Data.Function ((&))+import Data.Text (Text, unpack)+import qualified Data.Text as T+import qualified Data.Text.IO as T+import GHC.Exts (IsList (..))+import GHC.Stack+import Language.Halide+import System.IO (stderr)+import Test.HUnit+import Test.HUnit.Lang (FailureReason (..), HUnitFailure (..))+import Test.Hspec++shouldContainText :: Text -> Text -> Expectation+a `shouldContainText` b = T.unpack a `shouldContain` T.unpack b++shouldNotContainText :: Text -> Text -> Expectation+a `shouldNotContainText` b = T.unpack a `shouldNotContain` T.unpack b++appearsBeforeText :: Text -> Text -> Text -> Expectation+appearsBeforeText a b t = do+ t `shouldContainText` b+ fst (T.breakOn b t) `shouldContainText` a++testOnGpu :: (Target -> Expectation) -> Expectation+testOnGpu f =+ case gpuTarget of+ Just t -> f t+ Nothing -> pendingWith "no GPU target available"++class Num a => HasEpsilon a where+ eps :: a++instance HasEpsilon Float where+ eps = 1.1920929e-7++instance HasEpsilon Double where+ eps = 2.220446049250313e-16++approx :: (Ord a, Num a) => a -> a -> a -> a -> Bool+approx rtol atol a b = abs (a - b) <= max atol (rtol * max (abs a) (abs b))++approx' :: (Ord a, HasEpsilon a) => a -> a -> Bool+approx' a b = approx (2 * eps * max (abs a) (abs b)) (4 * eps) a b++shouldApproxBe :: (Ord a, Num a, Show a) => a -> a -> a -> a -> Expectation+shouldApproxBe rtol atol a b+ | approx rtol atol a b = pure ()+ | otherwise = assertFailure $ "expected " <> show a <> ", but got " <> show b++showInCodeLenses :: Text -> IO String+showInCodeLenses v = error (unpack v)++class EqForTesting a where+ equalForTesting :: a -> a -> Bool+ default equalForTesting :: Eq a => a -> a -> Bool+ a `equalForTesting` b = a == b++instance EqForTesting a => EqForTesting [a] where+ as `equalForTesting` bs = and $ zipWith equalForTesting as bs++instance EqForTesting (Expr Int32) where+ a `equalForTesting` b+ | (Just aInt, Just bInt) <- (toIntImm a, toIntImm b) = aInt == bInt+ | otherwise = show a == show b++instance EqForTesting SplitContents where+ a `equalForTesting` b =+ and+ [ a.splitOld == b.splitOld+ , a.splitOuter == b.splitOuter+ , a.splitInner == b.splitInner+ , a.splitFactor `equalForTesting` b.splitFactor+ , a.splitExact == b.splitExact+ , a.splitTail == b.splitTail+ ]++instance EqForTesting Split where+ (SplitVar a) `equalForTesting` (SplitVar b) = a `equalForTesting` b+ (FuseVars a) `equalForTesting` (FuseVars b) = a == b+ _ `equalForTesting` _ = False++instance EqForTesting ReductionVariable where+ a `equalForTesting` b = a.varName == b.varName && a.minExpr `equalForTesting` b.minExpr && a.extentExpr `equalForTesting` b.extentExpr++instance EqForTesting PrefetchDirective where+ a `equalForTesting` b =+ and+ [ a.prefetchFunc == b.prefetchFunc+ , a.prefetchAt == b.prefetchAt+ , a.prefetchFrom == b.prefetchFrom+ , a.prefetchOffset `equalForTesting` b.prefetchOffset+ , a.prefetchStrategy == b.prefetchStrategy+ ]++instance EqForTesting StageSchedule where+ a `equalForTesting` b =+ and+ [ a.rvars `equalForTesting` b.rvars+ , a.dims == b.dims+ , a.prefetches `equalForTesting` b.prefetches+ , a.fuseLevel == b.fuseLevel+ , a.fusedPairs == b.fusedPairs+ , a.allowRaceConditions == b.allowRaceConditions+ ]++shouldBeEqForTesting :: (HasCallStack, EqForTesting a, Show a) => a -> a -> Expectation+shouldBeEqForTesting actual expected =+ unless (actual `equalForTesting` expected) $ do+ throwIO (HUnitFailure location $ ExpectedButGot Nothing expectedMsg actualMsg)+ where+ expectedMsg = show expected+ actualMsg = show actual+ location = case reverse (toList callStack) of+ (_, loc) : _ -> Just loc+ [] -> Nothing