halide-haskell 0.0.1.0 → 0.0.2.0

raw patch · 25 files changed

+2094/−911 lines, 25 filesdep ~HUnitdep ~QuickCheckdep ~hspecnew-component:exe:halide-minimal-gpubinary-addedPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: HUnit, QuickCheck, hspec, vector

API changes (from Hackage documentation)

Files

• CHANGELOG.md +0/−11
• README.md +129/−13
• construction.png binary
• example/GettingStarted.hs +3/−3
• example/MinimalGPU.hs +28/−0
• halide-haskell.cabal +33/−13
• src/Language/Halide.hs +28/−5
• src/Language/Halide/BoundaryConditions.hs +44/−0
• src/Language/Halide/Buffer.hs +373/−57
• src/Language/Halide/Context.hs +23/−13
• src/Language/Halide/Expr.hs +606/−52
• src/Language/Halide/Func.hs +293/−353
• src/Language/Halide/Kernel.hs +268/−162
• src/Language/Halide/Prelude.hs +0/−41
• src/Language/Halide/Schedule.hs +11/−11
• src/Language/Halide/Target.hs +16/−6
• src/Language/Halide/Trace.hs +12/−9
• src/Language/Halide/Type.hs +16/−64
• test/Language/Halide/BoundaryConditionsSpec.hs +25/−0
• test/Language/Halide/BufferSpec.hs +36/−1
• test/Language/Halide/ExprSpec.hs +69/−54
• test/Language/Halide/FuncSpec.hs +34/−6
• test/Language/Halide/KernelSpec.hs +6/−9
• test/Language/Halide/ScheduleSpec.hs +15/−2
• test/Utils.hs +26/−26

@@ -1,11 +0,0 @@-# 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

@@ -1,9 +1,27 @@-# halide-haskell+<h1 align="center">+halide-haskell+</h1> -[![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)+<div align="center"> +<img src="assets/construction_1f6a7.png" width="32">This project is still under heavy development and might not be production-ready<img src="assets/construction_1f6a7.png" width="32"><br>+We encourage experimenting with it and reporting any issues you run into via+[Github issues](https://github.com/twesterhout/halide-haskell/issues).++<br />++[**Tutorials**](https://github.com/twesterhout/halide-haskell/tree/master/tutorials) | [**Documentation**](https://hackage.haskell.org/package/halide-haskell-0.0.1.0) | [**Showcases**](https://github.com/twesterhout/halide-haskell-examples)++[![license](https://img.shields.io/github/license/twesterhout/halide-haskell.svg?style=flat-square)](LICENSE)+[![build](https://img.shields.io/github/actions/workflow/status/twesterhout/halide-haskell/ci.yml?style=flat-square)](https://github.com/twesterhout/halide-haskell/actions/workflows/ci.yml)+[![Hackage](https://img.shields.io/hackage/v/halide-haskell?style=flat-square)](https://hackage.haskell.org/package/halide-haskell)++</div>++<table>+<tr>+<td>+ [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@@ -12,22 +30,31 @@ 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++** 😋.+</td>+</tr>+</table> -  - [Tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials)-  - [Reference documentation](https://hackage.haskell.org/package/halide-haskell-0.0.1.0)+<h4 align="center" >+This package provides Haskell bindings that allow to write Halide embedded in+Haskell without C++ 😋.+</h4> -## 🚀 Getting started+  - [Example usage](#-example-usage)+  - [Installing](#-installing)+  - [Motivation](#-motivation)+  - [Contributing](#-contributing) ++## 🚀 Example usage+ As a simple example, here's how you could implement array addition with halide-haskell:  ```haskell-{-# LANGUAGE AllowAmbiguousTypes, DataKinds, OverloadedStrings #-}+{-# LANGUAGE AllowAmbiguousTypes, DataKinds, OverloadedStrings, ViewPatterns #-} import Language.Halide  -- The algorithm-mkArrayPlus = compile $ \a b -> do+mkArrayPlus = compile $ \(buffer "a" -> a) (buffer "b" -> b) -> do   -- Create an index variable   i <- mkVar "i"   -- Define the resulting function. We call it "out".@@ -57,8 +84,96 @@         print =<< peekToList out' ``` -For more examples, have a look a the [tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials).+For more examples, have a look at the [tutorials](https://github.com/twesterhout/halide-haskell/tree/master/tutorials). +## 💻 Installing++Currently, the easiest way to install the library is using+[Nix](https://nixos.org/). It is not a fundamental limitation, because the+library itself is just a normal Cabal-based Haskell project, but installing &+patching (not all our bug fixes have been upstreamed yet) the system+dependencies is just too much work without Nix.++So, once you have Nix installed, you can add halide-haskell to your flake+inputs like [this project+demonstrates](https://github.com/twesterhout/halide-haskell-examples/blob/main/flake.nix#L27)+and then include it in your `build-depends` section in the Cabal file.++If you just want to try building the library, type++```sh+nix build+```++and to run an example, try++```sh+nix run+nix run .#ghc927-intel-ocl.halide-haskell # for Intel OpenCL support+nix run .#ghc927-cuda.halide-haskell      # for CUDA support+nix run .#ghc944.halide-haskell           # to build with GHC 9.4.4 instead+```++(for OpenCL and CUDA, you may need to set `NIXPKGS_ALLOW_UNFREE=1`)++## 🤩 Motivation++The availability of Deep Learning frameworks such as+[PyTorch](https://pytorch.org/) or [JAX](https://github.com/google/jax) has+revolutionized array processing, independently of whether one works on Machine+Learning tasks or other numerical algorithms. The ecosystem in Haskell has been+catching up as well, and there are now multiple good array+libraries ([hmatrix](https://github.com/haskell-numerics/hmatrix),+[massiv](https://github.com/lehins/massiv),+[Accelerate](https://www.acceleratehs.org/),+[arrayfire-haskell](https://github.com/arrayfire/arrayfire-haskell),+[Hasktorch](https://github.com/hasktorch/hasktorch), are all high-quality+libraries). To accommodate multiple domains, such libraries+have to support hundreds if not thousands of operations (e.g. there are more+than 3.5 thousand of so called [“native” functions in PyTorch](https://github.com/pytorch/pytorch/blob/6a09847c42bf7d33ba0aea5b083eebd846661ce1/aten/src/ATen/native/native_functions.yaml)),+and this count does not include specializations for different device and data+types).++To overcome this difficulty, we propose to build a common extension mechanism+for Haskell array libraries. The mechanism is based on embedding the+[Halide](https://halide-lang.org/) language into Haskell that allows to+just-in-time (JIT) compile computational kernels for various hardware.++### 🤨 Why not Accelerate?++One might wonder, why write another package instead of relying on+[Accelerate](https://www.acceleratehs.org/) for the JIT compilation of the+kernels. Accelerate is a Haskell eDSL (embedded Domain Specific Language) for+collective operations on dense multi-dimensional arrays. It relies on+[LLVM](https://llvm.org/) to JIT compile the computational kernels for the+target architecture including multicore CPUs and GPUs. Users have to rely on+Accelerate to generate high-performance kernels and have no way to force some+low-level optimizations. For example, [Trevor L. McDonell et+al.](https://doi.org/10.1145/2887747.2804313) explain that the reason why+hand-written [CUDA](https://www.nvidia.com/en-gb/geforce/technologies/cuda/)+implementation of the [N-body+problem](https://en.wikipedia.org/wiki/N-body_problem) outperforms Accelerate+is the use of on-chip shared memory. Another example would be the matrix-matrix+product where achieving maximal performance requires writing no fewer than six+nested loops instead of the naive three ([ACM Trans. Math. Softw. 34, 3,+Article 12 (May 2008)](https://doi.org/10.1145/1356052.1356053)).+Accelerate has no way of knowing that such optimizations have to be applied and+cannot perform them automatically, and this is precisely the gap that we are+trying to fill by embedding Halide into Haskell.++Halide is a C++ eDSL for high-performance image and array processing. Its core+idea is to decouple the *algorithm* (i.e. what is computed) from the *schedule*+(i.e. where and when it is computed). The eDSL allows to quickly prototype and+test the algorithm and then move on to the optimization. Optimizations such as+fusion, tiling, parallelism and vectorization can be freely explored without+the risk of breaking the original algorithm definition. Schedulers can also be+generated automatically by [advanced optimization+algorithms](https://halide-lang.org/papers/autoscheduler2019.html)++Halide provides a lower level interface than Accelerate and thus does not aim+to replace it. Instead, Halide can be used to extend Accelerate, and later on+one might even think about using Halide as a backend for Accelerate.+ ## 🔨 Contributing  Currently, the best way to get started is to use Nix:@@ -67,7 +182,8 @@ nix develop ``` -This will drop you into a shell with all the necessary tools to build the code such that you can do+This will drop you into a shell with all the necessary tools to build the code+such that you can do  ```sh cabal build

binary file changed (absent → 4123 bytes)

@@ -1,14 +1,14 @@ module Main (main) where -import qualified Data.Vector.Storable as S-import qualified Data.Vector.Storable.Mutable as SM+import Data.Vector.Storable qualified as S+import Data.Vector.Storable.Mutable qualified 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+  kernel <- compile $ \(buffer "a" -> a) b -> do     -- Create an index variable     i <- mkVar "i"     -- Define the resulting function. We call it "out".

@@ -0,0 +1,28 @@+module Main (main) where++import Control.Monad (unless, void)+import Language.Halide++main :: IO ()+main = do+  [x, y, xo, xi, yo, yi] <- mapM mkVar ["x", "y", "xo", "xi", "yo", "yi"]+  f <- define "f" (x, y) $ x + y++  void $+    split TailAuto x (xo, xi) 16 f+      >>= split TailAuto y (yo, yi) 16+      >>= reorder [xi, yi, xo, yo]+      >>= gpuBlocks DeviceDefaultGPU (xo, yo)+      >>= gpuThreads DeviceDefaultGPU (xi, yi)++  case gpuTarget of+    Nothing -> putStrLn "no GPU target found; skipping ..."+    Just target -> do+      r <- realizeOnTarget (setFeature FeatureDebug target) f [32, 32] peekToList+      let expected = [[i + j | i <- [0 .. 31]] | j <- [0 .. 31]]+      unless (r == expected) . error $+        "wrong result:"+          <> "\n          got: "+          <> show r+          <> ",\n but expected: "+          <> show expected

@@ -1,6 +1,6 @@ cabal-version:   3.0 name:            halide-haskell-version:         0.0.1.0+version:         0.0.2.0 synopsis:        Haskell bindings to Halide description:   Halide is a programming language designed to make it easier to write@@ -9,10 +9,10 @@   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'@@ -30,10 +30,10 @@ copyright:       2022-2023 Tom Westerhout build-type:      Simple extra-doc-files:-  CHANGELOG.md+  construction.png   README.md -tested-with:     GHC ==9.2.4 || ==9.2.5 || ==9.4.4+tested-with:     GHC ==9.2.7 || ==9.4.4 || ==9.4.5  source-repository head   type:     git@@ -63,6 +63,7 @@   hs-source-dirs:  src   exposed-modules: Language.Halide   other-modules:+    Language.Halide.BoundaryConditions     Language.Halide.Buffer     Language.Halide.Context     Language.Halide.Dimension@@ -70,7 +71,6 @@     Language.Halide.Func     Language.Halide.Kernel     Language.Halide.LoopLevel-    Language.Halide.Prelude     Language.Halide.RedundantConstraints     Language.Halide.Schedule     Language.Halide.Target@@ -78,6 +78,7 @@     Language.Halide.Type     Language.Halide.Utils +  -- Language.Halide.Prelude   build-depends:     , bytestring        >=0.11.1.0 && <0.12     , constraints       >=0.13.4   && <0.14@@ -98,10 +99,17 @@     cpp-options:   -DUSE_DLOPEN=1     build-depends: unix >=2.7.2.2 && <3.0 -  extra-libraries:-    Halide-    stdc+++  extra-libraries: Halide +  -- if os(linux)+  --   extra-libraries: stdc++++  -- if os(osx)+  --   extra-libraries: iconv++  cxx-options:     -std=c++17 -Wall+  ghc-options:     -optcxx-std=c++17 -optcxx-Wall+ executable halide-haskell   import:         common-options   hs-source-dirs: example@@ -118,12 +126,21 @@     , halide-haskell     , vector +executable halide-minimal-gpu+  import:         common-options+  hs-source-dirs: example+  main-is:        MinimalGPU.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.BoundaryConditionsSpec     Language.Halide.BufferSpec     Language.Halide.ExprSpec     Language.Halide.FuncSpec@@ -135,14 +152,17 @@    build-depends:     , halide-haskell-    , hspec-    , HUnit+    , hspec           >=2.9.7   && <3+    , HUnit           >=1.6.2.0 && <1.7     , inline-c     , inline-c-cpp-    , QuickCheck+    , QuickCheck      >=2.14.2  && <3     , text     , vector -  ghc-options:    -threaded -rtsopts -with-rtsopts=-N+  cxx-options:    -std=c++17 -Wall+  ghc-options:+    -threaded -rtsopts -with-rtsopts=-N -optcxx-std=c++17 -optcxx-Wall+    -Wno-missing-export-lists  -- build-tools-depends: hspec-discover:hspec-discover

@@ -15,18 +15,22 @@   , RVar   , VarOrRVar   , IsHalideType+  , ReductionDomain      -- ** Creating   , mkExpr   , mkVar   , mkRVar+  , mkRDom+  , toRVars   , undef   , cast-  , bool+  , ifThenElse      -- ** Inspecting   , toIntImm   , printed+  , printedWhen   , evaluate      -- ** Comparisons@@ -41,9 +45,19 @@   , lte   , gt   , gte+  , min+  , max+  , div+  , mod +    -- ** Boolean functions+  , and+  , or+     -- * Functions   , Func (..)+  , Function+  , Parameter   , FuncTy (..)   , Stage (..) @@ -51,6 +65,8 @@   , define   , update   , (!)+  , repeatEdge+  , constantExterior      -- ** Inspecting   , getArgs@@ -66,8 +82,10 @@     -- | To easily test out your pipeline, there are helper functions to create 'HalideBuffer's without     -- worrying about the low-level representation.   , allocaCpuBuffer+  , allocaBuffer     -- | Buffers can also be converted to lists to easily print them for debugging.   , IsListPeek (..)+  , peekScalar     -- | 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@@ -84,6 +102,7 @@     --     -- The simplest way to build a t'Func' and then call 'realize' to evaluate it over a rectangular domain.   , realize+  , realizeOnTarget   , 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@@ -109,6 +128,7 @@   , setExtent   , setStride   , setEstimate+  , setScalarEstimate      -- ** Targets   , Target (..)@@ -162,7 +182,6 @@   , IsTuple (..)   , ToTuple   , FromTuple-  , IndexTuple   , Length   , All @@ -178,18 +197,20 @@   , colMajorStrides   , isDeviceDirty   , isHostDirty+  , getBufferExtent   , bufferCopyToHost+  , withCopiedToHost+  , withCropped   , module Language.Halide.Schedule-  , IsFuncBuilder-  , ReturnsFunc   , FunctionArguments   , FunctionReturn   , Curry (..)   , UnCurry (..)-  , Lowered+  , LoweredSignature      -- ** inline-c helpers   , importHalide+  , testWriteToStderr   , CxxExpr   , CxxVar   , CxxRVar@@ -210,6 +231,7 @@  import Foreign.Ptr (Ptr) import GHC.TypeLits (KnownNat)+import Language.Halide.BoundaryConditions import Language.Halide.Buffer import Language.Halide.Context import Language.Halide.Dimension@@ -221,3 +243,4 @@ import Language.Halide.Target import Language.Halide.Trace import Language.Halide.Type+import Prelude ()

