diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -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 diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,80 @@ +# halide-haskell + +[![GitHub CI](https://github.com/twesterhout/halide-haskell/actions/workflows/ci.yml/badge.svg)](https://github.com/twesterhout/halide-haskell/actions/workflows/ci.yml) +[![Hackage](https://img.shields.io/hackage/v/halide-haskell.svg?logo=haskell)](https://hackage.haskell.org/package/halide-haskell-0.0.1.0/candidate) +[![BSD-3-Clause license](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)](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 +``` diff --git a/example/Example01.hs b/example/Example01.hs new file mode 100644 --- /dev/null +++ b/example/Example01.hs @@ -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 diff --git a/example/GettingStarted.hs b/example/GettingStarted.hs new file mode 100644 --- /dev/null +++ b/example/GettingStarted.hs @@ -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) diff --git a/halide-haskell.cabal b/halide-haskell.cabal new file mode 100644 --- /dev/null +++ b/halide-haskell.cabal @@ -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 diff --git a/src/Language/Halide.hs b/src/Language/Halide.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide.hs @@ -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 diff --git a/src/Language/Halide/Buffer.hs b/src/Language/Halide/Buffer.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Buffer.hs @@ -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)) diff --git a/src/Language/Halide/Context.hs b/src/Language/Halide/Context.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Context.hs @@ -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 +-- > auto handle_halide_exceptions(Func&& func); +-- > +-- > template +-- > 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 "" + , C.include "" + , C.include "" + , 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 \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 \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\ + \" diff --git a/src/Language/Halide/Dimension.hs b/src/Language/Halide/Dimension.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Dimension.hs @@ -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 diff --git a/src/Language/Halide/Expr.hs b/src/Language/Halide/Expr.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Expr.hs @@ -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($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::RDom{ + *$(const Halide::Expr* min'), + *$(const Halide::Expr* extent'), + std::string{$bs-ptr:s, static_cast($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(); + 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::Expr* e)->type()); } |] + peek t + getHalideType (Var fp) = + withForeignPtr fp $ \e -> alloca $ \t -> do + [CU.block| void { + *$(halide_type_t* t) = static_cast( + static_cast(*$(Halide::Var* e)).type()); } |] + peek t + getHalideType (RVar fp) = + withForeignPtr fp $ \e -> alloca $ \t -> do + [CU.block| void { + *$(halide_type_t* t) = static_cast( + static_cast(*$(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::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::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($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::Var* varPtr))} } |] +forceExpr (RVar fp) = + withForeignPtr fp $ \rvarPtr -> + cxxConstructExpr $ \ptr -> + [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{ + static_cast(*$(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* 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 diff --git a/src/Language/Halide/Func.hs b/src/Language/Halide/Func.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Func.hs @@ -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* 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($(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($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* vars'); + auto& stage = *$(Halide::Stage* stage'); + auto const device = static_cast($(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* vars'); + auto& stage = *$(Halide::Stage* stage'); + auto const device = static_cast($(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($(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($(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(Halide::TailStrategy::RoundUp) } |] + TailGuardWithIf -> [CU.pure| int { static_cast(Halide::TailStrategy::GuardWithIf) } |] + TailPredicate -> [CU.pure| int { static_cast(Halide::TailStrategy::Predicate) } |] + TailPredicateLoads -> [CU.pure| int { static_cast(Halide::TailStrategy::PredicateLoads) } |] + TailPredicateStores -> [CU.pure| int { static_cast(Halide::TailStrategy::PredicateStores) } |] + TailShiftInwards -> [CU.pure| int { static_cast(Halide::TailStrategy::ShiftInwards) } |] + TailAuto -> [CU.pure| int { static_cast(Halide::TailStrategy::Auto) } |] + toEnum k + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::RoundUp) } |] = TailRoundUp + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::GuardWithIf) } |] = TailGuardWithIf + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::Predicate) } |] = TailPredicate + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::PredicateLoads) } |] = TailPredicateLoads + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::PredicateStores) } |] = TailPredicateStores + | fromIntegral k == [CU.pure| int { static_cast(Halide::TailStrategy::ShiftInwards) } |] = TailShiftInwards + | fromIntegral k == [CU.pure| int { static_cast(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($(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($(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* 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* { + new std::vector{$(const Halide::Func* func')->args()} } |] + destroy v = [CU.exp| void { delete $(std::vector* v) } |] + bracket allocate destroy $ \v -> do + n <- [CU.exp| size_t { $(const std::vector* v)->size() } |] + forM [0 .. n - 1] $ \i -> + fmap Var . cxxConstruct $ \ptr -> + [CU.exp| void { + new ($(Halide::Var* ptr)) Halide::Var{$(const std::vector* 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($(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($(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* 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* 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($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::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($bs-len:s)}}; + f(*$(std::vector* 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* 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* 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::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::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) diff --git a/src/Language/Halide/Kernel.hs b/src/Language/Halide/Kernel.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Kernel.hs @@ -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* 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* 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* { + auto p = new std::vector{}; + p->reserve($(size_t count)); + return p; + } |] + destroy p = [CU.exp| void { delete $(std::vector* 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* 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(Halide::StmtOutputFormat::Text) } |] + StmtHTML -> [CU.pure| int { static_cast(Halide::StmtOutputFormat::HTML) } |] + toEnum k + | fromIntegral k == [CU.pure| int { static_cast(Halide::StmtOutputFormat::Text) } |] = StmtText + | fromIntegral k == [CU.pure| int { static_cast(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($bs-len:s)}, + *$(const std::vector* v), + static_cast($(int o)), + *$(Halide::Target* t)); + }); + } |] + T.readFile (dir <> "/code.stmt") diff --git a/src/Language/Halide/LoopLevel.hs b/src/Language/Halide/LoopLevel.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/LoopLevel.hs @@ -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(Halide::LoopAlignStrategy::AlignStart) } |] + LoopAlignEnd -> [CU.pure| int { static_cast(Halide::LoopAlignStrategy::AlignEnd) } |] + LoopNoAlign -> [CU.pure| int { static_cast(Halide::LoopAlignStrategy::NoAlign) } |] + LoopAlignAuto -> [CU.pure| int { static_cast(Halide::LoopAlignStrategy::Auto) } |] + toEnum k + | fromIntegral k == [CU.pure| int { static_cast(Halide::LoopAlignStrategy::AlignStart) } |] = LoopAlignStart + | fromIntegral k == [CU.pure| int { static_cast(Halide::LoopAlignStrategy::AlignEnd) } |] = LoopAlignEnd + | fromIntegral k == [CU.pure| int { static_cast(Halide::LoopAlignStrategy::NoAlign) } |] = LoopNoAlign + | fromIntegral k == [CU.pure| int { static_cast(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 diff --git a/src/Language/Halide/Prelude.hs b/src/Language/Halide/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Prelude.hs @@ -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 diff --git a/src/Language/Halide/RedundantConstraints.hs b/src/Language/Halide/RedundantConstraints.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/RedundantConstraints.hs @@ -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 _ = () diff --git a/src/Language/Halide/Schedule.hs b/src/Language/Halide/Schedule.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Schedule.hs @@ -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(Halide::Internal::ForType::Serial) } |] = + ForSerial + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::Parallel) } |] = + ForParallel + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::Vectorized) } |] = + ForVectorized + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::Unrolled) } |] = + ForUnrolled + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::Extern) } |] = + ForExtern + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::GPUBlock) } |] = + ForGPUBlock + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::ForType::GPUThread) } |] = + ForGPUThread + | fromIntegral k == [CU.pure| int { static_cast(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(Halide::Internal::DimType::PureVar) } |] = + DimPureVar + | fromIntegral k == [CU.pure| int { static_cast(Halide::Internal::DimType::PureRVar) } |] = + DimPureRVar + | fromIntegral k == [CU.pure| int { static_cast(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(Halide::PrefetchBoundStrategy::Clamp) } |] = + PrefetchClamp + | fromIntegral k == [CU.pure| int { static_cast(Halide::PrefetchBoundStrategy::GuardWithIf) } |] = + PrefetchGuardWithIf + | fromIntegral k == [CU.pure| int { static_cast(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($(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* { + &$(const Halide::Internal::StageSchedule* schedule)->rvars() } |] + +getSplits :: Ptr CxxStageSchedule -> IO [Split] +getSplits schedule = + peekCxxVector + =<< [CU.exp| const std::vector* { + &$(const Halide::Internal::StageSchedule* schedule)->splits() } |] + +getDims :: Ptr CxxStageSchedule -> IO [Dim] +getDims schedule = + peekCxxVector + =<< [CU.exp| const std::vector* { + &$(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* { + &$(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($(Halide::Internal::Dim* p)->for_type) } |] + device <- + toEnum . fromIntegral + <$> [CU.exp| int { static_cast($(Halide::Internal::Dim* p)->device_api) } |] + dimType' <- + toEnum . fromIntegral + <$> [CU.exp| int { static_cast($(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($(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($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] diff --git a/src/Language/Halide/Target.hs b/src/Language/Halide/Target.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Target.hs @@ -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($(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(Halide::DeviceAPI::None) } |] + DeviceHost -> [CU.pure| int { static_cast(Halide::DeviceAPI::Host) } |] + DeviceDefaultGPU -> [CU.pure| int { static_cast(Halide::DeviceAPI::Default_GPU) } |] + DeviceCUDA -> [CU.pure| int { static_cast(Halide::DeviceAPI::CUDA) } |] + DeviceOpenCL -> [CU.pure| int { static_cast(Halide::DeviceAPI::OpenCL) } |] + DeviceOpenGLCompute -> [CU.pure| int { static_cast(Halide::DeviceAPI::OpenGLCompute) } |] + DeviceMetal -> [CU.pure| int { static_cast(Halide::DeviceAPI::Metal) } |] + DeviceHexagon -> [CU.pure| int { static_cast(Halide::DeviceAPI::Hexagon) } |] + DeviceHexagonDma -> [CU.pure| int { static_cast(Halide::DeviceAPI::HexagonDma) } |] + DeviceD3D12Compute -> [CU.pure| int { static_cast(Halide::DeviceAPI::D3D12Compute) } |] + toEnum k + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::None) } |] = DeviceNone + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::Host) } |] = DeviceHost + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::Default_GPU) } |] = DeviceDefaultGPU + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::CUDA) } |] = DeviceCUDA + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::OpenCL) } |] = DeviceOpenCL + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::OpenGLCompute) } |] = DeviceOpenGLCompute + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::Metal) } |] = DeviceMetal + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::Hexagon) } |] = DeviceHexagon + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::HexagonDma) } |] = DeviceHexagonDma + | fromIntegral k == [CU.pure| int { static_cast(Halide::DeviceAPI::D3D12Compute) } |] = DeviceD3D12Compute + | fromIntegral k == [CU.pure| int { static_cast(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($(int api))); +-- return (feature == Halide::Target::FeatureEnd) ? (-1) : static_cast(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 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 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 diff --git a/src/Language/Halide/Trace.hs b/src/Language/Halide/Trace.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Trace.hs @@ -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) diff --git a/src/Language/Halide/Type.hs b/src/Language/Halide/Type.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Type.hs @@ -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*@ +-- 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 diff --git a/src/Language/Halide/Utils.hs b/src/Language/Halide/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Language/Halide/Utils.hs @@ -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 diff --git a/test/Language/Halide/BufferSpec.hs b/test/Language/Halide/BufferSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/BufferSpec.hs @@ -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 diff --git a/test/Language/Halide/ExprSpec.hs b/test/Language/Halide/ExprSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/ExprSpec.hs @@ -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 diff --git a/test/Language/Halide/FuncSpec.hs b/test/Language/Halide/FuncSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/FuncSpec.hs @@ -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" + s `shouldContainText` "gpu_block x" + + 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" + + 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" + s `shouldContainText` "gpu_thread x" + + 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" + s `shouldContainText` "gpu_thread x" + + 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]) diff --git a/test/Language/Halide/KernelSpec.hs b/test/Language/Halide/KernelSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/KernelSpec.hs @@ -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 {" diff --git a/test/Language/Halide/LoopLevelSpec.hs b/test/Language/Halide/LoopLevelSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/LoopLevelSpec.hs @@ -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 diff --git a/test/Language/Halide/ScheduleSpec.hs b/test/Language/Halide/ScheduleSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/ScheduleSpec.hs @@ -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 diff --git a/test/Language/Halide/TargetSpec.hs b/test/Language/Halide/TargetSpec.hs new file mode 100644 --- /dev/null +++ b/test/Language/Halide/TargetSpec.hs @@ -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 diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/test/Utils.hs b/test/Utils.hs new file mode 100644 --- /dev/null +++ b/test/Utils.hs @@ -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