@@ -0,0 +1,44 @@+{-# LANGUAGE QuasiQuotes #-}+{-# LANGUAGE TemplateHaskell #-}++-- |+-- Module      : Language.Halide.BoundaryConditions+-- Copyright   : (c) Tom Westerhout, 2023+module Language.Halide.BoundaryConditions+  ( repeatEdge+  , constantExterior+  )+where++import GHC.TypeLits+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.Type+import Prelude hiding (min, tail)++importHalide++-- | Impose a boundary condition such that the nearest edge sample is returned everywhere outside the given region.+--+-- For more information, see [@Halide::repeat_edge@](https://halide-lang.org/docs/namespace_halide_1_1_boundary_conditions.html#a0548f23db36e4a8a03690bc8bee1e850).+repeatEdge :: (KnownNat n, IsHalideType a) => Func 'ParamTy n (Expr a) -> IO (Func 'FuncTy n (Expr a))+repeatEdge source =+  withBufferParam source $ \source' ->+    wrapCxxFunc+      =<< [CU.exp| Halide::Func* { new Halide::Func{+            Halide::BoundaryConditions::repeat_edge(*$(const Halide::ImageParam* source'))} } |]++-- | Impose a boundary condition such that a given expression is returned everywhere outside the boundary.+--+-- For more information, see [@Halide::constant_exterior@](https://halide-lang.org/docs/namespace_halide_1_1_boundary_conditions.html#aa4ed713b5f9a6f13e6323f2a21d41d5e).+constantExterior :: (KnownNat n, IsHalideType a) => Expr a -> Func 'ParamTy n (Expr a) -> IO (Func 'FuncTy n (Expr a))+constantExterior value source =+  withBufferParam source $ \source' ->+    asExpr value $ \value' ->+      wrapCxxFunc+        =<< [CU.exp| Halide::Func* { new Halide::Func{+              Halide::BoundaryConditions::constant_exterior(+                *$(const Halide::ImageParam* source'), *$(const Halide::Expr* value'))} } |]

@@ -3,6 +3,7 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE UndecidableInstances #-}  -- | -- Module      : Language.Halide.Buffer@@ -23,8 +24,10 @@     -- | To easily test out your pipeline, there are helper functions to create 'HalideBuffer's without     -- worrying about the low-level representation.   , allocaCpuBuffer+  , allocaBuffer     -- | Buffers can also be converted to lists to easily print them for debugging.   , IsListPeek (..)+  , peekScalar     -- | 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@@ -43,31 +46,37 @@   , colMajorStrides   , isDeviceDirty   , isHostDirty+  , getBufferExtent   , bufferCopyToHost+  , withCopiedToHost+  , withCropped   ) where +import Control.Exception (bracket_) 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.List qualified as List import Data.Proxy-import qualified Data.Vector.Storable as S-import qualified Data.Vector.Storable.Mutable as SM+import Data.Vector.Storable qualified as S+import Data.Vector.Storable.Mutable qualified as SM import Data.Word+import Foreign.Marshal.Alloc (alloca, free, mallocBytes) 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.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.Context+import Language.Halide.Target import Language.Halide.Type+import Prelude hiding (min)  -- | Information about a dimension in a buffer. --@@ -291,6 +300,27 @@     S.unsafeWith v $ \cpuPtr ->       bufferFromPtrShape cpuPtr [d0, d1, d2] f +-- | Lists can also act as Halide buffers. __Use for testing only.__+instance IsHalideType a => IsHalideBuffer [[[[a]]]] 4 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))+        d3 = if d2 == 0 then 0 else length (head (head (head xs)))+        -- we want column-major ordering, so transpose first+        v =+          S.fromList+            . concat+            . concat+            . concatMap (fmap List.transpose . List.transpose . fmap List.transpose)+            . List.transpose+            . fmap (List.transpose . fmap List.transpose)+            $ xs+    when (S.length v /= d0 * d1 * d2 * d3) $+      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, d3] f+ whenM :: Monad m => m Bool -> m () -> m () whenM cond f =   cond >>= \case@@ -312,45 +342,151 @@   => [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+allocaCpuBuffer = allocaBuffer hostTarget++getTotalBytes :: Ptr RawHalideBuffer -> IO Int+getTotalBytes buf = do+  fromIntegral+    <$> [CU.block| size_t {+          auto const& b = *$(const halide_buffer_t* buf);+          auto const n = std::accumulate(b.dim, b.dim + b.dimensions, size_t{1},+                                         [](auto acc, auto const& dim) { return acc * dim.extent; });+          return n * (b.type.bits * b.type.lanes / 8);+        } |]++allocateHostMemory :: Ptr RawHalideBuffer -> IO ()+allocateHostMemory buf = do+  ptr <- mallocBytes =<< getTotalBytes buf+  [CU.block| void { $(halide_buffer_t* buf)->host = $(uint8_t* ptr); } |]++freeHostMemory :: Ptr RawHalideBuffer -> IO ()+freeHostMemory buf = do+  ptr <-+    [CU.block| uint8_t* {+      auto& b = *$(halide_buffer_t* buf);+      auto const p = b.host;+      b.host = nullptr;+      return p;+    } |]+  free ptr++allocateDeviceMemory :: Ptr HalideDeviceInterface -> Ptr RawHalideBuffer -> IO ()+allocateDeviceMemory interface buf = do+  [CU.block| void {+    auto const* interface = $(const halide_device_interface_t* interface);+    interface->device_malloc(nullptr, $(halide_buffer_t* buf), interface);+  } |]++freeDeviceMemory :: HasCallStack => Ptr RawHalideBuffer -> IO ()+freeDeviceMemory buf = do+  deviceInterface <-+    [CU.exp| const halide_device_interface_t* { $(const halide_buffer_t* buf)->device_interface } |]+  when (deviceInterface == nullPtr) $+    error "cannot free device memory: device_interface is NULL"+  [CU.block| void {+    $(halide_buffer_t* buf)->device_interface->device_free(nullptr, $(halide_buffer_t* buf));+    $(halide_buffer_t* buf)->device = 0;+  } |]++allocaBuffer+  :: forall n a b+   . (HasCallStack, KnownNat n, IsHalideType a)+  => Target+  -> [Int]+  -> (Ptr (HalideBuffer n a) -> IO b)+  -> IO b+allocaBuffer target shape action = do+  deviceInterface <- getDeviceInterface target+  let onHost = deviceInterface == nullPtr+  withArrayLen (zipWith simpleDimension shape (colMajorStrides shape)) $ \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 rawBuffer =+          RawHalideBuffer+            { halideBufferDevice = 0+            , halideBufferDeviceInterface = nullPtr+            , halideBufferHost = nullPtr+            , halideBufferFlags = 0+            , halideBufferType = halideTypeFor (Proxy :: Proxy a)+            , halideBufferDimensions = fromIntegral n+            , halideBufferDim = dim+            , halideBufferPadding = nullPtr+            }+    with rawBuffer $ \buf -> do+      let allocate+            | onHost = allocateHostMemory+            | otherwise = allocateDeviceMemory deviceInterface+      let deallocate+            | onHost = freeHostMemory+            | otherwise = freeDeviceMemory+      bracket_ (allocate buf) (deallocate buf) $ do+        r <- action (castPtr buf)+        isHostNull <- toBool <$> [CU.exp| bool { $(halide_buffer_t* buf)->host == nullptr } |]+        isDeviceNull <- toBool <$> [CU.exp| bool { $(halide_buffer_t* buf)->device == 0 } |]+        when (onHost && not isDeviceNull) . error $+          "buffer was allocated on host, but its device pointer is not NULL"+            <> "; did you forget a copyToHost in your pipeline?"+        when (not onHost && not isHostNull) . error $+          "buffer was allocated on device, but its host pointer is not NULL"+            <> "; did you add an extra copyToHost?"+        pure r++getDeviceInterface :: Target -> IO (Ptr HalideDeviceInterface)+getDeviceInterface target =+  case device of+    DeviceNone -> pure nullPtr+    DeviceHost -> pure nullPtr+    _ ->+      withCxxTarget target $ \target' ->+        [C.throwBlock| const halide_device_interface_t* {+          return handle_halide_exceptions([=](){+            auto const device = static_cast<Halide::DeviceAPI>($(int api));+            auto const& target = *$(const Halide::Target* target');+            return Halide::get_device_interface_for_device_api(device, target, "getDeviceInterface");+          });+        } |]   where-    numElements = foldl' (*) 1 shape+    device@(fromIntegral . fromEnum -> api) = deviceAPIForTarget target  -- | 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() } |] +-- | Set the @device_dirty@ flag to the given value.+setDeviceDirty :: Bool -> Ptr RawHalideBuffer -> IO ()+setDeviceDirty (fromIntegral . fromEnum -> b) p =+  [CU.exp| void { $(halide_buffer_t* p)->set_device_dirty($(bool b)) } |]+ -- | 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() } |] +-- | Set the @host_dirty@ flag to the given value.+setHostDirty :: Bool -> Ptr RawHalideBuffer -> IO ()+setHostDirty (fromIntegral . fromEnum -> b) p =+  [CU.exp| void { $(halide_buffer_t* p)->set_host_dirty($(bool b)) } |]+ -- | Copy the underlying memory from device to host.-bufferCopyToHost :: Ptr RawHalideBuffer -> IO ()-bufferCopyToHost p =-  [C.throwBlock| void {+bufferCopyToHost :: HasCallStack => Ptr RawHalideBuffer -> IO ()+bufferCopyToHost p = whenM (isDeviceDirty p) $ do+  raw <- peek p+  when (raw.halideBufferDeviceInterface == nullPtr) . error $+    "device_dirty is set, but device_interface is NULL"+  when (raw.halideBufferHost == nullPtr) . error $+    "host is NULL, did you forget to allocate memory?"+  [CU.block| 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);-    }+    buf.device_interface->copy_to_host(nullptr, &buf);   } |]+  whenM (isDeviceDirty p) . error $+    "device_dirty is set right after a copy_to_host; something went wrong..."  checkNumberOfDimensions :: forall n. (HasCallStack, KnownNat n) => RawHalideBuffer -> IO () checkNumberOfDimensions raw = do@@ -361,59 +497,239 @@         <> " != "         <> show raw.halideBufferDimensions +-- | Perform an action on a cropped buffer.+withCropped+  :: Ptr (HalideBuffer n a)+  -- ^ buffer+  -> Int+  -- ^ dimension+  -> Int+  -- ^ min+  -> Int+  -- ^ extent+  -> (Ptr (HalideBuffer n a) -> IO b)+  -- ^ what to do+  -> IO b+withCropped+  (castPtr -> src)+  (fromIntegral -> d)+  (fromIntegral -> min)+  (fromIntegral -> extent)+  action = do+    rank <- fromIntegral <$> [CU.exp| int { $(const halide_buffer_t* src)->dimensions } |]+    alloca $ \dst ->+      allocaArray rank $ \dstDim -> do+        [CU.block| void {+          auto const& src = *$(const halide_buffer_t* src);+          auto& dst = *$(halide_buffer_t* dst);+          auto const d = $(int d);++          dst = src;+          dst.dim = $(halide_dimension_t* dstDim);+          memcpy(dst.dim, src.dim, src.dimensions * sizeof(halide_dimension_t));++          if (dst.host != nullptr) {+            auto const shift = $(int min) - src.dim[d].min;+            dst.host += (shift * src.dim[d].stride) * ((src.type.bits + 7) / 8);+          }+          dst.dim[d].min = $(int min);+          dst.dim[d].extent = $(int extent);++          if (src.device != 0 && src.device_interface != nullptr) {+            src.device_interface->device_crop(nullptr, &src, &dst);+          }+        } |]+        action (castPtr dst)++getBufferExtent :: forall n a. KnownNat n => Ptr (HalideBuffer n a) -> Int -> IO Int+getBufferExtent (castPtr -> buf) (fromIntegral -> d)+  | d < fromIntegral (natVal (Proxy @n)) =+      fromIntegral <$> [CU.exp| int { $(const halide_buffer_t* buf)->dim[$(int d)].extent } |]+  | otherwise = error "index out of bounds"++peekScalar :: forall a. (HasCallStack, IsHalideType a) => Ptr (HalideBuffer 0 a) -> IO a+peekScalar p = withCopiedToHost p $ do+  raw <- peek (castPtr @_ @RawHalideBuffer p)+  checkNumberOfDimensions @0 raw+  when (raw.halideBufferHost == nullPtr) . error $ "host is NULL"+  peek $ castPtr @_ @a raw.halideBufferHost+ -- | 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]+-- class IsListPeek a where+--   type ListPeekElem a :: Type+--   peekToList :: HasCallStack => Ptr a -> IO [ListPeekElem a]+type family NestedList (n :: Nat) (a :: Type) where+  NestedList 0 a = a+  NestedList 1 a = [a]+  NestedList 2 a = [[a]]+  NestedList 3 a = [[[a]]]+  NestedList 4 a = [[[[a]]]]+  NestedList 5 a = [[[[[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)+type family NestedListLevel (a :: Type) :: Nat where+  NestedListLevel [a] = 1 + NestedListLevel a+  NestedListLevel a = 0 -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"+type family NestedListType (a :: Type) :: Type where+  NestedListType [a] = NestedListType a+  NestedListType a = a++class+  ( KnownNat n+  , IsHalideType a+  , NestedList n a ~ b+  , NestedListLevel b ~ n+  , NestedListType b ~ a+  ) =>+  IsListPeek n a b+    | n a -> b+    , n b -> a+    , a b -> n+  where+  peekToList :: HasCallStack => Ptr (HalideBuffer n a) -> IO b++instance+  (IsHalideType a, NestedListLevel [a] ~ 1, NestedListType [a] ~ a)+  => IsListPeek 1 a [a]+  where+  peekToList p = withCopiedToHost p $ do     raw <- peek (castPtr @_ @RawHalideBuffer p)     (HalideDimension min0 extent0 stride0 _) <- peekElemOff (halideBufferDim raw) 0     let ptr0 = castPtr @_ @a (halideBufferHost raw)+    when (ptr0 == nullPtr) . error $ "host is NULL"     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"+instance+  (IsHalideType a, NestedListLevel [[a]] ~ 2, NestedListType [[a]] ~ a)+  => IsListPeek 2 a [[a]]+  where+  peekToList p = withCopiedToHost p $ do     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)+    when (ptr0 == nullPtr) . error $ "host is NULL"     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"+instance+  (IsHalideType a, NestedListLevel [[[a]]] ~ 3, NestedListType [[[a]]] ~ a)+  => IsListPeek 3 a [[[a]]]+  where+  peekToList p = withCopiedToHost p $ do     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)+    when (ptr0 == nullPtr) . error $ "host is NULL"     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))++instance+  (IsHalideType a, NestedListLevel [[[[a]]]] ~ 4, NestedListType [[[[a]]]] ~ a)+  => IsListPeek 4 a [[[[a]]]]+  where+  peekToList p = withCopiedToHost p $ do+    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+    (HalideDimension min3 extent3 stride3 _) <- peekElemOff (halideBufferDim raw) 3+    let ptr0 = castPtr @_ @a (halideBufferHost raw)+    when (ptr0 == nullPtr) . error $ "host is NULL"+    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 -> do+          let ptr3 = ptr2 `advancePtr` fromIntegral (min2 + stride2 * i2)+          forM [0 .. extent3 - 1] $ \i3 ->+            peekElemOff ptr3 (fromIntegral (min3 + stride3 * i3))++-- | @withCopiedToHost buf action@ performs the action @action@ ensuring that @buf@ has been+-- copied to the host beforehand. If @buf@ is already on the host, no copying is performed.+withCopiedToHost :: Ptr (HalideBuffer n a) -> IO b -> IO b+withCopiedToHost (castPtr @_ @RawHalideBuffer -> buf) action = do+  raw <- peek buf+  let allocate = when (raw.halideBufferDevice /= 0) $ allocateHostMemory buf+      deallocate = when (raw.halideBufferDevice /= 0) $ freeHostMemory buf+  bracket_ allocate deallocate $ do+    when (raw.halideBufferDevice /= 0) $ do+      setDeviceDirty True buf+      bufferCopyToHost buf+    action++-- instance IsHalideType a => IsListPeek (HalideBuffer 0 a) where+--   type ListPeekElem (HalideBuffer 0 a) = a+--   peekToList p = withCopiedToHost p $ do+--     raw <- peek (castPtr @_ @RawHalideBuffer p)+--     checkNumberOfDimensions @0 raw+--     when (raw.halideBufferHost == nullPtr) . error $ "host is NULL"+--     fmap pure . peek $ castPtr @_ @a raw.halideBufferHost++-- instance IsHalideType a => IsListPeek (HalideBuffer 1 a) where+--   type ListPeekElem (HalideBuffer 1 a) = a+--   peekToList p = withCopiedToHost p $ do+--     raw <- peek (castPtr @_ @RawHalideBuffer p)+--     (HalideDimension min0 extent0 stride0 _) <- peekElemOff (halideBufferDim raw) 0+--     let ptr0 = castPtr @_ @a (halideBufferHost raw)+--     when (ptr0 == nullPtr) . error $ "host is NULL"+--     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 = withCopiedToHost p $ do+--     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)+--     when (ptr0 == nullPtr) . error $ "host is NULL"+--     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 = withCopiedToHost p $ do+--     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)+--     when (ptr0 == nullPtr) . error $ "host is NULL"+--     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))++-- instance IsHalideType a => IsListPeek (HalideBuffer 4 a) where+--   type ListPeekElem (HalideBuffer 4 a) = [[[a]]]+--   peekToList p = withCopiedToHost p $ do+--     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+--     (HalideDimension min3 extent3 stride3 _) <- peekElemOff (halideBufferDim raw) 3+--     let ptr0 = castPtr @_ @a (halideBufferHost raw)+--     when (ptr0 == nullPtr) . error $ "host is NULL"+--     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 -> do+--           let ptr3 = ptr2 `advancePtr` fromIntegral (min2 + stride2 * i2)+--           forM [0 .. extent3 - 1] $ \i3 ->+--             peekElemOff ptr3 (fromIntegral (min3 + stride3 * i3))

@@ -41,12 +41,12 @@   ) where -import qualified Language.C.Inline as C-import qualified Language.C.Inline.Cpp as C+import Language.C.Inline qualified as C+import Language.C.Inline.Cpp qualified 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+import Language.Haskell.TH qualified as TH  -- | One stop function to include all the neccessary machinery to call Halide functions via inline-c. --@@ -58,7 +58,9 @@     <$> sequence       [ C.context =<< halideCxt       , C.include "<Halide.h>"-      , C.include "<cxxabi.h>"+      , -- , C.include "<HalideRuntimeOpenCL.h>"+        -- , C.include "<HalideRuntimeCuda.h>"+        C.include "<cxxabi.h>"       , C.include "<dlfcn.h>"       , defineExceptionHandler       ]@@ -84,24 +86,28 @@         , ("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::Dim", "Dim")         , ("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::Internal::ReductionVariable", "ReductionVariable")+        , ("Halide::Internal::Split", "Split")+        , ("Halide::Internal::StageSchedule", "CxxStageSchedule")+        , ("Halide::Argument", "CxxArgument")+        , ("Halide::Buffer", "CxxBuffer")+        , ("Halide::LoopLevel", "CxxLoopLevel")+        , ("Halide::Stage", "CxxStage")+        , ("Halide::Range", "CxxRange")+        , ("Halide::RDom", "CxxRDom")+        , ("halide_buffer_t", "Language.Halide.Buffer.RawHalideBuffer")+        , ("halide_device_interface_t", "HalideDeviceInterface")+        , ("halide_dimension_t", "HalideDimension")         , ("halide_trace_event_t", "TraceEvent")         ]     optional :: (CIdentifier, String) -> Q [(CIdentifier, TypeQ)]@@ -136,4 +142,8 @@     \  stream << x;                                                   \n\     \  return new std::string{stream.str()};                          \n\     \}                                                                \n\+    \\n\+    \namespace Halide { namespace Internal {\n\+    \  std::string print_loop_nest(const std::vector<Function> &);\n\+    \} }\n\     \"

@@ -7,6 +7,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilyDependencies #-} {-# OPTIONS_GHC -Wno-orphans #-}  -- |@@ -18,10 +19,13 @@   , Var   , RVar   , VarOrRVar+  , ReductionDomain (..)   , Int32   , mkExpr   , mkVar   , mkRVar+  , mkRDom+  , toRVars   , cast   , eq   , neq@@ -29,7 +33,13 @@   , lte   , gt   , gte-  , bool+  , and+  , or+  , min+  , max+  , div+  , mod+  , ifThenElse   , 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.@@ -38,7 +48,9 @@     -- 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+  , printedWhen   , toIntImm+  , setScalarEstimate      -- * Internal   , exprToForeignPtr@@ -58,39 +70,88 @@   , binaryOp   , unaryOp   , checkType+  , testWriteToStderr+  , Solo (..)+  , IsTuple (..)+  , FromTuple+  , ToTuple+  , UniformTuple+  , UniformTupleProperties+  , proveUniformTupleProperties+  , IndexType+  , IndexTypeProperties+  , proveIndexTypeProperties+  , HasIndexType   ) where  import Control.Exception (bracket) import Control.Monad (unless)+import Data.Constraint import Data.IORef import Data.Int (Int32)+import Data.Kind import Data.Proxy import Data.Ratio (denominator, numerator)-import Data.Text (Text, unpack)+import Data.Text (Text, pack, unpack) import Data.Text.Encoding qualified as T+import Data.Tuple 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 Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr)+import Foreign.Storable 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.RedundantConstraints import Language.Halide.Type import Language.Halide.Utils import System.IO.Unsafe (unsafePerformIO)-import Prelude hiding (min)+import Type.Reflection+import Unsafe.Coerce+import Prelude hiding (and, div, max, min, mod, or) +-- | 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 single-dimensional span.+data Range = Range {rangeMin :: !(Expr Int32), rangeExtent :: !(Expr Int32)}++-- | Haskell counterpart of @Halide::Range@.+data CxxRange++-- | Haskell counterpart of @Halide::RDom@.+data CxxRDom+ importHalide  instanceCxxConstructible "Halide::Expr" instanceCxxConstructible "Halide::Var" instanceCxxConstructible "Halide::RVar" instanceCxxConstructible "Halide::VarOrRVar"+instanceCxxConstructible "Halide::Range"+instanceCxxConstructible "Halide::RDom"  defineIsHalideTypeInstances @@ -98,6 +159,7 @@ instanceHasCxxVector "Halide::Var" instanceHasCxxVector "Halide::RVar" instanceHasCxxVector "Halide::VarOrRVar"+instanceHasCxxVector "Halide::Range"  -- instanceCxxConstructible "Halide::Var" -- instanceCxxConstructible "Halide::RVar"@@ -109,27 +171,273 @@     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]+-- | Type family that maps @'Arguments' ts@ to the corresponding tuple type.+type family ToTuple t = s | s -> t where+  ToTuple '[] = ()+  ToTuple '[Expr a1] = Expr a1+  ToTuple '[a1, a2] = (a1, a2)+  ToTuple '[a1, a2, a3] = (a1, a2, a3)+  ToTuple '[a1, a2, a3, a4] = (a1, a2, a3, a4)+  ToTuple '[a1, a2, a3, a4, a5] = (a1, a2, a3, a4, a5)+  ToTuple '[a1, a2, a3, a4, a5, a6] = (a1, a2, a3, a4, a5, a6)+  ToTuple '[a1, a2, a3, a4, a5, a6, a7] = (a1, a2, a3, a4, a5, a6, a7)+  ToTuple '[a1, a2, a3, a4, a5, a6, a7, a8] = (a1, a2, a3, a4, a5, a6, a7, a8)+  ToTuple '[a1, a2, a3, a4, a5, a6, a7, a8, a9] = (a1, a2, a3, a4, a5, a6, a7, a8, a9)+  ToTuple '[a1, a2, a3, a4, a5, a6, a7, a8, a9, a10] = (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10) --- | A scalar expression in Halide.+-- | Type family that maps tuples to the corresponding @'Arguments' ts@ type. This is essentially the inverse+-- of 'ToTuple'.+type family FromTuple t = s | s -> t where+  FromTuple () = '[]+  FromTuple (Expr a1) = '[Expr a1]+  FromTuple (a1, a2) = '[a1, a2]+  FromTuple (a1, a2, a3) = '[a1, a2, a3]+  FromTuple (a1, a2, a3, a4) = '[a1, a2, a3, a4]+  FromTuple (a1, a2, a3, a4, a5) = '[a1, a2, a3, a4, a5]+  FromTuple (a1, a2, a3, a4, a5, a6) = '[a1, a2, a3, a4, a5, a6]+  FromTuple (a1, a2, a3, a4, a5, a6, a7) = '[a1, a2, a3, a4, a5, a6, a7]+  FromTuple (a1, a2, a3, a4, a5, a6, a7, a8) = '[a1, a2, a3, a4, a5, a6, a7, a8]+  FromTuple (a1, a2, a3, a4, a5, a6, a7, a8, a9) = '[a1, a2, a3, a4, a5, a6, a7, a8, a9]+  FromTuple (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10) = '[a1, a2, a3, a4, a5, a6, a7, a8, a9, a10]++-- | Specifies that there is an isomorphism between a type @a@ and a tuple @t@. ----- 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)))+-- 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 :: Arguments a -> t+  fromTuple :: t -> Arguments a +-- | Generates a tuple of @n@ elements of type @t@.+type family UniformTuple (n :: Nat) (t :: Type) = (tuple :: Type) | tuple -> n where+  UniformTuple 0 t = ()+  UniformTuple 1 t = Expr t+  UniformTuple 2 t = (Expr t, Expr t)+  UniformTuple 3 t = (Expr t, Expr t, Expr t)+  UniformTuple 4 t = (Expr t, Expr t, Expr t, Expr t)+  UniformTuple 5 t = (Expr t, Expr t, Expr t, Expr t, Expr t)+  UniformTuple 6 t = (Expr t, Expr t, Expr t, Expr t, Expr t, Expr t)+  UniformTuple 7 t = (Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t)+  UniformTuple 8 t = (Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t)+  UniformTuple 9 t = (Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t)+  UniformTuple 10 t = (Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t, Expr t)++type UniformTupleProperties n a =+  ( IsTuple (FromTuple (UniformTuple n a)) (UniformTuple n a)+  , All ((~) (Expr a)) (FromTuple (UniformTuple n a))+  )++type HasIndexType n = (KnownNat n, n <= 10)++proveUniformTupleProperties :: forall n a. (KnownNat n, n <= 10) :- UniformTupleProperties n a+proveUniformTupleProperties = Sub $+  case fromIntegral (natVal (Proxy @n)) :: Int of+    0 -> unsafeCoerce $ Dict @(UniformTupleProperties 0 a)+    1 -> unsafeCoerce $ Dict @(UniformTupleProperties 1 a)+    2 -> unsafeCoerce $ Dict @(UniformTupleProperties 2 a)+    3 -> unsafeCoerce $ Dict @(UniformTupleProperties 3 a)+    4 -> unsafeCoerce $ Dict @(UniformTupleProperties 4 a)+    5 -> unsafeCoerce $ Dict @(UniformTupleProperties 5 a)+    6 -> unsafeCoerce $ Dict @(UniformTupleProperties 6 a)+    7 -> unsafeCoerce $ Dict @(UniformTupleProperties 7 a)+    8 -> unsafeCoerce $ Dict @(UniformTupleProperties 8 a)+    9 -> unsafeCoerce $ Dict @(UniformTupleProperties 9 a)+    10 -> unsafeCoerce $ Dict @(UniformTupleProperties 10 a)+    _ -> error "cannot happen"+{-# NOINLINE proveUniformTupleProperties #-}++class CanPeekUniformTuple n where+  peekUniformTupleImpl :: CxxConstructible b => (Ptr b -> IO (Expr a)) -> Ptr b -> IO (UniformTuple n a)++instance CanPeekUniformTuple 0 where+  peekUniformTupleImpl _ _ = pure ()++instance CanPeekUniformTuple 1 where+  peekUniformTupleImpl f = f++instance CanPeekUniformTuple 2 where+  peekUniformTupleImpl f (p :: Ptr b) = (,) <$> f p <*> f (p `plusPtr` cxxSizeOf @b)++instance CanPeekUniformTuple 3 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,) <$> f p <*> f (p `plusPtr` cxxSizeOf @b) <*> f (p `plusPtr` (2 * cxxSizeOf @b))++instance CanPeekUniformTuple 4 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))++instance CanPeekUniformTuple 5 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))++instance CanPeekUniformTuple 6 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))+      <*> f (p `plusPtr` (5 * cxxSizeOf @b))++instance CanPeekUniformTuple 7 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))+      <*> f (p `plusPtr` (5 * cxxSizeOf @b))+      <*> f (p `plusPtr` (6 * cxxSizeOf @b))++instance CanPeekUniformTuple 8 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))+      <*> f (p `plusPtr` (5 * cxxSizeOf @b))+      <*> f (p `plusPtr` (6 * cxxSizeOf @b))+      <*> f (p `plusPtr` (7 * cxxSizeOf @b))++instance CanPeekUniformTuple 9 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))+      <*> f (p `plusPtr` (5 * cxxSizeOf @b))+      <*> f (p `plusPtr` (6 * cxxSizeOf @b))+      <*> f (p `plusPtr` (7 * cxxSizeOf @b))+      <*> f (p `plusPtr` (8 * cxxSizeOf @b))++instance CanPeekUniformTuple 10 where+  peekUniformTupleImpl f (p :: Ptr b) =+    (,,,,,,,,,)+      <$> f (p `plusPtr` (0 * cxxSizeOf @b))+      <*> f (p `plusPtr` (1 * cxxSizeOf @b))+      <*> f (p `plusPtr` (2 * cxxSizeOf @b))+      <*> f (p `plusPtr` (3 * cxxSizeOf @b))+      <*> f (p `plusPtr` (4 * cxxSizeOf @b))+      <*> f (p `plusPtr` (5 * cxxSizeOf @b))+      <*> f (p `plusPtr` (6 * cxxSizeOf @b))+      <*> f (p `plusPtr` (7 * cxxSizeOf @b))+      <*> f (p `plusPtr` (8 * cxxSizeOf @b))+      <*> f (p `plusPtr` (9 * cxxSizeOf @b))++peekUniformTuple+  :: forall n b a+   . (KnownNat n, n <= 10, CxxConstructible b)+  => (Ptr b -> IO (Expr a))+  -> Ptr b+  -> IO (UniformTuple n a)+peekUniformTuple f p+  | Just Refl <- sameNat (Proxy @0) (Proxy @n) = peekUniformTupleImpl @0 f p+  | Just Refl <- sameNat (Proxy @1) (Proxy @n) = peekUniformTupleImpl @1 f p+  | Just Refl <- sameNat (Proxy @2) (Proxy @n) = peekUniformTupleImpl @2 f p+  | Just Refl <- sameNat (Proxy @3) (Proxy @n) = peekUniformTupleImpl @3 f p+  | Just Refl <- sameNat (Proxy @4) (Proxy @n) = peekUniformTupleImpl @4 f p+  | Just Refl <- sameNat (Proxy @5) (Proxy @n) = peekUniformTupleImpl @5 f p+  | Just Refl <- sameNat (Proxy @6) (Proxy @n) = peekUniformTupleImpl @6 f p+  | Just Refl <- sameNat (Proxy @7) (Proxy @n) = peekUniformTupleImpl @7 f p+  | Just Refl <- sameNat (Proxy @8) (Proxy @n) = peekUniformTupleImpl @8 f p+  | Just Refl <- sameNat (Proxy @9) (Proxy @n) = peekUniformTupleImpl @9 f p+  | Just Refl <- sameNat (Proxy @10) (Proxy @n) = peekUniformTupleImpl @10 f p+  | otherwise = error "cannot happen"++type IndexType n = UniformTuple n Int32++type IndexTypeProperties n =+  ( IsTuple (FromTuple (IndexType n)) (IndexType n)+  , All ((~) (Expr Int32)) (FromTuple (IndexType n))+  )++proveIndexTypeProperties :: forall n. (KnownNat n, n <= 10) :- IndexTypeProperties n+proveIndexTypeProperties = Sub $+  case proveUniformTupleProperties @n @Int32 of+    Sub Dict -> Dict++instance IsTuple '[] () where+  toTuple Nil = ()+  fromTuple () = Nil++instance IsTuple '[Expr a1] (Expr a1) where+  toTuple (a1 ::: Nil) = a1+  {-# INLINE toTuple #-}+  fromTuple a1 = a1 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2] (a1, a2) where+  toTuple (a1 ::: a2 ::: Nil) = (a1, a2)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2) = a1 ::: a2 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3] (a1, a2, a3) where+  toTuple (a1 ::: a2 ::: a3 ::: Nil) = (a1, a2, a3)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3) = a1 ::: a2 ::: a3 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4] (a1, a2, a3, a4) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: Nil) = (a1, a2, a3, a4)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4) = a1 ::: a2 ::: a3 ::: a4 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5] (a1, a2, a3, a4, a5) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: Nil) = (a1, a2, a3, a4, a5)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5, a6] (a1, a2, a3, a4, a5, a6) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: Nil) = (a1, a2, a3, a4, a5, a6)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5, a6) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5, a6, a7] (a1, a2, a3, a4, a5, a6, a7) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: Nil) = (a1, a2, a3, a4, a5, a6, a7)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5, a6, a7) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5, a6, a7, a8] (a1, a2, a3, a4, a5, a6, a7, a8) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: Nil) = (a1, a2, a3, a4, a5, a6, a7, a8)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5, a6, a7, a8) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5, a6, a7, a8, a9] (a1, a2, a3, a4, a5, a6, a7, a8, a9) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: a9 ::: Nil) = (a1, a2, a3, a4, a5, a6, a7, a8, a9)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5, a6, a7, a8, a9) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: a9 ::: Nil+  {-# INLINE fromTuple #-}++instance IsTuple '[a1, a2, a3, a4, a5, a6, a7, a8, a9, a10] (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10) where+  toTuple (a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: a9 ::: a10 ::: Nil) = (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10)+  {-# INLINE toTuple #-}+  fromTuple (a1, a2, a3, a4, a5, a6, a7, a8, a9, a10) = a1 ::: a2 ::: a3 ::: a4 ::: a5 ::: a6 ::: a7 ::: a8 ::: a9 ::: a10 ::: Nil+  {-# INLINE fromTuple #-}++-- instance IsTuple (Arguments '[Expr a]) (Expr a) where+--   toTuple (x ::: Nil) = x+--   fromTuple () = Nil+ -- | A v'Var'. type Var = Expr Int32 @@ -139,6 +447,13 @@ -- | Either v'Var' or v'RVar'. type VarOrRVar = Expr Int32 +-- | A multi-dimensional box -- cartesian product of the 'Range's.+newtype Region = Region [Range]+  deriving stock (Show)++-- | An @n@-dimensional reduction domain.+newtype ReductionDomain (n :: Nat) = ReductionDomain (ForeignPtr CxxRDom)+ -- | Create a scalar expression from a Haskell value. mkExpr :: IsHalideType a => a -> Expr a mkExpr x = unsafePerformIO $! Expr <$> toCxxExpr x@@ -149,6 +464,115 @@   [CU.exp| void {     new ($(Halide::Var* ptr)) Halide::Var{std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}} } |] +withRange :: Range -> (Ptr CxxRange -> IO a) -> IO a+withRange r action =+  asExpr r.rangeMin $ \minPtr ->+    asExpr r.rangeExtent $ \extentPtr -> do+      fp <-+        cxxConstruct $ \destPtr ->+          [CU.exp| void { new ($(Halide::Range* destPtr)) Halide::Range{+            *$(const Halide::Expr* minPtr), *$(const Halide::Expr* extentPtr)} } |]+      withForeignPtr fp action++-- | Create a reduction domain. Use 'asRVar' to cast it into an index.+--+-- For more information about reduction variables, see [@Halide::RDom@](https://halide-lang.org/docs/class_halide_1_1_r_dom.html).+mkRDom+  :: forall n+   . HasIndexType n+  => Text+  -- ^ name+  -> IndexType n+  -- ^ mins+  -> IndexType n+  -- ^ extents+  -> IO (ReductionDomain n)+  -- ^ reduction variables+mkRDom (T.encodeUtf8 -> name) mins extents = fmap ReductionDomain $+  case proveIndexTypeProperties @n of+    Sub Dict ->+      asVectorOf @((~) (Expr Int32)) asExpr (fromTuple mins) $ \mins' ->+        asVectorOf @((~) (Expr Int32)) asExpr (fromTuple extents) $ \extents' ->+          cxxConstruct $ \destPtr ->+            [CU.block| void {+              auto const& mins = *$(const std::vector<Halide::Expr>* mins');+              auto const& extents = *$(const std::vector<Halide::Expr>* extents');+              std::vector<Halide::Range> region;+              for (auto i = size_t{0}; i < mins.size(); ++i) {+                region.emplace_back(mins.at(i), extents.at(i));+              }+              new ($(Halide::RDom* destPtr)) Halide::RDom{+                region, std::string{$bs-ptr:name, static_cast<size_t>($bs-len:name)}};+            } |]++withCxxRDom :: ReductionDomain n -> (Ptr CxxRDom -> IO a) -> IO a+withCxxRDom (ReductionDomain fp) = withForeignPtr fp++-- | Cast a reduction domain into a multi-dimensional index that can be used to+-- perform multi-dimensional reductions.+toRVars :: forall n. HasIndexType n => ReductionDomain n -> IO (IndexType n)+toRVars rdom = do+  let allocate =+        withCxxRDom rdom $ \rdom' ->+          [CU.block| std::vector<Halide::RVar>* {+            auto const& rdom = *$(const Halide::RDom* rdom');+            std::vector<Halide::RVar> rvars;+            rvars.reserve(rdom.dimensions());+            for (auto i = 0; i < rdom.dimensions(); ++i) {+              rvars.push_back(rdom[i]);+            }+            return new std::vector<Halide::RVar>{std::move(rvars)};+          } |]+  bracket allocate deleteCxxVector $ \v -> do+    n <- cxxVectorSize v+    unless (n == fromIntegral (natVal (Proxy @n))) $ error "wrong vector length"+    ptr <- cxxVectorData v+    peekUniformTuple @n peekRVar ptr++setScalarEstimate+  :: IsHalideType a+  => a+  -- ^ Estimate+  -> Expr a+  -- ^ Parameter+  -> IO ()+setScalarEstimate estimate param =+  asScalarParam param $ \param' ->+    asExpr (mkExpr estimate) $ \estimate' ->+      [CU.exp| void {+        $(Halide::Internal::Parameter* param')->set_estimate(*$(const Halide::Expr* estimate'))+      } |]++--   withMany withRange ranges $ \regionPtr -> do+--     forM [0 .. n - 1] $ \i ->+--       cxxConstructExpr $ \destPtr ->+--         let srcPtr = ptr ``+-- cxxVectorToList+--   :: (CxxConstructible a, HasCxxVector a)+--   => (Ptr a -> Ptr a -> IO ())+--   -- ^ Copy constructor: dest src+--   -> Ptr (CxxVector a)+--   -> IO [a]+-- cxxVectorToList construct v = do+--   n <- cxxVectorSize v+--   ptr <- cxxVectorData v+--   forM [0 .. n - 1] $ \i ->+--     cxxConstruct $+--       construct undefined++--     forM [0 .. k] $ \i ->+--       pure ()+--     wrapCxxRVar+--       =<< [CU.exp| Halide::RVar* {+--             new Halide::RVar{static_cast<Halide::RVar>(Halide::RDom{+--               *$(const Halide::Expr* min'),+--               *$(const Halide::Expr* extent'),+--               std::string{$bs-ptr:name, static_cast<size_t>($bs-len:name)}+--               })}+--           } |]+--+--   undefined+ -- | 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).@@ -160,20 +584,23 @@   -> Expr Int32   -- ^ extent   -> IO (Expr Int32)-mkRVar name min extent =-  asExpr min $ \min' ->-    asExpr extent $ \extent' ->-      wrapCxxRVar-        =<< [CU.exp| Halide::RVar* {-              new Halide::RVar{static_cast<Halide::RVar>(Halide::RDom{-                *$(const Halide::Expr* min'),-                *$(const Halide::Expr* extent'),-                std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}-                })}-            } |]-  where-    s = T.encodeUtf8 name+mkRVar name start extent = do+  rdom <- mkRDom name start extent+  toRVars rdom +-- asExpr start $ \min' ->+--   asExpr extent $ \extent' ->+--     wrapCxxRVar+--       =<< [CU.exp| Halide::RVar* {+--             new Halide::RVar{static_cast<Halide::RVar>(Halide::RDom{+--               *$(const Halide::Expr* min'),+--               *$(const Halide::Expr* extent'),+--               std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}+--               })}+--           } |]+-- where+--   s = T.encodeUtf8 name+ -- | Return an undef value of the given type. -- -- For more information, see [@Halide::undef@](https://halide-lang.org/docs/namespace_halide.html#a9389bcacbed602df70eae94826312e03).@@ -197,13 +624,64 @@         [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.+-- | Print all expressions to stdout when the result is evaluates. The first expression is returned. -- -- 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))} } |]+--+-- This function is similar to 'Text.Printf.printf' in that it accepts a variable number of arguments,+-- i.e the following is valid:+--+-- @+-- let x :: Expr Float+--     x = 1+--  in printed (sin x) ("<- sin(" :: Text) x (")" :: Text)+-- @+--+-- @:: Text@ specifications are only needed if you have the @OverloadedStrings@ extension enabled.+--+-- Arguments to @printed@ can be @'Expr' a@, 'String', or 'Text'.+printed :: forall a t. (IsHalideType a, PrintedType t (Expr a)) => Expr a -> t+printed = printedWhen @a @t (mkExpr True) +printedWhen :: forall a t. (IsHalideType a, PrintedType t (Expr a)) => Expr Bool -> Expr a -> t+printedWhen cond x = unsafePerformIO $ do+  v <- newCxxVector Nothing+  appendToPrintArgs v x+  pure $ printedWhenImpl @t @(Expr a) cond v++class PrintedType t r where+  printedWhenImpl :: Expr Bool -> Ptr (CxxVector CxxExpr) -> t++instance (IsHalideType a, r ~ Expr a) => PrintedType (Expr a) r where+  printedWhenImpl cond v = unsafePerformIO $+    asExpr cond $ \cond' ->+      cxxConstructExpr $ \expr ->+        [CU.exp| void { new ($(Halide::Expr* expr)) Halide::Expr{Halide::print_when(+          *$(const Halide::Expr* cond'), *$(const std::vector<Halide::Expr>* v))} } |]+  {-# NOINLINE printedWhenImpl #-}++instance (PrintedArg a, PrintedType t r) => PrintedType (a -> t) r where+  printedWhenImpl cond v x = unsafePerformIO $ do+    appendToPrintArgs v x+    pure (printedWhenImpl @t @r cond v)+  {-# NOINLINE printedWhenImpl #-}++class PrintedArg a where+  appendToPrintArgs :: Ptr (CxxVector CxxExpr) -> a -> IO ()++instance IsHalideType a => PrintedArg (Expr a) where+  appendToPrintArgs v expr =+    asExpr expr $ \expr' ->+      [CU.exp| void { $(std::vector<Halide::Expr>* v)->push_back(*$(const Halide::Expr* expr')) } |]++instance PrintedArg Text where+  appendToPrintArgs v (T.encodeUtf8 -> msg) =+    [CU.exp| void { $(std::vector<Halide::Expr>* v)->emplace_back(+      std::string{$bs-ptr:msg, static_cast<size_t>($bs-len:msg)}) } |]++instance PrintedArg String where+  appendToPrintArgs v (pack -> msg) = appendToPrintArgs v msg+ infix 4 `eq`, `neq`, `lt`, `lte`, `gt`, `gte`  -- | '==' but lifted to return an 'Expr'.@@ -242,10 +720,52 @@   [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+-- | '&&' but lifted to return an 'Expr'.+and :: Expr Bool -> Expr Bool -> Expr Bool+and = 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'.+or :: Expr Bool -> Expr Bool -> Expr Bool+or = binaryOp $ \a b ptr ->+  [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+    (*$(Halide::Expr* a)) || (*$(Halide::Expr* b))} } |]++-- | 'Prelude.min' but lifted to return an 'Expr'.+min :: IsHalideType a => Expr a -> Expr a -> Expr a+min = binaryOp $ \a b ptr ->+  [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+    Halide::min(*$(Halide::Expr* a), *$(Halide::Expr* b))} } |]++-- | 'Prelude.max' but lifted to return an 'Expr'.+max :: IsHalideType a => Expr a -> Expr a -> Expr a+max = binaryOp $ \a b ptr ->+  [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+    Halide::max(*$(Halide::Expr* a), *$(Halide::Expr* b))} } |]++-- | Divide two integers, rounding towards zero.+div :: forall a. (IsHalideType a, Integral a) => Expr a -> Expr a -> Expr a+div = binaryOp $ \a b ptr ->+  [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+    Halide::div_round_to_zero(*$(Halide::Expr* a), *$(Halide::Expr* b))} } |]+  where+    _ = keepRedundantConstraint (Proxy @(Integral a))++-- | Compute the remainder of dividing two integers, when division is rounding toward zero.+mod :: forall a. (IsHalideType a, Integral a) => Expr a -> Expr a -> Expr a+mod = binaryOp $ \a b ptr ->+  [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+    Halide::mod_round_to_zero(*$(Halide::Expr* a), *$(Halide::Expr* b))} } |]+  where+    _ = keepRedundantConstraint (Proxy @(Integral a))++-- | 'ifThenElse cond a b' is the analogue of @if cond then a else b@, but -- lifted to work with 'Expr' types.-bool :: IsHalideType a => Expr Bool -> Expr a -> Expr a -> Expr a-bool condExpr trueExpr falseExpr = unsafePerformIO $+--+-- See also the [RebindableSyntax](https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/rebindable_syntax.html#extension-RebindableSyntax) extension.+ifThenElse :: IsHalideType a => Expr Bool -> Expr a -> Expr a -> Expr a+ifThenElse condExpr trueExpr falseExpr = unsafePerformIO $   asExpr condExpr $ \p ->     asExpr trueExpr $ \t ->       asExpr falseExpr $ \f ->@@ -292,19 +812,37 @@       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 (Expr expr) =+    unpack . unsafePerformIO $! do+      withForeignPtr expr $ \x ->+        peekAndDeleteCxxString+          =<< [CU.block| std::string* {++              // std::cerr << "Calling to_string_via_iostream on "+              //           << static_cast<void const *>($(const Halide::Expr* x))+              //           << ": '"+              //           << Halide::Expr{123}+              //           << "'"+              //           << std::endl;+              std::cerr << "Calling IRPrinter::test() ...\n";+              Halide::Internal::IRPrinter::test();+              std::cerr << "1)\n";+              auto x = *$(const Halide::Expr* x);+              std::cerr << "2)\n";+              std::cerr << x.defined() << " " << x.type() << "'" << x << "'" << std::endl;+              std::cerr << "3)\n";++              auto expr = Halide::Expr{123};+              std::cerr << "Returning new std::string: '" << expr << "' ...\n";+              return to_string_via_iostream(expr);+            } |]   show (Var var) = unpack . unsafePerformIO $ do     withForeignPtr var $ \x ->       peekAndDeleteCxxString-        =<< [CU.exp| std::string* { to_string_via_iostream(*$(const Halide::Var* x)) } |]+        =<< [CU.block| std::string* {+              return to_string_via_iostream(*$(const Halide::Var* x));+            } |]   show (RVar rvar) = unpack . unsafePerformIO $ do     withForeignPtr rvar $ \x ->       peekAndDeleteCxxString@@ -396,6 +934,8 @@   atanh :: Expr a -> Expr a   atanh = unaryOp $ \a ptr -> [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{Halide::atanh(*$(Halide::Expr* a))} } |] +deriving stock instance Show Range+ -- | Wrap a raw @Halide::Expr@ pointer in a Haskell value. -- -- __Note:__ This function checks the runtime type of the expression.@@ -598,6 +1138,11 @@     go [] v = action v     go (x : xs) v = asPtr x $ \p -> cxxVectorPushBack v p >> go xs v +peekRVar :: Ptr CxxRVar -> IO RVar+peekRVar p =+  wrapCxxRVar+    =<< [CU.exp| Halide::RVar* { new Halide::RVar{*$(const Halide::RVar* p)} } |]+ -- | Use the underlying @Halide::Var@ in an 'IO' action. asVar :: HasCallStack => Expr Int32 -> (Ptr CxxVar -> IO b) -> IO b asVar (Var fp) = withForeignPtr fp@@ -655,3 +1200,12 @@   asExpr a $ \aPtr -> asExpr b $ \bPtr ->     cxxConstructExpr $ \destPtr ->       f aPtr bPtr destPtr++testWriteToStderr :: IO ()+testWriteToStderr = do+  [CU.block| void {+    Halide::Expr expr{123};+    std::ostringstream out;+    out << expr;+    std::cerr << "Output: '" << out.str() << "'" << std::endl;+  } |]

@@ -9,6 +9,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilyDependencies #-} {-# OPTIONS_GHC -Wno-redundant-constraints #-}  -- |@@ -20,10 +21,13 @@     Func (..)   , FuncTy (..)   , Stage (..)+  , Function+  , Parameter   , buffer   , scalar   , define   , (!)+  , realizeOnTarget   , realize      -- * Scheduling@@ -56,9 +60,9 @@   , prettyLoopNest      -- * Internal-  , IndexTuple   , asBufferParam   , withFunc+  , withCxxFunc   , withBufferParam   , wrapCxxFunc   , CxxStage@@ -69,6 +73,8 @@  import Control.Exception (bracket) import Control.Monad (forM)+import Data.Constraint+import Data.Functor ((<&>)) import Data.IORef import Data.Kind (Type) import Data.Proxy@@ -91,6 +97,7 @@ import Language.Halide.Type import Language.Halide.Utils import System.IO.Unsafe (unsafePerformIO)+import Unsafe.Coerce import Prelude hiding (min, tail)  -- | Haskell counterpart of [Halide::Stage](https://halide-lang.org/docs/class_halide_1_1_stage.html).@@ -101,17 +108,27 @@ -- | A function in Halide. Conceptually, it can be thought of as a lazy -- @n@-dimensional buffer of type @a@. --+-- Here, @a@ is most often @'Expr' t@ for a type @t@ that is an instance of 'IsHalideType'.+-- However, one can also define @Func@s that return multiple values. In this case, @a@ will+-- be a tuple of 'Expr's.+-- -- 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+  Param :: IsHalideType a => {-# UNPACK #-} !(IORef (Maybe (ForeignPtr CxxImageParam))) -> Func 'ParamTy n (Expr 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) +-- | Synonym for the most commonly used function type.+type Function n a = Func 'FuncTy n (Expr a)++-- | Synonym for the most commonly used parameter type.+type Parameter n a = Func 'ParamTy n (Expr a)+ -- | A single definition of a t'Func'. newtype Stage (n :: Nat) (a :: Type) = Stage (ForeignPtr CxxStage) @@ -194,13 +211,8 @@     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+class KnownNat n => Schedulable f (n :: Nat) (a :: Type) where   -- | Vectorize the dimension.   vectorize :: VarOrRVar -> f n a -> IO (f n a) @@ -208,7 +220,15 @@   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)+  --+  -- Note that @variables@ should only contain variables that belong to the function.+  -- If this is not the case, a runtime error will be thrown.+  reorder+    :: [VarOrRVar]+    -- ^ variables+    -> f n a+    -- ^ function or stage+    -> 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@.@@ -229,10 +249,17 @@   -- | Mark the dimension to be traversed in parallel   parallel :: VarOrRVar -> f n a -> IO (f n a) +  -- | Issue atomic updates for this Func.+  atomic+    :: Bool+    -- ^ whether to override the associativity test+    -> 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)+  gpuBlocks :: (KnownNat k, 1 <= k, k <= 3) => DeviceAPI -> IndexType k -> f n a -> IO (f n a)+  gpuThreads :: (KnownNat k, 1 <= k, k <= 3) => DeviceAPI -> IndexType k -> 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@@ -241,7 +268,11 @@   -- 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+-- | GHC is not able to automatically prove the transitivity property for type-level naturals. We help GHC out 😀.+proveTransitivityOfLessThanEqual :: (KnownNat k, KnownNat l, KnownNat m, k <= l, l <= m) => Dict (k <= m)+proveTransitivityOfLessThanEqual = unsafeCoerce $ Dict @(1 <= 2)++instance KnownNat n => Schedulable Stage n a where   vectorize var stage = do     withCxxStage stage $ \stage' ->       asVarOrRVar var $ \var' ->@@ -321,6 +352,15 @@           });         } |]     pure stage+  atomic (fromIntegral . fromEnum -> override) stage = do+    withCxxStage stage $ \stage' ->+      [C.throwBlock| void {+        handle_halide_exceptions([=](){+          $(Halide::Stage* stage')->atomic($(bool override));+        });+      } |]+    pure stage+   specialize cond stage = do     withCxxStage stage $ \stage' ->       asExpr cond $ \cond' ->@@ -339,46 +379,54 @@             std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)});         });       } |]-  gpuBlocks (fromIntegral . fromEnum -> api) vars stage = do-    withCxxStage stage $ \stage' ->-      asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do-        [C.throwBlock| void {-          handle_halide_exceptions([=](){-            auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');-            auto& stage = *$(Halide::Stage* stage');-            auto const device = static_cast<Halide::DeviceAPI>($(int api));-            switch (vars.size()) {-              case 1: stage.gpu_blocks(vars.at(0), device);-                      break;-              case 2: stage.gpu_blocks(vars.at(0), vars.at(1), device);-                      break;-              case 3: stage.gpu_blocks(vars.at(0), vars.at(1), vars.at(2), device);-                      break;-              default: throw std::runtime_error{"unexpected number of arguments in gpuBlocks"};-            }-          });-        } |]-    pure stage-  gpuThreads (fromIntegral . fromEnum -> api) vars stage = do-    withCxxStage stage $ \stage' ->-      asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do-        [C.throwBlock| void {-          handle_halide_exceptions([=](){-            auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');-            auto& stage = *$(Halide::Stage* stage');-            auto const device = static_cast<Halide::DeviceAPI>($(int api));-            switch (vars.size()) {-              case 1: stage.gpu_threads(vars.at(0), device);-                      break;-              case 2: stage.gpu_threads(vars.at(0), vars.at(1), device);-                      break;-              case 3: stage.gpu_threads(vars.at(0), vars.at(1), vars.at(2), device);-                      break;-              default: throw std::runtime_error{"unexpected number of arguments in gpuThreads"};-            }-          });-        } |]-    pure stage+  gpuBlocks :: forall k. (KnownNat k, 1 <= k, k <= 3) => DeviceAPI -> IndexType k -> Stage n a -> IO (Stage n a)+  gpuBlocks (fromIntegral . fromEnum -> api :: C.CInt) vars stage =+    case proveTransitivityOfLessThanEqual @k @3 @10 of+      Dict -> case proveIndexTypeProperties @k of+        Sub Dict ->+          withCxxStage stage $ \stage' ->+            asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do+              [C.throwBlock| void {+                handle_halide_exceptions([=](){+                  auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');+                  auto& stage = *$(Halide::Stage* stage');+                  auto const device = static_cast<Halide::DeviceAPI>($(int api));+                  switch (vars.size()) {+                    case 1: stage.gpu_blocks(vars.at(0), device);+                            break;+                    case 2: stage.gpu_blocks(vars.at(0), vars.at(1), device);+                            break;+                    case 3: stage.gpu_blocks(vars.at(0), vars.at(1), vars.at(2), device);+                            break;+                    default: throw std::runtime_error{"unexpected number of arguments in gpuBlocks"};+                  }+                });+              } |]+              pure stage+  gpuThreads :: forall k. (KnownNat k, 1 <= k, k <= 3) => DeviceAPI -> IndexType k -> Stage n a -> IO (Stage n a)+  gpuThreads (fromIntegral . fromEnum -> api :: C.CInt) vars stage =+    case proveTransitivityOfLessThanEqual @k @3 @10 of+      Dict -> case proveIndexTypeProperties @k of+        Sub Dict -> do+          withCxxStage stage $ \stage' ->+            asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple vars) $ \vars' -> do+              [C.throwBlock| void {+                handle_halide_exceptions([=](){+                  auto const& vars = *$(const std::vector<Halide::VarOrRVar>* vars');+                  auto& stage = *$(Halide::Stage* stage');+                  auto const device = static_cast<Halide::DeviceAPI>($(int api));+                  switch (vars.size()) {+                    case 1: stage.gpu_threads(vars.at(0), device);+                            break;+                    case 2: stage.gpu_threads(vars.at(0), vars.at(1), device);+                            break;+                    case 3: stage.gpu_threads(vars.at(0), vars.at(1), vars.at(2), device);+                            break;+                    default: throw std::runtime_error{"unexpected number of arguments in gpuThreads"};+                  }+                });+              } |]+          pure stage   gpuLanes (fromIntegral . fromEnum -> api) var stage = do     withCxxStage stage $ \stage' ->       asVarOrRVar var $ \var' ->@@ -402,7 +450,7 @@         } |]  viaStage1-  :: (KnownNat n, IsHalideType b)+  :: KnownNat n   => (a -> Stage n b -> IO (Stage n b))   -> a   -> Func t n b@@ -412,7 +460,7 @@   pure func  viaStage2-  :: (KnownNat n, IsHalideType b)+  :: (KnownNat n)   => (a1 -> a2 -> Stage n b -> IO (Stage n b))   -> a1   -> a2@@ -437,7 +485,7 @@ -}  viaStage4-  :: (KnownNat n, IsHalideType b)+  :: (KnownNat n)   => (a1 -> a2 -> a3 -> a4 -> Stage n b -> IO (Stage n b))   -> a1   -> a2@@ -449,7 +497,7 @@   _ <- f a1 a2 a3 a4 =<< getStage func   pure func -instance (KnownNat n, IsHalideType a) => Schedulable (Func t) n a where+instance KnownNat n => Schedulable (Func t) n a where   vectorize = viaStage1 vectorize   unroll = viaStage1 unroll   reorder = viaStage1 reorder@@ -457,6 +505,7 @@   fuse = viaStage2 fuse   serial = viaStage1 serial   parallel = viaStage1 parallel+  atomic = viaStage1 atomic   specialize cond func = getStage func >>= specialize cond   specializeFail msg func = getStage func >>= specializeFail msg   gpuBlocks = viaStage2 gpuBlocks@@ -484,74 +533,13 @@     | fromIntegral k == [CU.pure| int { static_cast<int>(Halide::TailStrategy::Auto) } |] = TailAuto     | otherwise = error $ "invalid TailStrategy: " <> show k --- vectorize---   :: (KnownNat n, IsHalideType a)---   => TailStrategy---   -> Func t n a---   -> Expr Int32---   -- ^ Variable to vectorize---   -> Expr Int32---   -- ^ Split factor---   -> IO ()--- vectorize strategy func var factor =---   withFunc func $ \f ->---     asVarOrRVar var $ \x ->---       asExpr factor $ \n ->---         [C.throwBlock| void {---           $(Halide::Func* f)->vectorize(*$(Halide::VarOrRVar* x), *$(Halide::Expr* n),---                                         static_cast<Halide::TailStrategy>($(int tail)));---         } |]---   where---     tail = fromIntegral (fromEnum strategy)---- | Split a dimension by the given factor, then unroll the inner dimension.------ This is how you unroll a loop of unknown size by some constant factor. After--- this call, @var@ refers to the outer dimension of the split.--- unroll---   :: (KnownNat n, IsHalideType a)---   => TailStrategy---   -> Func t n a---   -> Expr Int32---   -- ^ Variable @var@ to vectorize---   -> Expr Int32---   -- ^ Split factor---   -> IO ()--- unroll strategy func var factor =---   withFunc func $ \f ->---     asVarOrRVar var $ \x ->---       asExpr factor $ \n ->---         [C.throwBlock| void {---           $(Halide::Func* f)->unroll(*$(Halide::VarOrRVar* x), *$(Halide::Expr* n),---                                      static_cast<Halide::TailStrategy>($(int tail)));---         } |]---   where---     tail = fromIntegral (fromEnum strategy)---- | Reorder variables to have the given nesting order, from innermost out.--- reorder---   :: forall t n a i ts---    . ( IsTuple (Arguments ts) i---      , All ((~) (Expr Int32)) ts---      , Length ts ~ n---      , KnownNat n---      , IsHalideType a---      )---   => Func t n a---   -> i---   -> IO ()--- reorder func args =---   asVectorOf @((~) (Expr Int32)) asVarOrRVar (fromTuple args) $ \v -> do---     withFunc func $ \f ->---       [C.throwBlock| void { $(Halide::Func* f)->reorder(*$(std::vector<Halide::VarOrRVar>* v)); } |]- -- | Statically declare the range over which the function will be evaluated in the general case. -- -- This provides a basis for the auto scheduler to make trade-offs and scheduling decisions. -- The auto generated schedules might break when the sizes of the dimensions are very different from the -- estimates specified. These estimates are used only by the auto scheduler if the function is a pipeline output. estimate-  :: (KnownNat n, IsHalideType a)+  :: KnownNat n   => Expr Int32   -- ^ index variable   -> Expr Int32@@ -560,11 +548,14 @@   -- ^ @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)) } |]+estimate var start extent func =+  withFunc func $ \f ->+    asVar var $ \i ->+      asExpr start $ \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. --@@ -573,7 +564,7 @@ -- 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)+  :: KnownNat n   => Expr Int32   -- ^ index variable   -> Expr Int32@@ -582,16 +573,19 @@   -- ^ @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)) } |]+bound var start extent func =+  withFunc func $ \f ->+    asVar var $ \i ->+      asExpr start $ \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 :: KnownNat n => Func t n a -> IO [Var] getArgs func =   withFunc func $ \func' -> do     let allocate =@@ -608,7 +602,7 @@ -- | 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 :: KnownNat n => 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(); }); } |]@@ -621,7 +615,7 @@ -- -- 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)+  :: (KnownNat n, KnownNat m)   => Func t1 n a   -> Func 'FuncTy m b   -> IO (Func 'FuncTy n a)@@ -635,7 +629,7 @@ -- -- 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 :: KnownNat n => Func t n a -> IO (Func 'FuncTy n a) asUsed f =   withFunc f $ \fPtr ->     wrapCxxFunc@@ -647,7 +641,7 @@ -- 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 :: KnownNat n => DeviceAPI -> Func t n a -> IO (Func t n a) copyToDevice deviceApi func = do   withFunc func $ \f ->     [C.throwBlock| void {@@ -660,95 +654,9 @@     api = fromIntegral . fromEnum $ deviceApi  -- | Same as @'copyToDevice' 'DeviceHost'@-copyToHost :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+copyToHost :: KnownNat n => Func t n a -> IO (Func t n a) copyToHost = copyToDevice DeviceHost --- | Split a dimension into inner and outer subdimensions with the given names, where the inner dimension--- iterates from @0@ to @factor-1@.------ The inner and outer subdimensions can then be dealt with using the other scheduling calls. It's okay--- to reuse the old variable name as either the inner or outer variable. The first argument specifies--- how the tail should be handled if the split factor does not provably divide the extent.--- split---   :: (KnownNat n, IsHalideType a)---   => TailStrategy---   -- ^ how to treat the remainder---   -> Func t n a---   -> Expr Int32---   -- ^ loop variable to split---   -> Expr Int32---   -- ^ new outer loop variable---   -> Expr Int32---   -- ^ new inner loop variable---   -> Expr Int32---   -- ^ split factor---   -> IO (Func t n a)--- split tail func old outer inner factor = do---   withFunc func $ \f ->---     asVarOrRVar old $ \old' ->---       asVarOrRVar outer $ \outer' ->---         asVarOrRVar inner $ \inner' ->---           asExpr factor $ \factor' ->---             [C.throwBlock| void {---               handle_halide_exceptions([=](){---                 $(Halide::Func* f)->split(---                   *$(const Halide::VarOrRVar* old'),---                   *$(const Halide::VarOrRVar* outer'),---                   *$(const Halide::VarOrRVar* inner'),---                   *$(const Halide::Expr* factor'),---                   static_cast<Halide::TailStrategy>($(int t)));---               }); } |]---   pure func---   where---     t = fromIntegral . fromEnum $ tail---- | Join two dimensions into a single fused dimenion.------ The fused dimension covers the product of the extents of the inner and outer dimensions given.--- fuse---   :: (KnownNat n, IsHalideType a)---   => Func t n a---   -> Expr Int32---   -- ^ inner loop variable---   -> Expr Int32---   -- ^ outer loop variable---   -> Expr Int32---   -- ^ new fused loop variable---   -> IO (Func t n a)--- fuse func outer inner fused = do---   withFunc func $ \f ->---     asVarOrRVar outer $ \outer' ->---       asVarOrRVar inner $ \inner' ->---         asVarOrRVar fused $ \fused' ->---           [CU.exp| void {---                 $(Halide::Func* f)->fuse(---                   *$(const Halide::VarOrRVar* outer'),---                   *$(const Halide::VarOrRVar* inner'),---                   *$(const Halide::VarOrRVar* fused')) } |]---   pure func---- withVarOrRVarMany :: [Expr Int32] -> (Int -> Ptr (CxxVector CxxVarOrRVar) -> IO a) -> IO a--- withVarOrRVarMany xs f =---   bracket allocate destroy $ \v -> do---     let go !k [] = f k v---         go !k (y : ys) = withVarOrRVarMany y $ \p -> do---           [CU.exp| void { $(std::vector<Halide::Expr>* v)->push_back(*$(Halide::VarOrRVar* p)) } |]---           go (k + 1) ys---     go 0 xs---   where---     count = fromIntegral (length xs)----   withFunc func $ \f ->---     withVarOrRVarMany vars $ \count v -> do---       unless natVal (Proxy @n)---       handleHalideExceptionsM---         [C.tryBlock| void {---           $(Halide::Func* f)->reorder(*$(std::vector<Halide::VarOrRVar>* v));---         } |]------ class Curry (args :: [Type]) (r :: Type) (f :: Type) | args r -> f where---   curryG :: (Arguments args -> r) -> f- mkBufferParameter   :: forall n a. (KnownNat n, IsHalideType a) => Maybe Text -> IO (ForeignPtr CxxImageParam) mkBufferParameter maybeName = do@@ -785,42 +693,62 @@ withBufferParam   :: forall n a b    . (HasCallStack, KnownNat n, IsHalideType a)-  => Func 'ParamTy n a+  => Func 'ParamTy n (Expr 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)+withFunc :: KnownNat n => Func t n a -> (Ptr CxxFunc -> IO b) -> IO b+withFunc f action = case f of+  Func fp -> withForeignPtr fp action+  p@(Param _) -> forceFunc p >>= \(Func fp) -> withForeignPtr fp action +withCxxFunc :: KnownNat n => Func 'FuncTy n a -> (Ptr CxxFunc -> IO b) -> IO b+withCxxFunc (Func fp) = withForeignPtr fp+ wrapCxxFunc :: Ptr CxxFunc -> IO (Func 'FuncTy n a) wrapCxxFunc = fmap Func . newForeignPtr deleter   where     deleter = [C.funPtr| void deleteFunc(Halide::Func *x) { delete x; } |] -forceFunc :: forall t n a. (KnownNat n, IsHalideType a) => Func t n a -> IO (Func 'FuncTy n a)-forceFunc x@(Func _) = pure x-forceFunc (Param r) = do-  fp <- getBufferParameter @n @a Nothing r-  withForeignPtr fp $ \p ->-    wrapCxxFunc-      =<< [CU.exp| Halide::Func* {-            new Halide::Func{static_cast<Halide::Func>(*$(Halide::ImageParam* p))} } |]+forceFunc :: forall t n a. KnownNat n => Func t n (Expr a) -> IO (Func 'FuncTy n (Expr a))+forceFunc = \case+  x@(Func _) -> pure x+  (Param r) -> do+    fp <- getBufferParameter @n @a Nothing r+    withForeignPtr fp $ \p ->+      wrapCxxFunc+        =<< [CU.exp| Halide::Func* {+              new Halide::Func{static_cast<Halide::Func>(*$(Halide::ImageParam* p))} } |] -funcToForeignPtr :: (KnownNat n, IsHalideType a) => Func t n a -> ForeignPtr CxxFunc-funcToForeignPtr x = unsafePerformIO $! forceFunc x >>= \(Func fp) -> pure fp+class IsFuncDefinition d where+  definitionToExprList :: d -> [ForeignPtr CxxExpr]+  exprListToDefinition :: [ForeignPtr CxxExpr] -> d +instance IsHalideType a => IsFuncDefinition (Expr a) where+  definitionToExprList = pure . exprToForeignPtr+  exprListToDefinition [x1] = unsafePerformIO $ withForeignPtr x1 (checkType @a) >> pure (Expr x1)+  exprListToDefinition _ = error "should never happen"++instance (IsHalideType a1, IsHalideType a2) => IsFuncDefinition (Expr a1, Expr a2) where+  definitionToExprList (x1, x2) = [exprToForeignPtr x1, exprToForeignPtr x2]+  exprListToDefinition [x1, x2] = unsafePerformIO $ do+    withForeignPtr x1 (checkType @a1)+    withForeignPtr x2 (checkType @a2)+    pure (Expr x1, Expr x2)+  exprListToDefinition _ = error "should never happen"++instance (IsHalideType a1, IsHalideType a2, IsHalideType a3) => IsFuncDefinition (Expr a1, Expr a2, Expr a3) where+  definitionToExprList (x1, x2, x3) = [exprToForeignPtr x1, exprToForeignPtr x2, exprToForeignPtr x3]+  exprListToDefinition [x1, x2, x3] = unsafePerformIO $ do+    withForeignPtr x1 (checkType @a1)+    withForeignPtr x2 (checkType @a2)+    withForeignPtr x3 (checkType @a3)+    pure (Expr x1, Expr x2, Expr x3)+  exprListToDefinition _ = error "should never happen"+ -- | Define a Halide function. -- -- @define "f" i e@ defines a Halide function called "f" such that @f[i] = e@.@@ -832,81 +760,105 @@ -- >>> 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-     )+  :: forall n d+   . (HasIndexType n, IsFuncDefinition d)   => Text-  -> i-  -> Expr a-  -> IO (Func 'FuncTy n a)-define name args expr =-  asVectorOf @((~) (Expr Int32)) asVar (fromTuple args) $ \x -> do-    let s = T.encodeUtf8 name-    asExpr expr $ \y ->-      wrapCxxFunc-        =<< [CU.block| Halide::Func* {-              Halide::Func f{std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}};-              f(*$(std::vector<Halide::Var>* x)) = *$(Halide::Expr* y);-              return new Halide::Func{f};-            } |]+  -> IndexType n+  -> d+  -> IO (Func 'FuncTy n d)+define name args definition =+  case proveIndexTypeProperties @n of+    Sub Dict -> asVectorOf @((~) (Expr Int32)) asVar (fromTuple args) $ \x -> do+      let s = T.encodeUtf8 name+      withMany withForeignPtr (definitionToExprList definition) $ \v ->+        wrapCxxFunc+          =<< [CU.block| Halide::Func* {+                Halide::Func f{std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)}};+                auto const& args = *$(const std::vector<Halide::Var>* x);+                auto const& def = *$(const std::vector<Halide::Expr>* v);+                if (def.size() == 1) {+                  f(args) = def.at(0);+                }+                else {+                  f(args) = Halide::Tuple{def};+                }+                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+  :: forall n d+   . (HasIndexType n, IsFuncDefinition d)+  => Func 'FuncTy n d+  -> IndexType n+  -> d   -> IO ()-update func args expr =+update func args definition =   withFunc func $ \f ->-    asVectorOf @((~) (Expr Int32)) asExpr (fromTuple args) $ \x ->-      asExpr expr $ \y ->-        [C.throwBlock| void {-          handle_halide_exceptions([=](){-            $(Halide::Func* f)->operator()(*$(std::vector<Halide::Expr>* x)) = *$(Halide::Expr* y);-          });-        } |]+    case proveIndexTypeProperties @n of+      Sub Dict -> asVectorOf @((~) (Expr Int32)) asExpr (fromTuple args) $ \index ->+        withMany withForeignPtr (definitionToExprList definition) $ \value ->+          [C.throwBlock| void {+            handle_halide_exceptions([=](){+              auto& f = *$(Halide::Func* f);+              auto const& index = *$(const std::vector<Halide::Expr>* index);+              auto const& value = *$(const std::vector<Halide::Expr>* value);+              if (value.size() == 1) {+                f(index) = value.at(0);+              }+              else {+                f(index) = Halide::Tuple{value};+              }+            });+          } |]  infix 9 ! +withExprIndices :: forall n a. HasIndexType n => IndexType n -> (Ptr (CxxVector CxxExpr) -> IO a) -> IO a+withExprIndices indices action =+  case proveIndexTypeProperties @n of+    Sub Dict -> asVectorOf @((~) (Expr Int32)) asExpr (fromTuple indices) $ \x ->+      action x++indexFunc :: forall n a t. HasIndexType n => Func t n a -> IndexType n -> IO [ForeignPtr CxxExpr]+indexFunc func indices = withExprIndices indices $ \x ->+  withFunc func $ \f -> do+    let allocate =+          [CU.block| std::vector<Halide::Expr>* {+            Halide::FuncRef ref = $(Halide::Func* f)->operator()(*$(std::vector<Halide::Expr>* x));+            std::vector<Halide::Expr> v;+            if (ref.size() == 1) {+              v.push_back(static_cast<Halide::Expr>(ref));+            }+            else {+              for (auto i = size_t{0}; i < ref.size(); ++i) {+                v.push_back(ref[i]);+              }+            }+            return new std::vector<Halide::Expr>{std::move(v)};+          } |]+    bracket allocate deleteCxxVector $ \v -> do+      size <- fromIntegral <$> cxxVectorSize v+      forM [0 .. size - 1] $ \i ->+        cxxConstruct $ \ptr ->+          [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{+                $(const std::vector<Halide::Expr>* v)->at($(size_t i))} } |]+ -- | Apply a Halide function. Conceptually, @f ! i@ is equivalent to @f[i]@, i.e. -- indexing into a lazy array.-(!)-  :: ( IsTuple (Arguments ts) i-     , All ((~) (Expr Int32)) ts-     , Length ts ~ n-     , KnownNat n-     , IsHalideType a-     )-  => Func t n a-  -> i-  -> Expr a-(!) func args =-  unsafePerformIO $-    withFunc func $ \f ->-      asVectorOf @((~) (Expr Int32)) asExpr (fromTuple args) $ \x ->-        cxxConstructExpr $ \ptr ->-          [CU.exp| void { new ($(Halide::Expr* ptr)) Halide::Expr{-            $(Halide::Func* f)->operator()(*$(std::vector<Halide::Expr>* x))} } |]+(!) :: (HasIndexType n, IsFuncDefinition a) => Func t n a -> IndexType n -> a+(!) func args = unsafePerformIO $ indexFunc func args <&> exprListToDefinition  -- | Get a particular dimension of a pipeline parameter. dim   :: forall n a-   . (HasCallStack, KnownNat n, IsHalideType a)+   . (HasCallStack, KnownNat n)   => Int-  -> Func 'ParamTy n a+  -> Func 'ParamTy n (Expr a)   -> IO Dimension-dim k func+dim k func@(Param _)   | 0 <= k && k < fromIntegral (natVal (Proxy @n)) =       let n = fromIntegral k        in withBufferParam func $ \f ->@@ -921,23 +873,13 @@           <> 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 :: KnownNat n => Func t n r -> IO Text prettyLoopNest func = withFunc func $ \f ->   peekAndDeleteCxxString     =<< [C.throwBlock| std::string* {@@ -947,11 +889,11 @@           });         } |] --- | Evaluate this function over a rectangular domain.+-- | Similar to 'realizeOnTarget' except that the pipeline is run on 'hostTarget'. realize   :: forall n a t b    . (KnownNat n, IsHalideType a)-  => Func t n a+  => Func t n (Expr a)   -- ^ Function to evaluate   -> [Int]   -- ^ Domain over which to evaluate@@ -959,37 +901,35 @@   -- ^ 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+realize = realizeOnTarget hostTarget --- \| 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+-- | Evaluate this function over a rectangular domain.+--+-- If your target is a GPU, this function will not automatically copy data back from the GPU.+realizeOnTarget+  :: forall n a t b+   . (KnownNat n, IsHalideType a)+  => Target+  -- ^ Target on which to run the pipeline+  -> Func t n (Expr 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+realizeOnTarget target func shape action =+  withFunc func $ \func' ->+    withCxxTarget target $ \target' ->+      allocaBuffer target shape $ \buf -> do+        let raw = castPtr buf+        [C.throwBlock| void {+          handle_halide_exceptions([=](){+            $(Halide::Func* func')->realize($(halide_buffer_t* raw), *$(const Halide::Target* target'));+          });+        } |]+        action buf  -- | A view pattern to specify the name of a buffer argument. --@@ -1008,7 +948,7 @@ --   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 :: forall n a. (KnownNat n, IsHalideType a) => Text -> Func 'ParamTy n (Expr a) -> Func 'ParamTy n (Expr a) buffer name p@(Param r) = unsafePerformIO $ do   _ <- getBufferParameter @n @a (Just name) r   pure p@@ -1032,29 +972,29 @@   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 :: 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 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 :: KnownNat n => Func t n a -> IO (Stage n a) getStage func =   withFunc func $ \func' ->     [CU.exp| Halide::Stage* { new Halide::Stage{static_cast<Halide::Stage>(*$(Halide::Func* func'))} } |]       >>= wrapCxxStage  -- | Return 'True' when the function has update definitions, 'False' otherwise.-hasUpdateDefinitions :: (KnownNat n, IsHalideType a) => Func t n a -> IO Bool+hasUpdateDefinitions :: KnownNat n => 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 :: KnownNat n => Int -> Func 'FuncTy n a -> IO (Stage n a) getUpdateStage k func =   withFunc func $ \func' ->     let k' = fromIntegral k@@ -1063,7 +1003,7 @@  -- | Identify the loop nest corresponding to some dimension of some function. getLoopLevelAtStage-  :: (KnownNat n, IsHalideType a)+  :: KnownNat n   => Func t n a   -> Expr Int32   -> Int@@ -1087,7 +1027,7 @@     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 :: KnownNat n => 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.@@ -1096,7 +1036,7 @@ -- 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 :: KnownNat n => Func 'FuncTy n a -> LoopLevel t -> IO (Func 'FuncTy n a) storeAt func level = do   withFunc func $ \f ->     withCxxLoopLevel level $ \l ->@@ -1106,7 +1046,7 @@ -- | 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 :: KnownNat n => Func 'FuncTy n a -> LoopLevel t -> IO (Func 'FuncTy n a) computeAt func level = do   withFunc func $ \f ->     withCxxLoopLevel level $ \l ->@@ -1129,7 +1069,7 @@    . IsHalideBuffer t n a   => t   -- ^ Object to treat as a buffer-  -> (Func 'ParamTy n a -> IO b)+  -> (Func 'ParamTy n (Expr a) -> IO b)   -- ^ What to do with the __temporary__ buffer   -> IO b asBufferParam arr action =

@@ -14,9 +14,7 @@   , compileToCallable   , compileToLoweredStmt   , StmtOutputFormat (..)-  , IsFuncBuilder-  , ReturnsFunc-  , Lowered+  , LoweredSignature   ) where @@ -35,7 +33,8 @@ import Foreign.C.Types (CUIntPtr (..)) import Foreign.ForeignPtr import Foreign.ForeignPtr.Unsafe-import Foreign.Ptr (FunPtr, Ptr, castPtr)+import Foreign.Marshal.Array (allocaArray, peekArray)+import Foreign.Ptr (Ptr, castPtr) import Foreign.Storable import GHC.TypeNats import Language.C.Inline qualified as C@@ -49,10 +48,14 @@ import Language.Halide.Target import Language.Halide.Type import System.IO.Temp (withSystemTempDirectory)-import Unsafe.Coerce (unsafeCoerce) +-- | Haskell counterpart of @Halide::Argument@.+data CxxArgument+ importHalide +instanceHasCxxVector "Halide::Argument"+ data ArgvStorage s   = ArgvStorage       {-# UNPACK #-} !(MutablePrimArray s (Ptr ()))@@ -62,24 +65,23 @@ newArgvStorage n = ArgvStorage <$> P.newPinnedPrimArray n <*> P.newPinnedPrimArray n  setArgvStorage-  :: (All ValidArgument inputs, All ValidArgument outputs)+  :: All ValidArgument args   => ArgvStorage RealWorld-  -> Arguments inputs-  -> Arguments outputs+  -> Ptr CxxUserContext+  -> Arguments args   -> IO ()-setArgvStorage (ArgvStorage argv scalarStorage) inputs outputs = do+setArgvStorage (ArgvStorage argv scalarStorage) context inputs = 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+      go :: All ValidArgument args' => Int -> Arguments args' -> IO ()+      go !_ Nil = pure ()+      go !i (x ::: 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+  go 0 (context ::: inputs)   touch argv   touch scalarStorage @@ -107,162 +109,106 @@     poke (castPtr argv :: Ptr (Ptr (HalideBuffer n a))) x   {-# INLINE fillSlot #-} -class ValidArgument (Lowered t) => ValidParameter (t :: Type) where-  appendToArgList :: Ptr (CxxVector CxxArgument) -> t -> IO ()-  prepareParameter :: IO t--instance IsHalideType a => ValidParameter (Expr a) where-  appendToArgList :: Ptr (CxxVector CxxArgument) -> Expr a -> IO ()-  appendToArgList v expr =-    asScalarParam expr $ \p ->-      [CU.exp| void { $(std::vector<Halide::Argument>* v)->emplace_back(-        $(Halide::Internal::Parameter const* p)->name(),-        Halide::Argument::InputScalar,-        $(Halide::Internal::Parameter const* p)->type(),-        $(Halide::Internal::Parameter const* p)->dimensions(),-        $(Halide::Internal::Parameter const* p)->get_argument_estimates()) } |]-  prepareParameter :: IO (Expr a)-  prepareParameter = ScalarParam <$> newIORef Nothing--instance (KnownNat n, IsHalideType a) => ValidParameter (Func t n a) where-  appendToArgList :: Ptr (CxxVector CxxArgument) -> Func t n a -> IO ()-  appendToArgList v func@(Param _) =-    withBufferParam func $ \p ->-      [CU.exp| void { $(std::vector<Halide::Argument>* v)->push_back(-        *$(Halide::ImageParam const* p)) } |]-  appendToArgList _ _ = error "appendToArgList called on Func; this should never happen"-  prepareParameter :: IO (Func t n a)-  prepareParameter = unsafeCoerce $ Param <$> newIORef Nothing--class PrepareParameters ts where-  prepareParameters :: IO (Arguments ts)--instance PrepareParameters '[] where-  prepareParameters :: IO (Arguments '[])-  prepareParameters = pure Nil--instance (ValidParameter t, PrepareParameters ts) => PrepareParameters (t ': ts) where-  prepareParameters :: IO (Arguments (t : ts))-  prepareParameters = do-    t <- prepareParameter @t-    ts <- prepareParameters @ts-    pure $ t ::: ts--prepareCxxArguments-  :: forall ts b-   . (All ValidParameter ts, KnownNat (Length ts))-  => Arguments ts-  -> (Ptr (CxxVector CxxArgument) -> IO b)-  -> IO b-prepareCxxArguments args action = do-  let count = fromIntegral (natVal (Proxy @(Length ts)))-      allocate =-        [CU.block| std::vector<Halide::Argument>* {-          auto p = new std::vector<Halide::Argument>{};-          p->reserve($(size_t count));-          return p;-        } |]-      destroy p = [CU.exp| void { delete $(std::vector<Halide::Argument>* p) } |]-  bracket allocate destroy $ \v -> do-    let go :: All ValidParameter ts' => Arguments ts' -> IO ()-        go Nil = pure ()-        go (x ::: xs) = appendToArgList v x >> go xs-    go args-    action v--deleteCxxUserContext :: FunPtr (Ptr CxxUserContext -> IO ())-deleteCxxUserContext = [C.funPtr| void deleteUserContext(Halide::JITUserContext* p) { delete p; } |]--wrapCxxUserContext :: Ptr CxxUserContext -> IO (ForeignPtr CxxUserContext)-wrapCxxUserContext = newForeignPtr deleteCxxUserContext+-- class ValidArgument (Lowered t) => ValidParameter (t :: Type) where+--   appendToArgList :: Ptr (CxxVector CxxArgument) -> t -> IO ()+--   prepareParameter :: IO t -newEmptyCxxUserContext :: IO (ForeignPtr CxxUserContext)-newEmptyCxxUserContext =-  wrapCxxUserContext =<< [CU.exp| Halide::JITUserContext* { new Halide::JITUserContext{} } |]+-- instance IsHalideType a => ValidParameter (Expr a) where+--   appendToArgList :: Ptr (CxxVector CxxArgument) -> Expr a -> IO ()+--   appendToArgList v expr =+--     asScalarParam expr $ \p ->+--       [CU.exp| void { $(std::vector<Halide::Argument>* v)->emplace_back(+--         $(Halide::Internal::Parameter const* p)->name(),+--         Halide::Argument::InputScalar,+--         $(Halide::Internal::Parameter const* p)->type(),+--         $(Halide::Internal::Parameter const* p)->dimensions(),+--         $(Halide::Internal::Parameter const* p)->get_argument_estimates()) } |]+--   prepareParameter :: IO (Expr a)+--   prepareParameter = ScalarParam <$> newIORef Nothing -wrapCxxCallable :: Ptr CxxCallable -> IO (Callable inputs outputs)-wrapCxxCallable = fmap Callable . newForeignPtr deleter-  where-    deleter = [C.funPtr| void deleteCallable(Halide::Callable* p) { delete p; } |]+-- instance (KnownNat n, IsHalideType a, t ~ 'ParamTy) => ValidParameter (Func t n (Expr a)) where+--   appendToArgList :: Ptr (CxxVector CxxArgument) -> Func 'ParamTy n (Expr a) -> IO ()+--   appendToArgList v func@(Param _) =+--     withBufferParam func $ \p ->+--       [CU.exp| void { $(std::vector<Halide::Argument>* v)->push_back(+--         *$(Halide::ImageParam const* p)) } |]+--   prepareParameter = Param <$> newIORef Nothing -type Lowered :: forall k. k -> k+class KnownNat n => FuncBuilder f (n :: Nat) (a :: Type) | f -> n a where+  buildFunc :: Ptr (CxxVector CxxArgument) -> f -> IO (Func 'FuncTy n a) --- | 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)+instance (k ~ 'ParamTy, KnownNat m, IsHalideType t, FuncBuilder r n a) => FuncBuilder (Func k m (Expr t) -> r) n a where+  buildFunc v f = do+    param <- Param <$> newIORef Nothing+    func <- buildFunc v (f param)+    withBufferParam param $ \p ->+      [CU.exp| void { $(std::vector<Halide::Argument>* v)->push_back(*$(Halide::ImageParam const* p)) } |]+    pure func --- | 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 (IsHalideType t, FuncBuilder r n a) => FuncBuilder (Expr t -> r) n a where+  buildFunc v f = do+    param <- ScalarParam <$> newIORef Nothing+    func <- buildFunc v (f param)+    asScalarParam param $ \p ->+      [CU.block| void {+        auto const& p = *$(Halide::Internal::Parameter const* p);+        $(std::vector<Halide::Argument>* v)->emplace_back(+          p.name(),+          Halide::Argument::InputScalar,+          p.type(),+          p.dimensions(),+          p.get_argument_estimates());+      } |]+    pure func -instance (FunctionReturn f ~ IO (Func t n a), IsHalideType a, KnownNat n) => ReturnsFunc f t n a+instance (KnownNat n, t ~ 'FuncTy, n' ~ n, a' ~ a) => FuncBuilder (IO (Func t n' a')) n a where+  buildFunc _ action = action -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)))-  )+type family LoweredSignature f where+  LoweredSignature (Expr a -> r) = a -> LoweredSignature r+  LoweredSignature (Func t n (Expr a) -> r) = Ptr (HalideBuffer n a) -> LoweredSignature r+  LoweredSignature (IO (Func t n (Expr a))) = Ptr (HalideBuffer n a) -> IO ()+  LoweredSignature (IO (Func t n (Expr a1, Expr a2))) = Ptr (HalideBuffer n a1) -> Ptr (HalideBuffer n a2) -> IO () -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)+type IsHalideKernel f = (KnownNat (Length (FunctionArguments f)), All ValidArgument (FunctionArguments f), Curry (FunctionArguments f) (IO ()) f) -newtype Callable (inputs :: [Type]) (output :: Type) = Callable (ForeignPtr CxxCallable)+newtype Callable (signature :: 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-     )+  :: forall n a f+   . (FuncBuilder f n a, IsHalideKernel (LoweredSignature f))   => Target   -> f-  -> IO (Callable inputs output)-compileToCallable target builder = do-  (args, func) <- buildFunc builder-  prepareCxxArguments args $ \args' ->-    withFunc func $ \func' ->+  -> IO (Callable (LoweredSignature f))+compileToCallable target builder =+  bracket (newCxxVector Nothing) deleteCxxVector $ \v -> do+    func <- buildFunc @f @n @a v builder+    withCxxFunc func $ \func' ->       withCxxTarget target $ \target' ->         wrapCxxCallable           =<< [C.throwBlock| Halide::Callable* {                 return handle_halide_exceptions([=]() {-                  return new Halide::Callable{-                    $(Halide::Func* func')->compile_to_callable(-                      *$(const std::vector<Halide::Argument>* args'),-                      *$(const Halide::Target* target'))};+                  auto& func = *$(Halide::Func* func');+                  auto& args = *$(std::vector<Halide::Argument>* v);+                  auto const& target = *$(const Halide::Target* target');+                  std::reverse(std::begin(args), std::end(args));+                  return new Halide::Callable{func.compile_to_callable(args, target)};                 });               } |]   where-    _ = keepRedundantConstraint (Proxy @(Ptr (HalideBuffer n a) ~ output))+    _ = keepRedundantConstraint @(IsHalideKernel (LoweredSignature f)) -callableToFunction-  :: forall inputs output kernel-   . ( Curry inputs (output -> IO ()) kernel-     , KnownNat (Length inputs)-     , All ValidArgument inputs-     , ValidArgument output-     )-  => Callable inputs output-  -> IO kernel+callableToFunction :: forall f. IsHalideKernel f => Callable f -> IO f callableToFunction (Callable callable) = do   context <- newEmptyCxxUserContext-  -- +1 comes from CxxUserContext and another +1 comes from output-  let argc = 2 + fromIntegral (natVal (Proxy @(Length inputs)))+  -- +1 comes from CxxUserContext+  let argc = 1 + fromIntegral (natVal (Proxy @(Length (FunctionArguments f))))   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)+      kernel args = do+        setArgvStorage storage contextPtr args         [CU.exp| void {           handle_halide_exceptions([=]() {             return $(Halide::Callable* callablePtr)->call_argv_fast(@@ -273,8 +219,167 @@         touch scalarStorage         touch context         touch callable-  pure $ curryG @inputs @(output -> IO ()) kernel+  pure $ curryG @(FunctionArguments f) @(IO ()) kernel +-- 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+--    . (ValidParameters' ts, All ValidParameter ts, KnownNat (Length ts))+--   => Arguments ts+--   -> (Ptr (CxxVector CxxArgument) -> IO b)+--   -> IO b+-- prepareCxxArguments args action = do+--   let count = fromIntegral (natVal (Proxy @(Length ts)))+--       allocate =+--         [CU.block| std::vector<Halide::Argument>* {+--           auto p = new std::vector<Halide::Argument>{};+--           p->reserve($(size_t count));+--           return p;+--         } |]+--       destroy p = [CU.exp| void { delete $(std::vector<Halide::Argument>* p) } |]+--   bracket allocate destroy $ \v -> do+--     let go :: (All ValidParameter ts') => Arguments ts' -> IO ()+--         go Nil = pure ()+--         go (x ::: xs) = appendToArgList v x >> go xs+--     go args+--     action v++wrapCxxUserContext :: Ptr CxxUserContext -> IO (ForeignPtr CxxUserContext)+wrapCxxUserContext = newForeignPtr deleter+  where+    deleter = [C.funPtr| void deleteUserContext(Halide::JITUserContext* p) { delete p; } |]++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; } |]++wrapCxxCallable :: Ptr CxxCallable -> IO (Callable signature)+wrapCxxCallable = fmap Callable . newForeignPtr deleter+  where+    deleter = [C.funPtr| void deleteCallable(Halide::Callable* p) { delete p; } |]++-- class All ValidArgument (LoweredOutputs t) => IsOutput t++-- 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 (Expr a)) = Ptr (HalideBuffer n a)+--   Lowered '[] = '[]+--   Lowered (t ': ts) = (Lowered t ': Lowered ts)++-- | A constraint that specifies that the function @f@ returns @'IO' ('Func' t n a)@.+-- class (FunctionReturn f ~ IO (Func 'FuncTy n a), KnownNat n) => ReturnsFunc f n a | f -> n a++-- instance (FunctionReturn f ~ IO (Func 'FuncTy n a), KnownNat n) => ReturnsFunc f n a++-- type family ValidParameters' (p :: [Type]) :: Constraint where+--   ValidParameters' (Expr a ': rest) = (IsHalideType a, ValidParameter (Expr a), ValidParameters' rest)+--   ValidParameters' (Func t n (Expr a) ': rest) = (t ~ 'ParamTy, IsHalideType a, ValidParameter (Func 'ParamTy n (Expr a)), ValidParameters' rest)+--   ValidParameters' (a ': rest) = (Bottom, ValidParameters' rest)+--   ValidParameters' '[] = ()++-- type IsFuncBuilder f n a =+--   ( ValidParameters' (FunctionArguments f)+--   , All ValidParameter (FunctionArguments f)+--   , All ValidArgument (Concat (Lowered (FunctionArguments f)) (LoweredOutputs (Func 'FuncTy n a)))+--   , UnCurry f (FunctionArguments f) (FunctionReturn f)+--   , PrepareParameters (FunctionArguments f)+--   , ReturnsFunc f n a+--   , KnownNat (Length (FunctionArguments f))+--   , KnownNat (Length (Lowered (FunctionArguments f)))+--   , KnownNat (Length (LoweredOutputs (Func 'FuncTy n a)))+--   )++-- buildFunc :: (IsFuncBuilder f n a) => f -> IO (Arguments (FunctionArguments f), Func 'FuncTy n a)+-- buildFunc builder = do+--   parameters <- prepareParameters+--   func <- uncurryG builder parameters+--   pure (parameters, func)++-- newtype Callable (inputs :: [Type]) (outputs :: [Type]) = Callable (ForeignPtr CxxCallable)++-- compileToCallable+--   :: forall n a f inputs outputs+--    . ( IsFuncBuilder f n a+--      , Lowered (FunctionArguments f) ~ inputs+--      , LoweredOutputs (Func 'FuncTy n a) ~ outputs+--      )+--   => Target+--   -> f+--   -> IO (Callable inputs outputs)+-- compileToCallable target builder = do+--   (args, func) <- buildFunc builder+--   prepareCxxArguments args $ \args' ->+--     case func of+--       Func fp ->+--         withForeignPtr fp $ \func' ->+--           withCxxTarget target $ \target' ->+--             wrapCxxCallable+--               =<< [C.throwBlock| Halide::Callable* {+--                     return handle_halide_exceptions([=]() {+--                       return new Halide::Callable{+--                         $(Halide::Func* func')->compile_to_callable(+--                           *$(const std::vector<Halide::Argument>* args'),+--                           *$(const Halide::Target* target'))};+--                     });+--                   } |]+--   where+--     _ = keepRedundantConstraint (Proxy @(LoweredOutputs (Func 'FuncTy n a) ~ outputs))++-- callableToFunction+--   :: forall inputs outputs kernel+--    . ( Curry (Concat inputs outputs) (IO ()) kernel+--      , KnownNat (Length inputs)+--      , KnownNat (Length outputs)+--      , All ValidArgument (Concat inputs outputs)+--      )+--   => Callable inputs outputs+--   -> IO kernel+-- callableToFunction (Callable callable) = do+--   context <- newEmptyCxxUserContext+--   -- +1 comes from CxxUserContext+--   let argc =+--         1+--           + fromIntegral (natVal (Proxy @(Length inputs)))+--           + fromIntegral (natVal (Proxy @(Length outputs)))+--   storage@(ArgvStorage argv scalarStorage) <- newArgvStorage (fromIntegral argc)+--   let argvPtr = P.mutablePrimArrayContents argv+--       contextPtr = unsafeForeignPtrToPtr context+--       callablePtr = unsafeForeignPtrToPtr callable+--       kernel args = do+--         setArgvStorage storage (contextPtr ::: args)+--         [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 @(Concat inputs outputs) @(IO ()) kernel+ -- | Convert a function that builds a Halide 'Func' into a normal Haskell function acccepting scalars and -- 'HalideBuffer's. --@@ -301,25 +406,21 @@ --     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-     )+  :: forall f n a+   . (FuncBuilder f n a, IsHalideKernel (LoweredSignature f))   => f   -- ^ Function to compile-  -> IO kernel+  -> IO (LoweredSignature f)   -- ^ 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-     )+  :: forall f n a+   . (FuncBuilder f n a, IsHalideKernel (LoweredSignature f))   => Target   -> f-  -> IO kernel+  -> IO (LoweredSignature f) compileForTarget target builder = compileToCallable target builder >>= callableToFunction  -- | Format in which to return the lowered code.@@ -344,22 +445,27 @@ -- -- 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+  :: forall n a f. (FuncBuilder f 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 ->+    bracket (newCxxVector Nothing) deleteCxxVector $ \v -> do+      func <- buildFunc @f @n @a v builder+      withCxxFunc func $ \func' ->+        withCxxTarget target $ \target' ->           [C.throwBlock| void {             handle_halide_exceptions([=]() {-              $(Halide::Func* f)->compile_to_lowered_stmt(+              auto& func = *$(Halide::Func* func');+              auto& args = *$(std::vector<Halide::Argument>* v);+              auto const& target = *$(const Halide::Target* target');+              std::reverse(std::begin(args), std::end(args));++              func.compile_to_lowered_stmt(                 std::string{$bs-ptr:s, static_cast<size_t>($bs-len:s)},-                *$(const std::vector<Halide::Argument>* v),+                args,                 static_cast<Halide::StmtOutputFormat>($(int o)),-                *$(Halide::Target* t));+                target);             });           } |]     T.readFile (dir <> "/code.stmt")

@@ -1,41 +0,0 @@-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

@@ -38,7 +38,7 @@  import Control.Monad (void) import Data.Text (Text, pack, unpack)-import qualified Data.Text as T+import Data.Text qualified as T import Data.Text.Encoding (encodeUtf8) import Foreign.C.Types (CInt (..)) import Foreign.ForeignPtr@@ -46,9 +46,9 @@ 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.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.Context import Language.Halide.Expr import Language.Halide.Func@@ -430,7 +430,7 @@ --         std::cout << "deleting ..." << std::endl; --         delete p; } |] -getStageSchedule :: (KnownNat n, IsHalideType a) => Stage n a -> IO StageSchedule+getStageSchedule :: Stage n a -> IO StageSchedule getStageSchedule stage =   withCxxStage stage $ \stage' ->     peekStageSchedule@@ -482,7 +482,7 @@           Mullapudi2016 -> "autoschedule_mullapudi2016"   loadLibrary path -applyAutoScheduler :: (KnownNat n, IsHalideType a) => AutoScheduler -> Target -> Func t n a -> IO Text+applyAutoScheduler :: KnownNat n => AutoScheduler -> Target -> Func t n a -> IO Text applyAutoScheduler scheduler target func = do   let s = encodeUtf8 . pack . show $ scheduler   withFunc func $ \f ->@@ -501,7 +501,7 @@ makeUnqualified :: Text -> Text makeUnqualified = snd . T.breakOnEnd "." -applySplit :: (KnownNat n, IsHalideType a) => Split -> Stage n a -> IO ()+applySplit :: KnownNat n => Split -> Stage n a -> IO () applySplit (SplitVar x) stage = do   oldVar <- mkVar (makeUnqualified x.splitOld)   outerVar <- mkVar (makeUnqualified x.splitOuter)@@ -513,10 +513,10 @@   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 :: KnownNat n => [Split] -> Stage n a -> IO () applySplits xs stage = mapM_ (`applySplit` stage) xs -applyDim :: (KnownNat n, IsHalideType a) => Dim -> Stage n a -> IO ()+applyDim :: KnownNat n => Dim -> Stage n a -> IO () applyDim x stage = do   var' <- mkVar (makeUnqualified x.var)   void $@@ -530,13 +530,13 @@       ForGPUThread -> gpuThreads x.deviceApi var' stage       ForGPULane -> gpuLanes x.deviceApi var' stage -applyDims :: (KnownNat n, IsHalideType a) => [Dim] -> Stage n a -> IO ()+applyDims :: KnownNat n => [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 :: KnownNat n => StageSchedule -> Stage n a -> IO () applySchedule schedule stage = do   applySplits schedule.splits stage   applyDims schedule.dims stage

@@ -26,6 +26,7 @@   , TargetFeature (..)   , DeviceAPI (..)   -- , targetFeatureForDeviceAPI+  , deviceAPIForTarget      -- * Internal   , withCxxTarget@@ -38,9 +39,9 @@ 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.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.Context import Language.Halide.Type import Language.Halide.Utils@@ -188,6 +189,13 @@ withCxxTarget :: Target -> (Ptr CxxTarget -> IO a) -> IO a withCxxTarget (Target fp) = withForeignPtr fp +deviceAPIForTarget :: Target -> DeviceAPI+deviceAPIForTarget target = unsafePerformIO $+  withCxxTarget target $ \target' ->+    toEnum . fromIntegral+      <$> [CU.exp| int { static_cast<int>(Halide::get_default_device_api_for_target(+            *$(Halide::Target* target'))) } |]+ -- targetFeatureForDeviceAPI :: DeviceAPI -> Maybe TargetFeature -- targetFeatureForDeviceAPI deviceAPI = --   toFeature . unsafePerformIO $@@ -354,7 +362,8 @@   | FeatureSanitizerCoverage   | FeatureProfileByTimer   | FeatureSPIRV-  | FeatureSemihosting+  -- removed in v15.0.0+  -- \| FeatureSemihosting   deriving stock (Eq, Show, Ord)  instance Enum TargetFeature where@@ -437,7 +446,8 @@       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 } |]++  -- 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@@ -514,5 +524,5 @@     | 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+    -- \| fromIntegral k == [CU.pure| int { halide_target_feature_semihosting } |] = FeatureSemihosting     | otherwise = error $ "unknown Target feature: " <> show k

@@ -29,8 +29,8 @@ 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.C.Inline qualified as C+import Language.C.Inline.Unsafe qualified as CU import Language.Halide.Buffer import Language.Halide.Context import Language.Halide.Dimension@@ -130,11 +130,14 @@     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+  :: KnownNat n+  => (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@@ -152,20 +155,20 @@           func.jit_handlers().custom_trace = nullptr;         } |] -traceStores :: (KnownNat n, IsHalideType a) => Func t n a -> IO (Func t n a)+traceStores :: KnownNat n => 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 :: KnownNat n => 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)+  :: KnownNat n   => (TraceEventCode -> Bool)   -> Func t n a   -> IO b

@@ -6,6 +6,7 @@ {-# LANGUAGE PolyKinds #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TypeFamilyDependencies #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -Wno-unused-local-binds -Wno-unused-matches #-} @@ -23,7 +24,6 @@   , CxxVarOrRVar   , CxxFunc   , CxxParameter-  , CxxArgument   , CxxImageParam   , CxxVector   , CxxUserContext@@ -34,6 +34,7 @@   , Arguments (..)   , Length   , Append+  , Concat   , argumentsAppend   , FunctionArguments   , FunctionReturn@@ -43,9 +44,6 @@   , defineIsHalideTypeInstances   , instanceHasCxxVector   , HasCxxVector (..)-  , IsTuple (..)-  , FromTuple-  , ToTuple   , instanceCxxConstructible   , CxxConstructible (..)   -- defineCastableInstances,@@ -60,16 +58,17 @@ import Data.Constraint import Data.Int import Data.Kind (Type)-import qualified Data.Text as T+import Data.Text qualified as T import Data.Word import Foreign.C.Types import Foreign.ForeignPtr import Foreign.Ptr import Foreign.Storable+import GHC.ForeignPtr (mallocForeignPtrAlignedBytes) 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.C.Inline qualified as C+import Language.C.Inline.Unsafe qualified as CU+import Language.Haskell.TH qualified as TH import Language.Haskell.TH.Syntax (Lift)  -- | Haskell counterpart of @Halide::Expr@.@@ -87,9 +86,6 @@ -- | Haskell counterpart of @Halide::Internal::Parameter@. data CxxParameter --- | Haskell counterpart of @Halide::Argument@.-data CxxArgument- -- | Haskell counterpart of @Halide::ImageParam@. data CxxImageParam @@ -120,10 +116,12 @@  cxxConstructWithDeleter :: Int -> FinalizerPtr a -> (Ptr a -> IO ()) -> IO (ForeignPtr a) cxxConstructWithDeleter size deleter constructor = do-  fp <- mallocForeignPtrBytes size+  fp <- mallocForeignPtrAlignedBytes size align   withForeignPtr fp constructor   addForeignPtrFinalizer deleter fp   pure fp+  where+    align = 64  -- data Split = --   SplitVar !Text !Text !Text !(Expr Int32) !@@ -309,6 +307,10 @@   Append '[] y = '[y]   Append (x ': xs) y = x ': Append xs y +type family Concat (xs :: [k]) (ys :: [k]) :: [k] where+  Concat '[] ys = ys+  Concat (x ': xs) ys = x ': Concat xs ys+ -- | Append a value to 'Arguments' argumentsAppend :: Arguments xs -> t -> Arguments (Append xs t) argumentsAppend = go@@ -328,7 +330,7 @@   FunctionReturn a = a  -- | Apply constraint to all types in a list.-type family All (c :: Type -> Constraint) (ts :: [Type]) :: Constraint where+type family All (c :: Type -> Constraint) (ts :: [Type]) = (p :: Constraint) | p -> ts where   All c '[] = ()   All c (t ': ts) = (c t, All c ts) @@ -337,7 +339,7 @@ -- -- 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+class UnCurry (f :: Type) (args :: [Type]) (r :: Type) | args r -> f, args f -> r where   uncurryG :: f -> Arguments args -> r  instance (FunctionArguments f ~ '[], FunctionReturn f ~ r, f ~ r) => UnCurry f '[] r where@@ -362,53 +364,3 @@  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

@@ -0,0 +1,25 @@+module Language.Halide.BoundaryConditionsSpec (spec) where++import Language.Halide+import Test.Hspec++spec :: Spec+spec = do+  it "extends the domain by repeating the edges" $ do+    let mat :: [[Float]]+        mat = [[1, 2, 3], [4, 5, 6]]+    asBufferParam @2 @Float mat $ \mat' -> do+      [i, j] <- mapM mkVar ["i", "j"]+      extended <- repeatEdge mat'+      f <- define "f" (i, j) $ extended ! (i - 1, j - 1)+      realize f [4, 5] peekToList+        `shouldReturn` [[1, 1, 2, 3, 3], [1, 1, 2, 3, 3], [4, 4, 5, 6, 6], [4, 4, 5, 6, 6]]+  it "extends the domain by filling with a constant value" $ do+    let mat :: [[Float]]+        mat = [[1, 2, 3], [4, 5, 6]]+    asBufferParam @2 @Float mat $ \mat' -> do+      [i, j] <- mapM mkVar ["i", "j"]+      extended <- constantExterior 0 mat'+      f <- define "f" (i, j) $ extended ! (i - 1, j - 1)+      realize f [4, 5] peekToList+        `shouldReturn` [[0, 0, 0, 0, 0], [0, 1, 2, 3, 0], [0, 4, 5, 6, 0], [0, 0, 0, 0, 0]]

@@ -3,7 +3,7 @@ module Language.Halide.BufferSpec (spec) where  import Data.Int (Int64)-import Foreign.Ptr (Ptr, nullPtr)+import Foreign.Ptr (nullPtr) import Language.Halide import Test.Hspec import Test.Hspec.QuickCheck@@ -18,6 +18,9 @@ newtype ListTensor3D a = ListTensor3D [[[a]]]   deriving stock (Show) +newtype ListTensor4D a = ListTensor4D [[[[a]]]]+  deriving stock (Show)+ instance Arbitrary a => Arbitrary (ListVector a) where   arbitrary = ListVector <$> listOf arbitrary @@ -34,6 +37,14 @@     d2 <- chooseInt (0, 30)     ListTensor3D <$> vectorOf d0 (vectorOf d1 (vector d2)) +instance Arbitrary a => Arbitrary (ListTensor4D a) where+  arbitrary = do+    d0 <- chooseInt (0, 30)+    d1 <- chooseInt (0, 30)+    d2 <- chooseInt (0, 30)+    d3 <- chooseInt (0, 30)+    ListTensor4D <$> vectorOf d0 (vectorOf d1 (vectorOf d2 (vector d3)))+ spec :: Spec spec = do   it "rowMajorStrides" $ do@@ -52,3 +63,27 @@     withHalideBuffer @2 @Int64 xs peekToList `shouldReturn` xs   prop "works with [[[a]]]" $ \(ListTensor3D xs :: ListTensor3D Double) ->     withHalideBuffer @3 @Double xs peekToList `shouldReturn` xs+  modifyMaxSuccess (const 20) $+    prop "works with [[[[a]]]]" $ \(ListTensor4D @Double xs) ->+      withHalideBuffer @4 @Double xs peekToList `shouldReturn` xs+  it "creates cropped buffers" $ do+    let mkFill2D (scalar @Float "value" -> value) = do+          [i, j] <- mapM mkVar ["i", "j"]+          define "fill" (i, j) value+    fill <- compile mkFill2D+    allocaCpuBuffer [4, 4] $ \buf -> do+      fill 0 buf+      peekToList buf `shouldReturn` [[0, 0, 0, 0],+                                     [0, 0, 0, 0],+                                     [0, 0, 0, 0],+                                     [0, 0, 0, 0]]+      withCropped buf 1 0 1 $ fill 1+      peekToList buf `shouldReturn` [[1, 0, 0, 0],+                                     [1, 0, 0, 0],+                                     [1, 0, 0, 0],+                                     [1, 0, 0, 0]]+      withCropped buf 0 1 2 $ fill 2+      peekToList buf `shouldReturn` [[1, 0, 0, 0],+                                     [2, 2, 2, 2],+                                     [2, 2, 2, 2],+                                     [1, 0, 0, 0]]

@@ -2,12 +2,12 @@  import Control.Monad (unless, when) import Data.Int+import Data.Text (Text) 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 Test.QuickCheck import Type.Reflection import Utils @@ -21,48 +21,46 @@         || op (toInteger x) (toInteger y) < toInteger (minBound @Int64)   | otherwise = False -infix 1 `evaluatesTo`+-- infix 1 `evaluatesTo` -evaluatesTo :: (Eq a, IsHalideType a) => Expr a -> a -> PropertyM IO ()-evaluatesTo expr expected =-  assert . (expected ==) =<< (run . evaluate) expr+-- evaluatesTo :: (Eq a, IsHalideType a) => Expr a -> a -> PropertyM IO ()+-- evaluatesTo expr expected =+--   assert . (expected ==) =<< (run . evaluate) expr+-- -infix 1 `evaluatesToApprox`+infix 1 `shouldEvaluateToApprox` -evaluatesToApprox :: (Ord a, IsHalideType a, HasEpsilon a) => Expr a -> a -> PropertyM IO ()-evaluatesToApprox expr expected =-  assert . approx' expected =<< (run . evaluate) expr+shouldEvaluateToApprox :: (Ord a, IsHalideType a, HasEpsilon a, Show a) => Expr a -> a -> Expectation+shouldEvaluateToApprox expr expected = do+  e <- evaluate expr+  e `shouldBeApprox` expected  infix 1 `shouldEvaluateTo`  shouldEvaluateTo :: (Eq a, IsHalideType a, Show a) => Expr a -> a -> Expectation-shouldEvaluateTo expr expected =-  evaluate expr `shouldReturn` expected+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+    prop "Bool" $ \x ->+      mkExpr (x :: Bool) `shouldEvaluateTo` x    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+        p :: forall a. (IsHalideType a, Eq a, Num a, Typeable a, Show a) => a -> a -> Expectation+        p x y = do+          whenNotOverflowing (+) x y $+            mkExpr x + mkExpr y `shouldEvaluateTo` x + y+          whenNotOverflowing (-) x y $+            mkExpr x - mkExpr y `shouldEvaluateTo` x - y+          whenNotOverflowing (*) x y $+            mkExpr x * mkExpr y `shouldEvaluateTo` x * y+          unless (x == -128) $+            abs (mkExpr x) `shouldEvaluateTo` abs x+          negate (mkExpr x) `shouldEvaluateTo` negate x     prop "Int8" $ p @Int8     prop "Int16" $ p @Int16     prop "Int32" $ p @Int32@@ -75,38 +73,55 @@     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+    let p :: forall a. (IsHalideType a, Eq a, Fractional a, Show a) => a -> a -> Expectation+        p x y = do+          unless (x == 0 && y == 0) $+            mkExpr x / mkExpr y `shouldEvaluateTo` x / y+    prop "Float" $ property (p @Float)+    prop "Double" $ property (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+    let p :: forall a. (IsHalideType a, Ord a, Floating a, HasEpsilon a, Show a) => a -> Expectation+        p x = do+          let y = mkExpr x           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+            log y `shouldEvaluateToApprox` log x+            sqrt y `shouldEvaluateToApprox` sqrt x+          when (x < 50) $+            exp y `shouldEvaluateToApprox` exp x+          sin y `shouldEvaluateToApprox` sin x+          cos y `shouldEvaluateToApprox` cos x+          tan y `shouldEvaluateToApprox` 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+            asin y `shouldEvaluateToApprox` asin x+            acos y `shouldEvaluateToApprox` acos x+            atan y `shouldEvaluateToApprox` atan x+          when (abs x < 50) $ do+            sinh y `shouldEvaluateToApprox` sinh x+            cosh y `shouldEvaluateToApprox` cosh x+          tanh y `shouldEvaluateToApprox` tanh x+          asinh y `shouldEvaluateToApprox` asinh x           when (x >= 1) $-            acosh (mkExpr x) `evaluatesToApprox` acosh x+            acosh y `shouldEvaluateToApprox` acosh x           when (-1 <= x && x <= 1) $-            atanh (mkExpr x) `evaluatesToApprox` atanh x+            atanh y `shouldEvaluateToApprox` atanh x     prop "Float" $ p @Float     prop "Double" $ p @Double     it "defines pi" $ do-      (pi :: Expr Float) `shouldEvaluateTo` pi-      (pi :: Expr Double) `shouldEvaluateTo` pi+      (pi :: Expr Float) `shouldEvaluateToApprox` pi+      (pi :: Expr Double) `shouldEvaluateToApprox` pi++  describe "printed" $+    it "prints expressions when evaluated" $ do+      printed (1 :: Expr Int32) `shouldEvaluateTo` 1+      printed (1 :: Expr Int32) ("<- when" :: String) ("haha" :: String) `shouldEvaluateTo` 1+      let x :: Expr Float+          x = 1+       in printed (sin x) ("<- sin(" :: Text) x (")" :: Text) `shouldEvaluateToApprox` sin 1++  -- describe "Show" $+  --   it "shows 123" $ do+  --     show (123 :: Expr Int32) `shouldBe` "123"+  describe "testWriteToStderr" $ do+    it "shows 123" $ do+      testWriteToStderr

@@ -7,7 +7,7 @@ module Language.Halide.FuncSpec (spec) where  import Control.Monad.ST (RealWorld)-import qualified Data.Vector.Storable.Mutable as SM+import Data.Vector.Storable.Mutable qualified as SM import Language.Halide import Test.Hspec hiding (parallel) import Utils@@ -33,7 +33,7 @@       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]+      realize g [] peekScalar `shouldReturn` 20.6    describe "vectorize" $ do     it "vectorizes loops" $ do@@ -240,13 +240,41 @@         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+        f <- define "sum" i (mkExpr 0)+        update f (0 :: Expr Int32) $ f ! 0 + src ! r         realize f [1] peekToList `shouldReturn` ([15] :: [Int32])+    it "computes multi-dimensional reductions" $ do+      asBufferParam @2 @Int32 ([[1, 2, 3], [4, 5, 6]] :: [[Int32]]) $ \src -> do+        d0 <- (.extent) <$> dim 0 src+        d1 <- (.extent) <$> dim 1 src+        r <- toRVars =<< mkRDom "r" (0, 0) (d0, d1)+        f <- define "sum" () (mkExpr 0)+        update f () $ f ! () + src ! r+        realize f [] peekScalar `shouldReturn` (21 :: 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+      f <- define "f" i $ ifThenElse (i `gt` 5) i 0+      update f i $ ifThenElse ((f ! i) `eq` 0) (2 * i) undef       realize f [10] peekToList `shouldReturn` ([0, 2, 4, 6, 8, 10] <> [6 .. 9] :: [Int32])++  describe "Tuples" $ do+    it "defines tuple Funcs" $ do+      i <- mkVar "i"+      f <- define "f" i (i, i)+      _ <- computeRoot f+      g <- define "g" i $ fst (f ! i)+      _ <- computeRoot g+      realize g [5] peekToList `shouldReturn` [0 .. 4]+      pure ()+    it "compiles tuple Funcs" $ do+      kernel <- compile $ do+        i <- mkVar "i"+        define "f" i (i, i + 1)++      allocaCpuBuffer [5] $ \a ->+        allocaCpuBuffer [5] $ \b -> do+          kernel a b+          peekToList a `shouldReturn` [0 .. 4]+          peekToList b `shouldReturn` [1 .. 5]

@@ -13,13 +13,13 @@     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+        define "out" i $ (a ! i :: Expr Float) + 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+      withHalideBuffer @1 @Float a $ \a' ->+        withHalideBuffer @_ @Float b $ \b' ->+          allocaCpuBuffer @_ @Float [n] $ \out' -> do             vectorPlus a' b' out'             peekToList out' `shouldReturn` zipWith (+) a b @@ -28,10 +28,7 @@         i <- mkVar "i"         j <- mkVar "j"         define "out" (i, j) $-          bool-            (i `eq` j)-            (v ! i / scale)-            0+          ifThenElse (i `eq` j) (v ! i / scale) 0       let a :: [Double]           a = [1.0, 2.0, 3.0]       withHalideBuffer a $ \a' ->@@ -43,7 +40,7 @@       scaledDiagonal <- compile $ \(scale :: Expr Double) v -> do         i <- mkVar "i"         j <- mkVar "j"-        out <- define "out" (i, j) 0+        out <- define "out" (i, j) (mkExpr 0)         update out (i, i) (v ! i / scale)         pure out       let a :: [Double]

@@ -3,14 +3,14 @@ module Language.Halide.ScheduleSpec (spec) where  import Control.Monad (forM_)-import qualified Data.Text.IO as T+import Data.Text.IO qualified 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 :: (KnownNat n, IsHalideType a) => IO (Func t n (Expr a)) -> (Func t n (Expr a) -> IO ()) -> Expectation checkScheduleRoundTrip prepare schedule = do   f1 <- prepare   f2 <- prepare@@ -146,6 +146,19 @@         checkScheduleRoundTrip prepare (schedule hostTarget (Just Adams2019))         checkScheduleRoundTrip prepare (schedule hostTarget (Just Li2018))         checkScheduleRoundTrip prepare (schedule hostTarget (Just Mullapudi2016))++    it "supports scalar estimates" $ do+      let f scheduler (scalar @Float "c" -> c) = do+            [x, y] <- mapM mkVar ["x", "y"]+            f <- define "f" (x, y) c+            estimate x 0 100 f+            estimate y 0 100 f+            setScalarEstimate 5.0 c+            loadAutoScheduler scheduler+            void $ applyAutoScheduler scheduler hostTarget f+            pure f+      _ <- compile $ f Adams2019+      pure ()  -- (x, y, z, xInner, f1) <- prepare -- split TailAuto x (x, xInner) 2 f1

@@ -6,10 +6,11 @@   ( shouldContainText   , shouldNotContainText   , appearsBeforeText-  , shouldApproxBe+  , shouldBeApprox+  , shouldBeEqForTesting   , testOnGpu   , approx-  , approx'+  , approxWith   , (&)   , void   , T.hPutStrLn@@ -18,7 +19,6 @@   , eps   , showInCodeLenses   , EqForTesting (..)-  , shouldBeEqForTesting   ) where @@ -26,13 +26,12 @@ 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 Data.Text qualified as T+import Data.Text.IO qualified as T import GHC.Exts (IsList (..)) import GHC.Stack-import Language.Halide+import Language.Halide hiding (and, max) import System.IO (stderr)-import Test.HUnit import Test.HUnit.Lang (FailureReason (..), HUnitFailure (..)) import Test.Hspec @@ -53,6 +52,17 @@     Just t -> f t     Nothing -> pendingWith "no GPU target available" +compareWith :: (HasCallStack, Show a) => (a -> a -> Bool) -> a -> a -> Expectation+compareWith comparator result expected =+  unless (comparator result expected) $ do+    throwIO (HUnitFailure location $ ExpectedButGot Nothing expectedMsg actualMsg)+  where+    expectedMsg = show expected+    actualMsg = show result+    location = case reverse (toList callStack) of+      (_, loc) : _ -> Just loc+      [] -> Nothing+ class Num a => HasEpsilon a where   eps :: a @@ -62,17 +72,18 @@ 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))+approxWith :: (Ord a, Num a) => a -> a -> a -> a -> Bool+approxWith 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+approx :: (Ord a, HasEpsilon a) => a -> a -> Bool+approx a b = approxWith (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+shouldBeApprox :: (Ord a, HasEpsilon a, Show a) => a -> a -> Expectation+shouldBeApprox = compareWith approx +shouldBeEqForTesting :: (HasCallStack, EqForTesting a, Show a) => a -> a -> Expectation+shouldBeEqForTesting = compareWith equalForTesting+ showInCodeLenses :: Text -> IO String showInCodeLenses v = error (unpack v) @@ -128,14 +139,3 @@       , 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