packages feed

futhark 0.26.3 → 0.26.4

raw patch · 69 files changed

+2496/−1456 lines, 69 filesdep ~futhark-manifest

Dependency ranges changed: futhark-manifest

Files

CHANGELOG.md view
@@ -5,6 +5,41 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.26.4]++### Added++* `futhark repl` has a new command: `:string`.++* The `hip` backend previously simulated `f16` operations with `f32`, but now it+  uses the hardware support for `f16`, similarly to the CUDA backend.+  Implemented by Jérôme Wagner. (#2470)++* Vector AD, exposed through the functions `jmp` and `mjp`.++* All opaque values available over the C API can now be decomposed into their+  constituents.++* The manifest now contains documentation for entry points and opaque types.++### Fixed++* Non-exhaustive pattern match warnings were not always emitted when wildcard+  patterns and explicit constructors were mixed. (#2483)++* Invalid fusion that could cause compiler crash. (#2474)++* GPU code generation of segmented reductions with array operands. (#2227,+  properly this time, and #2482)++* Use ascripted element type in API functions for arrays of records. (#2485)++* Consumption checking of certain local polymorphic functions (in practice,+  polymorphic functions that can only be written via holes). (#2488)++* A regression in fusion of forms such as `scatter dest (flatten inds) (flatten+  vals)`. (#2452)+ ## [0.26.3]  ### Added
docs/man/futhark-hip.rst view
@@ -99,6 +99,10 @@ compiler to compile the generated C program into a binary.  This only works if the C compiler can find the necessary HIP libraries. +At runtime (not just compile-time!), the ``HIP_PATH`` environment variable (if+set) must point to the HIP installation directory, which most importantly must+contain an ``include`` subdirectory with the HIP header files.+ SEE ALSO ========
futhark.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name:           futhark-version:        0.26.3+version:        0.26.4 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -116,7 +116,9 @@       Futhark.Actions       Futhark.AD.Derivatives       Futhark.AD.Fwd+      Futhark.AD.Shared       Futhark.AD.Rev+      Futhark.AD.Rev.Acc       Futhark.AD.Rev.Loop       Futhark.AD.Rev.Hist       Futhark.AD.Rev.Map@@ -299,7 +301,6 @@       Futhark.Internalise.Entry       Futhark.Internalise.Exps       Futhark.Internalise.FullNormalise-      Futhark.Internalise.Lambdas       Futhark.Internalise.LiftLambdas       Futhark.Internalise.Monad       Futhark.Internalise.Monomorphise@@ -480,7 +481,7 @@     , free >=5.1.10     , futhark-data >= 1.1.3.0     , futhark-server >= 1.4.1.0-    , futhark-manifest == 1.8.0.0+    , futhark-manifest == 1.9.0.0     , githash >=0.1.6.1     , half >= 0.3     , haskeline
prelude/ad.fut view
@@ -14,104 +14,130 @@ -- -- Futhark's AD support includes the following: -----   * Differentiation operators for forward-mode (`jvp`) and reverse-mode---     (`vjp`).+--   * Differential operators for forward-mode (`jvp`@term) and reverse-mode+--     (`vjp`@term). -----   * Arbitrary control flow in differentiable code.+--   * Almost arbitrary control flow in differentiable code (some limitations+--     apply when using GPU backends, see below). -- --   * Higher order derivatives by nesting differentiation operators, including --     arbitrary mixing of forward- and reverse mode (although using multiple --     rounds of reverse mode is rarely useful and often slow). -----   * Custom derivatives (`with_vjp`).+--   * Custom derivatives (`with_vjp`@term). --+--   * Vector AD (`mjp`@term, `jmp`@term), sometimes also known as "batched" or+--    "multi-directional" AD.+-- --   * Checkpointing of sequential loops. -- -- ## Jacobians ----- For a differentiable function *f* whose input comprise *n* scalars--- and whose output comprises *m* scalars, the--- [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant)--- for a given input point is an *m* by *n* matrix of scalars that--- each represent a [partial--- derivatives](https://en.wikipedia.org/wiki/Partial_derivative).--- Intuitively, position *(i,j)* of the Jacobian describes how--- sensitive output *i* is to input *j*. The notion of Jacobian--- generalises to functions that accept or produce compound structures--- such as arrays, records, sums, and so on, simply by "flattening--- out" the values and considering only their constituent scalars.+-- For a differentiable function *f* whose input comprise *n* scalars and whose+-- output comprises *m* scalars, the+-- [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) for+-- a given input point is an *m* by *n* matrix of scalars that each represent a+-- [partial derivative](https://en.wikipedia.org/wiki/Partial_derivative).+-- Intuitively, position *(i,j)* of the Jacobian describes how sensitive output+-- *i* is to input *j*. The notion of Jacobian generalises to functions that+-- accept or produce compound structures such as arrays, records, sums, and so+-- on, simply by "flattening out" the values and considering only their+-- constituent scalars. ----- Computing the full Jacobian is usually costly and sometimes not--- necessary, and it is not part of the AD facility provided by--- Futhark. Instead it is possible to parts of the Jacobian.+-- Computing the full Jacobian is usually costly and sometimes not necessary,+-- and it is not part of the AD facility provided by Futhark. Instead it is+-- possible to compute parts of the Jacobian, which semantically (but not+-- operationally) can be seen as multiplying the Jacobian with a vector,+-- producing a vector. However, it is important to understand that the full+-- Jacobian is *not* constructed as an intermediate step. ----- We can take the product of an an *m* by *n* Jacobian with an--- *n*-element *tangent vector* to produce an *m*-element vector--- (*Jacobian-vector product*). Such a product can be computed in a--- single (augmented) execution of the function *f*, and by choosing--- the tangent vector appropriately we can use this to compute the--- full Jacobian. This is provided by the function `jvp`.+-- We can take the product of an *m* by *n* Jacobian with an *n*-element+-- *tangent vector* to produce an *m*-element vector (*Jacobian-vector+-- product*). Such a product can be computed in a single (augmented) execution+-- of the function *f*. This is provided by the function `jvp`. -- -- We can also take the product of an *m*-element vector *cotangent -- vector* with the *m* by *n* Jacobian to produce an *n*-element--- vector (*Vector-Jacobian product*). This too can be computed in a+-- vector (*vector-Jacobian product*). This too can be computed in a -- single execution of *f*, with `vjp`. ----- We can use the `jvp` function to produce a *column* of the full--- Jacobian, and `vjp` to produce a *row*. Which is superior for a--- given situation depends on whether the function has more inputs or--- outputs.+-- A tangent has the same structure as the input and represents a direction in+-- input space. A cotangent has the same structure as the output and represents+-- sensitivities flowing backwards through the computation. ----- You can freely nest `vjp` and `jvp` to compute higher-order--- derivatives.+-- Using an elementary (co-)tangent vector, we can use the `jvp` function to+-- produce a *column* of the full Jacobian, and `vjp` to produce a *row*, with+-- the nonzero element of the vector identifying which column or row is+-- extracted. Which is superior for a given situation depends on whether the+-- function has more inputs or outputs. --+-- We can freely nest `vjp` and `jvp` to compute higher-order derivatives.+-- -- ## Efficiency -- -- Both `jvp` and `vjp` work by transforming the program to carry -- along extra information associated with each scalar value. ----- In the case of `jvp`, this extra information takes the form of an--- additional scalar representing the tangent, which is then--- propagated in each scalar computation using essentially the [chain--- rule](https://en.wikipedia.org/wiki/Chain_rule). Therefore, `jvp`--- has a memory overhead of approximately *2x*, and a computational--- overhead of slightly more, but usually less than *4x*.+-- In the case of `jvp` ("forward mode", or "tangent mode"), this extra+-- information takes the form of an additional scalar representing the tangent,+-- which is then propagated in each scalar computation using essentially the+-- [chain rule](https://en.wikipedia.org/wiki/Chain_rule). Therefore, `jvp` has+-- a memory overhead of approximately *2x*, and a computational overhead of+-- slightly more, but usually less than *4x*. ----- In the case of `vjp`, since our starting point is a *cotangent*,--- the function is essentially first run forward, then backwards (the--- *return sweep*) to propagate the cotangent. During the return--- sweep, all intermediate results computed during the forward sweep--- must still be available, and must therefore be stored in memory--- during the forward sweep. This means that the memory usage of `vjp`--- is proportional to the number of sequential steps of the original--- function (essentially turning *time* into *space*). The compiler--- does a nontrivial amount of optimisation to ameliorate this--- overhead (see [AD for an Array Language with Nested--- Parallelism](https://futhark-lang.org/publications/sc22-ad.pdf)),--- but it can still be substantial for programs with deep sequential--- loops.+-- In the case of `vjp` ("reverse mode" or "adjoint mode"), since our starting+-- point is a *cotangent*, the function is essentially first run forward, then+-- backwards (the *return sweep*) to propagate the cotangent. During the return+-- sweep, all intermediate results computed during the forward sweep must still+-- be available, and must therefore be stored in memory during the forward sweep+-- - this is called "the tape". This means that the memory usage of `vjp` is+-- proportional to the number of sequential steps of the original function+-- (essentially turning *time* into *space*). The compiler does a nontrivial+-- amount of optimisation to ameliorate this overhead (see [AD for an Array+-- Language with Nested+-- Parallelism](https://futhark-lang.org/publications/sc22-ad.pdf)), but it can+-- still be substantial for programs with deep sequential loops. --+-- Nesting `vjp`, understood as applying `vjp` to the result of `vjp`, is+-- usually a bad idea, as the code structure produced by `vjp` is fairly+-- complicated, due to the tape management. Passing the output of `jvp` to+-- `vjp`, or the other way, is however fine. As a rule of thumb, whenever you+-- stack multiple differential operators, make sure only one of them is `vjp` or+-- related ones.+--+-- When using vector AD (`mjp`@term/`jmp`@term), each scalar is associated with+-- a vector of tangents or cotangents, and the space overhead for storing these+-- is therefore multiplied with the vector size. However, in the case of `vjp`,+-- the intermediate results are only stored once. It varies on a case-by-case+-- basis whether vector AD is faster than using `map` on top of+-- `vjp`@term/`jvp`@term. Vector AD essentially converts propagation of+-- (co-)tangents from scalar to array operations, which can have a significant+-- impact on memory accesses, depending on how the compiler manages to optimise+-- the resulting code. It is hard to predict whether this offsets the reduction+-- in primal work. If the vector size is a constant, and the `#[unroll]`+-- attribute is put on the AD operator, then the vectors become unrolled (turned+-- into tuples, essentially), although this should only be done when the vector+-- size is quite small, as the increase in code size is substantial.+-- -- ## Differentiable functions ----- AD only gives meaningful results for differentiable functions. The--- Futhark type system does not distinguish differentiable or--- non-differentiable operations. As a rule of thumb, a function is--- differentiable if its results are computed using a composition of--- primitive floating-point operations, without ever converting to or--- from integers.+-- AD only gives meaningful results for differentiable functions. The Futhark+-- type system does not distinguish differentiable from non-differentiable+-- operations. As a rule of thumb, a function is differentiable if its results+-- are computed using a composition of primitive floating-point operations,+-- without ever converting to or from integers. Most functions will also have+-- discontinuities around values that influence control flow. ----- Note that a function whose input or output is a sum type with more--- than one constructor is *not* differentiable (or at least the--- sum-typed part is not). This is because the choice of constructor--- is not a continuous quantity.+-- Note that a function whose input or output is a sum type with more than one+-- constructor is *not* differentiable (or at least the sum-typed part is not).+-- This is because the choice of constructor is not a continuous quantity. -- -- ## Limitations ----- `jvp` is expected to work in all cases. `vjp` has limitations when--- using the GPU backends similar to those for irregular flattening.--- Specifically, you should avoid structures with variant sizes, such--- as loops that carry an array that changes size through the--- execution of the loop.+-- `jvp` is expected to work in all cases. `vjp` has limitations when using the+-- GPU backends similar to those for irregular flattening. Specifically, you+-- should avoid structures with variant sizes, such as loops that carry an array+-- that changes size through the execution of the loop.  -- | Jacobian-Vector Product ("forward mode"), producing also the -- primal result as the first element of the result tuple.@@ -123,6 +149,20 @@ def vjp2 'a 'b (f: a -> b) (x: a) (y': b) : (b, a) =   intrinsics.vjp2 f x y' +-- | Jacobian-Matrix Product, returning also the primal result. As `jvp2`, but+-- accepts an array of seed vectors (hence "matrix", although transposed).+-- Semantically equivalent to mapping, but may be more efficient. If used with+-- `#[unroll]`, tangent calculations are unrolled when possible.+def jmp2 'a 'b [n] (f: a -> b) (x: a) (x': [n]a) : (b, [n]b) =+  intrinsics.jmp2 f x x'++-- | Matrix-Jacobian Product, returning also the primal result. As `vjp2`, but+-- accepts an array of seed vectors (hence "matrix"). Semantically equivalent to+-- mapping, but may be more efficient. If used with `#[unroll]`, adjoint+-- calculations are unrolled when possible.+def mjp2 'a 'b [n] (f: a -> b) (x: a) (y': [n]b) : (b, [n]a) =+  intrinsics.mjp2 f x y'+ -- | Jacobian-Vector Product ("forward mode"). def jvp 'a 'b (f: a -> b) (x: a) (x': a) : b =   (jvp2 f x x').1@@ -131,6 +171,16 @@ def vjp 'a 'b (f: a -> b) (x: a) (y': b) : a =   (vjp2 f x y').1 +-- | Jacobian-Matrix Product. As `jvp`, but accepts a vector of seed values.+-- Semantically equivalent to mapping, but may be more efficient.+def jmp 'a 'b [n] (f: a -> b) (x: a) (x': [n]a) : [n]b =+  (jmp2 f x x').1++-- | Matrix-Jacobian product. As `vjp`, but accepts a vector of seed values.+-- Semantically equivalent to mapping, but may be more efficient.+def mjp 'a 'b [n] (f: a -> b) (x: a) (y': [n]b) : [n]a =+  (mjp2 f x y').1+ -- | Provide custom reverse-mode adjoint code for a given function. This is -- useful when the adjoint synthesised by AD is not as good as one that is known -- analytically.@@ -144,8 +194,7 @@ -- primal result of `with_vjp`, and some part is only used in `f'`. -- -- **Beware:** if `f` uses any free variables, these will not be taken into--- **account when computing the adjoint. Make these part of the argument--- **instead.+-- account when computing the adjoint. Make these part of the argument instead. def with_vjp 'a 'b (f: a -> b) (f': (res: b) -> (b_adj: b) -> a) (x: a) : b =   intrinsics.with_vjp f f' x
rts/c/backends/hip.h view
@@ -444,6 +444,10 @@   opts[i++] = msgprintf("-DLOCKSTEP_WIDTH=%zu", ctx->lockstep_width);   opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_thread_block_size); +  if (getenv("HIP_PATH") != NULL) {+    opts[i++] = msgprintf("-I%s/include", getenv("HIP_PATH"));+  }+   for (int j = 0; extra_opts[j] != NULL; j++) {     opts[i++] = strdup(extra_opts[j]);   }
rts/c/scalar_f16.h view
@@ -9,7 +9,10 @@ // under emulation, so the compiler will have to be careful when // generating reads or writes. -#if !defined(cl_khr_fp16) && !(defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 600) && !(defined(ISPC))+#if !defined(cl_khr_fp16) && \+     !(defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 600) && \+     !(defined(__HIP_DEVICE_COMPILE__)) && \+     !(defined(ISPC)) #define EMULATE_F16 #endif @@ -30,6 +33,8 @@  #ifdef __CUDA_ARCH__ #include <cuda_fp16.h>+#elif defined(__HIP_DEVICE_COMPILE__)+#include <hip/hip_fp16.h> #endif  typedef half f16;@@ -41,8 +46,6 @@ SCALAR_FUN_ATTR f16 fadd16(f16 x, f16 y) { return x + y; } SCALAR_FUN_ATTR f16 fsub16(f16 x, f16 y) { return x - y; } SCALAR_FUN_ATTR f16 fmul16(f16 x, f16 y) { return x * y; }-SCALAR_FUN_ATTR bool cmplt16(f16 x, f16 y) { return x < y; }-SCALAR_FUN_ATTR bool cmple16(f16 x, f16 y) { return x <= y; } SCALAR_FUN_ATTR f16 sitofp_i8_f16(int8_t x) { return (f16) x; } SCALAR_FUN_ATTR f16 sitofp_i16_f16(int16_t x) { return (f16) x; } SCALAR_FUN_ATTR f16 sitofp_i32_f16(int32_t x) { return (f16) x; }@@ -66,6 +69,20 @@  SCALAR_FUN_ATTR bool futrts_isnan16(f16 x) { return isnan((float)x); } +// compare needs a custom implementation for HIP because of+// https://github.com/ROCm/clr/issues/274+#if defined(__HIP_DEVICE_COMPILE__)++SCALAR_FUN_ATTR bool cmplt16(f16 x, f16 y) { return __hlt(x, y); }+SCALAR_FUN_ATTR bool cmple16(f16 x, f16 y) { return __hle(x, y); }++#else++SCALAR_FUN_ATTR bool cmplt16(f16 x, f16 y) { return x < y; }+SCALAR_FUN_ATTR bool cmple16(f16 x, f16 y) { return x <= y; }++#endif+ #ifdef __OPENCL_VERSION__  SCALAR_FUN_ATTR f16 fabs16(f16 x) { return fabs(x); }@@ -80,7 +97,7 @@ SCALAR_FUN_ATTR f16 fmin16(f16 x, f16 y) { return futrts_isnan16(x) ? y : futrts_isnan16(y) ? x : min(x, y); } SCALAR_FUN_ATTR f16 fpow16(f16 x, f16 y) { return pow(x, y); } -#else // Assuming CUDA.+#else // Assuming CUDA or HIP.  SCALAR_FUN_ATTR f16 fabs16(f16 x) { return fabsf(x); } SCALAR_FUN_ATTR f16 fmax16(f16 x, f16 y) { return fmaxf(x, y); }@@ -303,6 +320,8 @@  #else // No native f16 - emulate. +SCALAR_FUN_ATTR bool cmplt16(f16 x, f16 y) { return x < y; }+SCALAR_FUN_ATTR bool cmple16(f16 x, f16 y) { return x <= y; } SCALAR_FUN_ATTR f16 fabs16(f16 x) { return fabs32(x); } SCALAR_FUN_ATTR f16 fmax16(f16 x, f16 y) { return fmax32(x, y); } SCALAR_FUN_ATTR f16 fmin16(f16 x, f16 y) { return fmin32(x, y); }
src/Futhark/AD/Fwd.hs view
@@ -3,21 +3,22 @@ module Futhark.AD.Fwd (fwdJVP) where  import Control.Monad-import Control.Monad.RWS.Strict+import Control.Monad.Identity+import Control.Monad.Reader import Control.Monad.State.Strict-import Data.Bifunctor (second)-import Data.List (transpose)+import Data.Bifunctor (bimap, second)+import Data.Foldable+import Data.Functor.Product import Data.List.NonEmpty (NonEmpty (..)) import Data.Map qualified as M+import Data.Tuple (Solo (..), getSolo) import Futhark.AD.Derivatives+import Futhark.AD.Shared import Futhark.Analysis.PrimExp.Convert import Futhark.Builder-import Futhark.Construct import Futhark.IR.SOACS--zeroTan :: Type -> ADM SubExp-zeroTan (Prim t) = pure $ constant $ blankPrimValue t-zeroTan t = error $ "zeroTan on non-primitive type: " ++ prettyString t+import Futhark.Tools+import Futhark.Util (interleave, splitAt3, unterleave)  zeroExp :: Type -> Exp SOACS zeroExp (Prim pt) =@@ -26,11 +27,18 @@   BasicOp $ Replicate shape $ Constant $ blankPrimValue pt zeroExp t = error $ "zeroExp: " ++ show t -tanType :: TypeBase s u -> ADM (TypeBase s u)+tanType :: (ArrayShape s, Monoid u) => TypeBase s u -> ADM (TypeBase s u) tanType (Acc acc ispace ts u) = do-  ts_tan <- mapM tanType ts-  pure $ Acc acc ispace (ts ++ ts_tan) u-tanType t = pure t+  acc_tan <- tangent acc+  tan_shape <- askShape+  pure $ Acc acc_tan (tan_shape <> ispace) ts u+tanType t = do+  shape <- askShape+  pure $ arrayOf (Prim (elemType t)) (shape `prependShape` arrayShape t) u+  where+    u = case t of+      Array _ _ u' -> u'+      _ -> mempty  slocal' :: ADM a -> ADM a slocal' = slocal id@@ -48,12 +56,18 @@     stateNameSource :: VNameSource   } -newtype ADM a = ADM (BuilderT SOACS (State RState) a)+data FEnv = FEnv+  { envTanShape :: Shape,+    envAttrs :: Attrs+  }++newtype ADM a = ADM (BuilderT SOACS (ReaderT FEnv (State RState)) a)   deriving     ( Functor,       Applicative,       Monad,       MonadState RState,+      MonadReader FEnv,       MonadFreshNames,       HasScope SOACS,       LocalScope SOACS@@ -72,12 +86,18 @@   getNameSource = gets stateNameSource   putNameSource src = modify (\env -> env {stateNameSource = src}) -runADM :: (MonadFreshNames m) => ADM a -> m a-runADM (ADM m) =+askShape :: ADM Shape+askShape = ADM $ lift $ asks envTanShape++runADM :: (MonadFreshNames m) => Shape -> Attrs -> ADM a -> m a+runADM shape attrs (ADM m) =   modifyNameSource $ \vn ->     second stateNameSource $       runState-        (fst <$> runBuilderT m mempty)+        ( runReaderT+            (fst <$> runBuilderT m mempty)+            (FEnv shape attrs)+        )         (RState mempty vn)  tanVName :: VName -> ADM VName@@ -89,27 +109,20 @@  class TanBuilder a where   newTan :: a -> ADM a-  bundleNew :: a -> ADM [a]+  bundleNew :: a -> ADM (a, a)  bundleNewList :: (TanBuilder a) => [a] -> ADM [a]-bundleNewList = fmap mconcat . mapM bundleNew+bundleNewList = fmap (uncurry interleave . unzip) . mapM bundleNew -instance TanBuilder (PatElem (TypeBase s u)) where-  newTan (PatElem p t)-    | isAcc t = do-        insertTan p p-        t' <- tanType t-        pure $ PatElem p t'-    | otherwise = do-        p' <- tanVName p-        insertTan p p'-        t' <- tanType t-        pure $ PatElem p' t'-  bundleNew pe@(PatElem _ t) = do+instance (ArrayShape s, Monoid u) => TanBuilder (PatElem (TypeBase s u)) where+  newTan (PatElem p t) = do+    p' <- tanVName p+    insertTan p p'+    t' <- tanType t+    pure $ PatElem p' t'+  bundleNew pe = do     pe' <- newTan pe-    if isAcc t-      then pure [pe']-      else pure [pe, pe']+    pure (pe, pe')  newTanPat :: (TanBuilder (PatElem t)) => Pat t -> ADM (Pat t) newTanPat (Pat pes) = Pat <$> mapM newTan pes@@ -117,41 +130,33 @@ bundleNewPat :: (TanBuilder (PatElem t)) => Pat t -> ADM (Pat t) bundleNewPat (Pat pes) = Pat <$> bundleNewList pes -instance TanBuilder (Param (TypeBase s u)) where+instance (ArrayShape s, Monoid u) => TanBuilder (Param (TypeBase s u)) where   newTan (Param _ p t) = do     PatElem p' t' <- newTan $ PatElem p t     pure $ Param mempty p' t'-  bundleNew param@(Param _ _ (Prim Unit)) =-    pure [param]-  bundleNew param@(Param _ _ t) = do+  bundleNew param = do     param' <- newTan param-    if isAcc t-      then pure [param']-      else pure [param, param']+    pure (param, param') -instance (Tangent a) => TanBuilder (Param (TypeBase s u), a) where+instance (TanBuilder a, Tangent b) => TanBuilder (a, b) where   newTan (p, x) = (,) <$> newTan p <*> tangent x   bundleNew (p, x) = do-    b <- bundleNew p+    p' <- newTan p     x_tan <- tangent x-    pure $ zip b [x, x_tan]+    pure ((p, x), (p', x_tan))  class Tangent a where   tangent :: a -> ADM a-  bundleTan :: a -> ADM [a]+  bundleTan :: a -> ADM (a, a) -instance Tangent (TypeBase s u) where+instance (ArrayShape s, Monoid u) => Tangent (TypeBase s u) where   tangent = tanType-  bundleTan t-    | isAcc t = do-        t' <- tangent t-        pure [t']-    | otherwise = do-        t' <- tangent t-        pure [t, t']+  bundleTan t = do+    t' <- tangent t+    pure (t, t')  bundleTangents :: (Tangent a) => [a] -> ADM [a]-bundleTangents = (mconcat <$>) . mapM bundleTan+bundleTangents = fmap (uncurry interleave . unzip) . mapM bundleTan  instance Tangent VName where   tangent v = do@@ -160,27 +165,121 @@       Just v_tan -> pure v_tan       Nothing -> do         t <- lookupType v-        letExp (baseName v <> "_implicit_tan") $ zeroExp t+        when (isAcc t) $+          error $+            "Missing tangent for accumulator " <> prettyString v+        tan_shape <- askShape+        letExp (baseName v <> "_implicit_tan") $ zeroExp $ t `arrayOfShape` tan_shape   bundleTan v = do-    t <- lookupType v-    if isAcc t-      then pure [v]-      else do-        v_tan <- tangent v-        pure [v, v_tan]+    v_tan <- tangent v+    pure (v, v_tan)  instance Tangent SubExp where-  tangent (Constant c) = zeroTan $ Prim $ primValueType c+  tangent (Constant c) = do+    tan_shape <- askShape+    if tan_shape == mempty+      then pure $ constant $ blankPrimValue pt+      else letSubExp "const_implicit_tan" $ zeroExp $ Prim pt `arrayOfShape` tan_shape+    where+      pt = primValueType c   tangent (Var v) = Var <$> tangent v   bundleTan c@Constant {} = do     c_tan <- tangent c-    pure [c, c_tan]-  bundleTan (Var v) = fmap Var <$> bundleTan v+    pure (c, c_tan)+  bundleTan (Var v) = bimap Var Var <$> bundleTan v  instance Tangent SubExpRes where   tangent (SubExpRes cs se) = SubExpRes cs <$> tangent se-  bundleTan (SubExpRes cs se) = map (SubExpRes cs) <$> bundleTan se+  bundleTan (SubExpRes cs se) = bimap (SubExpRes cs) (SubExpRes cs) <$> bundleTan se +withTan ::+  SubExp ->+  (SubExp -> ADM (Exp SOACS)) ->+  ADM (Exp SOACS)+withTan x f = do+  shape <- askShape+  x_tan <- tangent x+  mapNest shape (MkSolo x_tan) (f . getSolo)++withTansI ::+  VName ->+  [SubExp] ->+  ([SubExp] -> VName -> [SubExp] -> ADM (Exp SOACS)) ->+  ADM (Exp SOACS)+withTansI x ys f = do+  shape <- askShape+  x_tan <- tangent x+  ys_tan <- mapM tangent ys+  if shape == mempty+    then f [] x_tan ys_tan+    else do+      let w = shapeSize 0 shape+      ys_tan_vs <- mapM asVName ys_tan+      iota_p <- newParam "iota_p" $ Prim int64+      x_tan_p <- newParam "x_tanp" . rowType =<< lookupType x_tan+      ys_tan_ps <- mapM (newParam "y_tanp" . rowType <=< lookupType) ys_tan_vs+      lam <- mkLambda (iota_p : x_tan_p : ys_tan_ps) $ do+        fmap (subExpsRes . pure) . letSubExp "tan"+          =<< f+            [Var $ paramName iota_p]+            (paramName x_tan_p)+            (map (Var . paramName) ys_tan_ps)+      iota_v <- letExp "iota" $ iota64 w+      Op . Screma w (iota_v : x_tan : ys_tan_vs) <$> mapSOAC lam++withTans ::+  PrimType ->+  SubExp ->+  SubExp ->+  (PrimExp VName -> PrimExp VName -> PrimExp VName) ->+  ADM (Exp SOACS)+withTans t x y f = do+  shape <- askShape+  x_tan <- tangent x+  y_tan <- tangent y+  mapNest shape (Pair (Identity x_tan) (Identity y_tan)) $ \xy -> do+    Pair (Identity x_tan_v) (Identity y_tan_v) <- traverse asVName xy+    toExp $ f (LeafExp x_tan_v t) (LeafExp y_tan_v t)++withAnyTans ::+  (Traversable f) =>+  f SubExp ->+  ([PrimExp VName] -> PrimExp VName) ->+  ADM (Exp SOACS)+withAnyTans xs f = do+  shape <- askShape+  xs_tan <- traverse tangent xs+  mapNest shape xs_tan $ \xs_tan' -> do+    xs_tan'' <- forM xs_tan' $ \se -> do+      ~(Prim t) <- subExpType se+      pure $ primExpFromSubExp t se+    toExp $ f $ toList xs_tan''++bindTanPat :: Pat Type -> StmAux () -> Exp SOACS -> ADM ()+bindTanPat pat_tan aux e = do+  attrs <- asks envAttrs+  auxing aux . attributing attrs . letBind pat_tan $ e++bindTan ::+  Pat Type ->+  StmAux () ->+  SubExp ->+  (SubExp -> ADM (Exp SOACS)) ->+  ADM ()+bindTan pat_tan aux x f = do+  bindTanPat pat_tan aux =<< withTan x f++bindTans ::+  Pat Type ->+  StmAux () ->+  PrimType ->+  SubExp ->+  SubExp ->+  (PrimExp VName -> PrimExp VName -> PrimExp VName) ->+  ADM ()+bindTans pat_tan aux t x y f = do+  bindTanPat pat_tan aux =<< withTans t x y f+ basicFwd :: Pat Type -> StmAux () -> BasicOp -> ADM () basicFwd pat aux op = do   pat_tan <- newTanPat pat@@ -192,136 +291,238 @@       se_tan <- tangent se       addStm $ Let pat_tan aux $ BasicOp $ Opaque opaqueop se_tan     ArrayLit ses t -> do+      tan_shape <- askShape       ses_tan <- mapM tangent ses-      addStm $ Let pat_tan aux $ BasicOp $ ArrayLit ses_tan t+      if tan_shape == mempty+        then+          addStm $ Let pat_tan aux $ BasicOp $ ArrayLit ses_tan t+        else do+          pat_tan_tr <- letExp "pat_tan_tr" $ BasicOp $ ArrayLit ses_tan $ t `arrayOfShape` tan_shape+          pat_tan_tr_t <- lookupType pat_tan_tr+          let perm = vecPerm tan_shape pat_tan_tr_t+          addStm $ Let pat_tan aux $ BasicOp $ Rearrange pat_tan_tr perm     UnOp unop x -> do       let t = unOpType unop           x_pe = primExpFromSubExp t x           dx = pdUnOp unop x_pe-      x_tan <- primExpFromSubExp t <$> tangent x-      auxing aux $ letBindNames (patNames pat_tan) <=< toExp $ x_tan ~*~ dx+      bindTan pat_tan aux x $ \x_tan ->+        toExp $ primExpFromSubExp t x_tan ~*~ dx     BinOp bop x y -> do       let t = binOpType bop-      x_tan <- primExpFromSubExp t <$> tangent x-      y_tan <- primExpFromSubExp t <$> tangent y-      let (wrt_x, wrt_y) =-            pdBinOp bop (primExpFromSubExp t x) (primExpFromSubExp t y)-      auxing aux $-        letBindNames (patNames pat_tan) <=< toExp $-          x_tan ~*~ wrt_x ~+~ y_tan ~*~ wrt_y-    CmpOp {} ->-      addStm $ Let pat_tan aux $ zeroExp $ Prim Bool-    ConvOp cop x -> do-      x_tan <- tangent x-      addStm $ Let pat_tan aux $ BasicOp $ ConvOp cop x_tan+      bindTans pat_tan aux t x y $ \x_tan y_tan ->+        let (wrt_x, wrt_y) =+              pdBinOp bop (primExpFromSubExp t x) (primExpFromSubExp t y)+         in x_tan ~*~ wrt_x ~+~ y_tan ~*~ wrt_y+    CmpOp {} -> do+      tan_shape <- askShape+      addStm $ Let pat_tan aux $ zeroExp $ Prim Bool `arrayOfShape` tan_shape+    ConvOp cop x ->+      bindTan pat_tan aux x $ \x_tan ->+        pure $ BasicOp $ ConvOp cop x_tan     Assert {} -> pure ()     Index arr slice -> do+      dims <- shapeDims <$> askShape       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Index arr_tan slice+      let slice' = Slice $ map sliceDim dims <> unSlice slice+      addStm $ Let pat_tan aux $ BasicOp $ Index arr_tan slice'     Update safety arr slice se -> do+      dims <- shapeDims <$> askShape       arr_tan <- tangent arr       se_tan <- tangent se-      addStm $ Let pat_tan aux $ BasicOp $ Update safety arr_tan slice se_tan+      let slice' = Slice $ map sliceDim dims <> unSlice slice+      addStm $ Let pat_tan aux $ BasicOp $ Update safety arr_tan slice' se_tan     Concat d (arr :| arrs) w -> do+      r <- shapeRank <$> askShape       arr_tan <- tangent arr       arrs_tans <- mapM tangent arrs-      addStm $ Let pat_tan aux $ BasicOp $ Concat d (arr_tan :| arrs_tans) w+      addStm $ Let pat_tan aux $ BasicOp $ Concat (d + r) (arr_tan :| arrs_tans) w     Manifest arr ds -> do+      r <- shapeRank <$> askShape       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Manifest arr_tan ds+      addStm . Let pat_tan aux . BasicOp $+        Manifest arr_tan ([0 .. r - 1] ++ map (+ r) ds)     Iota n _ _ it -> do-      addStm $ Let pat_tan aux $ BasicOp $ Replicate (Shape [n]) (intConst it 0)-    Replicate n x -> do-      x_tan <- tangent x-      addStm $ Let pat_tan aux $ BasicOp $ Replicate n x_tan-    Scratch t shape ->-      addStm $ Let pat_tan aux $ BasicOp $ Scratch t shape+      shape <- askShape+      addStm . Let pat_tan aux . BasicOp $+        Replicate (shape <> Shape [n]) (intConst it 0)+    Replicate n x ->+      bindTan pat_tan aux x $ \x_tan ->+        pure $ BasicOp $ Replicate n x_tan+    Scratch t shape -> do+      tan_shape <- askShape+      addStm $ Let pat_tan aux $ BasicOp $ Scratch t $ shapeDims tan_shape <> shape     Reshape arr reshape -> do+      shape <- askShape       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Reshape arr_tan reshape+      addStm $ Let pat_tan aux $ BasicOp $ Reshape arr_tan (newshapeInner shape reshape)     Rearrange arr perm -> do+      r <- shapeRank <$> askShape       arr_tan <- tangent arr-      addStm $ Let pat_tan aux $ BasicOp $ Rearrange arr_tan perm+      addStm . Let pat_tan aux . BasicOp $+        Rearrange arr_tan ([0 .. r - 1] <> map (+ r) perm)     _ -> error $ "basicFwd: Unsupported op " ++ prettyString op  fwdLambda :: Lambda SOACS -> ADM (Lambda SOACS)-fwdLambda l@(Lambda params ret body) =-  Lambda <$> bundleNewList params <*> bundleTangents ret <*> inScopeOf l (fwdBody body)+fwdLambda (Lambda params _ body) = do+  params' <- bundleNewList params+  mkLambda params' $ bodyBind =<< fwdBody body -fwdStreamLambda :: Lambda SOACS -> ADM (Lambda SOACS)-fwdStreamLambda l@(Lambda params ret body) =-  Lambda <$> ((take 1 params ++) <$> bundleNewList (drop 1 params)) <*> bundleTangents ret <*> inScopeOf l (fwdBody body)+fwdWithAccLambda :: [WithAccInput SOACS] -> Lambda SOACS -> ADM (Lambda SOACS)+fwdWithAccLambda inputs (Lambda params _ body) = do+  let (cert_params, acc_params) = splitAt (length inputs) params+  cert_params_tan <- replicateM (length inputs) $ newParam "acc_cert_tan" $ Prim Unit+  acc_params_tan <- zipWithM mkAccParam (map paramName cert_params_tan) inputs -interleave :: [a] -> [a] -> [a]-interleave xs ys = concat $ transpose [xs, ys]+  mkLambda (cert_params <> cert_params_tan <> acc_params <> acc_params_tan) $ do+    zipWithM_+      insertTan+      (map paramName (cert_params <> acc_params))+      (map paramName (cert_params_tan <> acc_params_tan))+    bodyBind =<< fwdBody body+  where+    mkAccParam c (shape, arrs, _) = do+      tan_shape <- askShape+      ts <- map (stripArray (shapeRank shape)) <$> mapM lookupType arrs+      newParam "acc_p_tan" $ Acc c (tan_shape <> shape) ts NoUniqueness -zeroFromSubExp :: SubExp -> ADM VName-zeroFromSubExp (Constant c) =-  letExp "zero" . BasicOp . SubExp . Constant $-    blankPrimValue (primValueType c)-zeroFromSubExp (Var v) = do-  t <- lookupType v-  letExp "zero" $ zeroExp t+fwdStreamLambda :: Int -> Lambda SOACS -> ADM (Lambda SOACS)+fwdStreamLambda num_accs (Lambda params _ body) = do+  tan_shape <- askShape+  let (chunk_params, acc_params, arr_params) = splitAt3 1 num_accs params+  acc_params' <- bundleNewList acc_params+  (arr_params', arr_params'_tan) <- mapAndUnzipM onArrParam arr_params+  let params' =+        chunk_params <> acc_params' <> interleave arr_params' arr_params'_tan+  mkLambda params' $ do+    zipWithM_ (trArrParamTan tan_shape) arr_params' arr_params'_tan+    (acc_res, map_res) <- fmap (splitAt (num_accs * 2)) . bodyBind =<< fwdBody body+    let (map_res_primal, map_res_tan) = unterleave map_res+    map_res_tan' <- mapM (trMapResTan tan_shape) map_res_tan+    pure $ acc_res <> interleave map_res_primal map_res_tan'+  where+    -- Array parameters need to be treated specially as the chunk parameter+    -- must always be outermost.+    onArrParam p = do+      shape <- askShape+      (p', p_tan) <- bundleNew p+      let perm = vecPerm shape $ paramType p_tan+      pure (p', p_tan {paramDec = rearrangeType perm (paramType p_tan)}) +    -- Put the tangent shape back in the outermost position.+    trArrParamTan tan_shape p p_tan = do+      let perm = rearrangeInverse $ vecPerm tan_shape $ paramType p_tan+      v <-+        letExp (baseName (paramName p_tan)) . BasicOp $+          Rearrange (paramName p_tan) perm+      insertTan (paramName p) v++    -- Put the chunk size back in the outermost position.+    trMapResTan tan_shape (SubExpRes cs ~(Var v)) = do+      v_t <- lookupType v+      let perm = vecPerm tan_shape v_t+      fmap varRes . certifying cs $ letExp (baseName v) . BasicOp $ Rearrange v perm++pushTanShape :: VName -> ADM VName+pushTanShape v = do+  tan_shape <- askShape+  v_t <- lookupType v+  if tan_shape == mempty || arrayShape v_t == tan_shape || isAcc v_t+    then pure v+    else do+      let perm = vecPerm tan_shape v_t+      letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm++soacInputsWithTangents :: [VName] -> ADM [VName]+soacInputsWithTangents xs = do+  xs_tans <- mapM (pushTanShape <=< tangent) xs+  pure $ interleave xs xs_tans++soacResPat :: Int -> Int -> Pat Type -> ADM (Pat Type, [(Pat Type, VName)])+soacResPat scan_res red_res (Pat pes) = do+  pes_tan <- mapM newTan pes+  bimap (Pat . interleave pes) mconcat . unzip <$> zipWithM tweakPatElem [0 ..] pes_tan+  where+    isRedRes i = i >= scan_res && i < scan_res + red_res+    tweakPatElem i pe@(PatElem v v_t) = do+      tan_shape <- askShape+      if isRedRes i || tan_shape == mempty || arrayShape v_t == tan_shape || isAcc v_t+        then pure (pe, [])+        else do+          let perm = vecPerm tan_shape v_t+          v' <- newName v+          pure (PatElem v' $ rearrangeType perm v_t, [(Pat [pe], v')])+ fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () fwdSOAC pat aux (Screma size xs (ScremaForm f scs reds post_lam)) = do-  pat' <- bundleNewPat pat-  xs' <- bundleTangents xs+  (pat', to_transpose) <- soacResPat (scanResults scs) (redResults reds) pat+  xs' <- soacInputsWithTangents xs   f' <- fwdLambda f   scs' <- mapM fwdScan scs   reds' <- mapM fwdRed reds   post_lam' <- fwdLambda post_lam   addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds' post_lam'+  tan_shape <- askShape+  forM_ to_transpose $ \(rpat, v) -> do+    v_t <- lookupType v+    let perm = rearrangeInverse $ vecPerm tan_shape v_t+    letBind rpat $ BasicOp $ Rearrange v perm   where+    zeroTans lam =+      mapM (letSubExp "zero" . zeroExp <=< tanType) $ lambdaReturnType lam+     fwdScan :: Scan SOACS -> ADM (Scan SOACS)     fwdScan sc = do       op' <- fwdLambda $ scanLambda sc-      neutral_tans <- mapM zeroFromSubExp $ scanNeutral sc+      neutral_tans <- zeroTans $ scanLambda sc       pure $         Scan-          { scanNeutral = scanNeutral sc `interleave` map Var neutral_tans,+          { scanNeutral = scanNeutral sc `interleave` neutral_tans,             scanLambda = op'           }     fwdRed :: Reduce SOACS -> ADM (Reduce SOACS)     fwdRed red = do       op' <- fwdLambda $ redLambda red-      neutral_tans <- mapM zeroFromSubExp $ redNeutral red+      neutral_tans <- zeroTans $ redLambda red       pure $         Reduce           { redComm = redComm red,             redLambda = op',-            redNeutral = redNeutral red `interleave` map Var neutral_tans+            redNeutral = redNeutral red `interleave` neutral_tans           }-fwdSOAC pat aux (Stream size xs nes lam) = do+fwdSOAC pat aux (Stream size xs accs lam) = do   pat' <- bundleNewPat pat-  lam' <- fwdStreamLambda lam-  xs' <- bundleTangents xs-  nes_tan <- mapM (fmap Var . zeroFromSubExp) nes-  let nes' = interleave nes nes_tan-  addStm $ Let pat' aux $ Op $ Stream size xs' nes' lam'+  lam' <- fwdStreamLambda (length accs) lam+  xs' <- soacInputsWithTangents xs+  accs_tan <- mapM (letSubExp "zero" . zeroExp <=< tanType <=< subExpType) accs+  let accs' = interleave accs accs_tan+  addStm $ Let pat' aux $ Op $ Stream size xs' accs' lam' fwdSOAC pat aux (Hist w arrs ops bucket_fun) = do-  pat' <- bundleNewPat pat+  -- TODO: this is probably not very efficient in the vector case as we end up+  -- with a dreadful update operator that involves arrays.+  (pat', to_transpose) <- soacResPat 0 0 pat   ops' <- mapM fwdHist ops   bucket_fun' <- fwdHistBucket bucket_fun-  arrs' <- bundleTangents arrs+  arrs' <- soacInputsWithTangents arrs   addStm $ Let pat' aux $ Op $ Hist w arrs' ops' bucket_fun'+  tan_shape <- askShape+  forM_ to_transpose $ \(rpat, v) -> do+    v_t <- lookupType v+    let perm = rearrangeInverse $ vecPerm tan_shape v_t+    letBind rpat $ BasicOp $ Rearrange v perm   where     n_indices = sum $ map (shapeRank . histShape) ops     fwdBodyHist (Body _ stms res) = buildBody_ $ do       mapM_ fwdStm stms       let (res_is, res_vs) = splitAt n_indices res       (res_is ++) <$> bundleTangents res_vs-    fwdHistBucket l@(Lambda params ret body) =-      let (r_is, r_vs) = splitAt n_indices ret-       in Lambda-            <$> bundleNewList params-            <*> ((r_is ++) <$> bundleTangents r_vs)-            <*> inScopeOf l (fwdBodyHist body)+    fwdHistBucket (Lambda params _ body) = do+      params' <- bundleNewList params+      mkLambda params' $ bodyBind =<< fwdBodyHist body      fwdHist :: HistOp SOACS -> ADM (HistOp SOACS)     fwdHist (HistOp shape rf dest nes op) = do-      dest' <- bundleTangents dest-      nes_tan <- mapM (fmap Var . zeroFromSubExp) nes+      dest' <- soacInputsWithTangents dest+      nes_tan <- mapM (letSubExp "zero" . zeroExp <=< tanType) $ lambdaReturnType op       op' <- fwdLambda op       pure $         HistOp@@ -344,20 +545,17 @@  fwdStm :: Stm SOACS -> ADM () fwdStm (Let pat aux (BasicOp (UpdateAcc safety acc i x))) = do-  pat' <- bundleNewPat pat-  x' <- bundleTangents x-  acc_tan <- tangent acc-  addStm $ Let pat' aux $ BasicOp $ UpdateAcc safety acc_tan i x'+  pat_tan <- newTanPat pat+  addStm $ Let pat aux $ BasicOp $ UpdateAcc safety acc i x+  addStm . Let pat_tan aux <=< withTansI acc x $ \is acc_tan x_tan' -> do+    pure $ BasicOp $ UpdateAcc safety acc_tan (is <> i) x_tan' fwdStm stm@(Let pat aux (BasicOp e)) = do   -- XXX: this has to be too naive.-  unless (any isAcc $ patTypes pat) $-    addStm stm+  unless (any isAcc $ patTypes pat) $ addStm stm   basicFwd pat aux e-fwdStm stm@(Let pat _ (Apply f args _ _))+fwdStm stm@(Let pat aux (Apply f args _ _))   | Just (ret, argts) <- M.lookup f builtInFunctions = do       addStm stm-      arg_tans <--        zipWith primExpFromSubExp argts <$> mapM (tangent . fst) args       pat_tan <- newTanPat pat       let arg_pes = zipWith primExpFromSubExp argts (map fst args)       case pdBuiltin f arg_pes of@@ -375,8 +573,10 @@                       _ -> error $ "fwdStm.convertTo: " ++ prettyString (f, tt, e_t)                 where                   e_t = primExpType e-          letBindNames (patNames pat_tan)-            =<< toExp (foldl1 (~+~) $ zipWith (~*~) (map (convertTo ret) arg_tans) derivs)++          auxing aux . letBind pat_tan <=< withAnyTans (map fst args) $+            \arg_tans' ->+              foldl1 (~+~) $ zipWith (~*~) (map (convertTo ret) arg_tans') derivs fwdStm (Let pat aux (Match ses cases defbody (MatchDec ret ifsort))) = do   cases' <- slocal' $ mapM (traverse fwdBody) cases   defbody' <- slocal' $ fwdBody defbody@@ -387,36 +587,35 @@   val_pats' <- bundleNewList val_pats   pat' <- bundleNewPat pat   body' <--    localScope (scopeOfFParams (map fst val_pats) <> scopeOfLoopForm loop) . slocal' $+    localScope (scopeOfFParams (map fst val_pats') <> scopeOfLoopForm loop) . slocal' $       fwdBody body   addStm $ Let pat' aux $ Loop val_pats' (WhileLoop v) body' fwdStm (Let pat aux (Loop val_pats loop@(ForLoop i it bound) body)) = do   pat' <- bundleNewPat pat   val_pats' <- bundleNewList val_pats   body' <--    localScope (scopeOfFParams (map fst val_pats) <> scopeOfLoopForm loop) . slocal' $+    localScope (scopeOfFParams (map fst val_pats') <> scopeOfLoopForm loop) . slocal' $       fwdBody body   addStm $ Let pat' aux $ Loop val_pats' (ForLoop i it bound) body' fwdStm (Let pat aux (WithAcc inputs lam)) = do-  inputs' <- forM inputs $ \(shape, arrs, op) -> do+  inputs_tan <- forM inputs $ \(shape, arrs, op) -> do     arrs_tan <- mapM tangent arrs+    tan_shape <- askShape     op' <- case op of       Nothing -> pure Nothing       Just (op_lam, nes) -> do-        nes_tan <- mapM (fmap Var . zeroFromSubExp) nes-        op_lam' <- fwdLambda op_lam-        case op_lam' of-          Lambda ps ret body -> do-            let op_lam'' = Lambda (removeIndexTans (shapeRank shape) ps) ret body-            pure $ Just (op_lam'', interleave nes nes_tan)-    pure (shape, arrs <> arrs_tan, op')+        -- We assume that op_lam has unit partial derivatives (i.e., is some+        -- kind of addition). This is the case for all WithAccs produced by VJP.+        lams <- mapM addLambda $ lambdaReturnType op_lam+        -- Horizontally fuse the lambdas to produce a single one.+        idx_params <- replicateM (shapeRank shape) $ newParam "idx" $ Prim int64+        let (xs, ys) = bimap concat concat $ unzip $ map (splitAt 1 . lambdaParams) lams+        op_lam' <- mkLambda (idx_params <> xs <> ys) $ mconcat <$> mapM (bodyBind . lambdaBody) lams+        pure $ Just (op_lam', nes)+    pure (tan_shape <> shape, arrs_tan, op')   pat' <- bundleNewPat pat-  lam' <- fwdLambda lam-  addStm $ Let pat' aux $ WithAcc inputs' lam'-  where-    removeIndexTans 0 ps = ps-    removeIndexTans i (p : _ : ps) = p : removeIndexTans (i - 1) ps-    removeIndexTans _ ps = ps+  lam' <- fwdWithAccLambda inputs lam+  addStm $ Let pat' aux $ WithAcc (interleave inputs inputs_tan) lam' fwdStm (Let pat aux (Op soac)) = fwdSOAC pat aux soac fwdStm stm =   error $ "unhandled forward mode AD for Stm: " ++ prettyString stm ++ "\n" ++ show stm@@ -431,10 +630,22 @@   mapM_ fwdStm stms   (res <>) <$> mapM tangent res -fwdJVP :: (MonadFreshNames m) => Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)-fwdJVP scope l@(Lambda params ret body) =-  runADM . localScope scope . inScopeOf l $ do+fwdJVP ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Shape ->+  Attrs ->+  Lambda SOACS ->+  m (Lambda SOACS)+fwdJVP scope shape attrs (Lambda params _ body) =+  runADM shape attrs . localScope scope $ do     params_tan <- mapM newTan params-    body_tan <- fwdBodyTansLast body-    ret_tan <- mapM tangent ret-    pure $ Lambda (params ++ params_tan) (ret <> ret_tan) body_tan+    mkLambda (params <> params_tan) $+      bodyBind =<< fwdBodyTansLast body++-- Note [Forward-Mode vector AD]+--+-- An primal variable of type 't' has a tangent of type '[tan_shape]t', where+-- 'tan_shape' is the vector shape (which may be empty in the non-vector case).+-- This requires some care for SOACs, which always map across the outermost+-- dimension: basically we have to transpose the inputs and the outputs.
src/Futhark/AD/Rev.hs view
@@ -9,21 +9,22 @@ module Futhark.AD.Rev (revVJP) where  import Control.Monad-import Control.Monad.Identity-import Data.List ((\\)) import Data.List.NonEmpty (NonEmpty (..)) import Data.Map qualified as M+import Data.Tuple import Futhark.AD.Derivatives+import Futhark.AD.Rev.Acc import Futhark.AD.Rev.Loop import Futhark.AD.Rev.Monad import Futhark.AD.Rev.SOAC+import Futhark.AD.Shared import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS import Futhark.Tools import Futhark.Transform.Rename import Futhark.Transform.Substitute-import Futhark.Util (chunks, takeLast)+import Futhark.Util (takeLast)  patName :: Pat Type -> ADM VName patName (Pat [pe]) = pure $ patElemName pe@@ -57,8 +58,12 @@     ConvOp op x -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do-        contrib <--          letExp "contrib" $ BasicOp $ ConvOp (flipConvOp op) $ Var pat_adj+        adj_shape <- askShape++        contrib <- letExp "convop_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $+          \(MkSolo pat_adj') ->+            pure $ BasicOp $ ConvOp (flipConvOp op) pat_adj'+         updateSubExpAdj x contrib     --     UnOp op x -> do@@ -66,12 +71,13 @@        returnSweepCode $ do         let t = unOpType op-        contrib <- do-          let x_pe = primExpFromSubExp t x-              pat_adj' = primExpFromSubExp t (Var pat_adj)-              dx = pdUnOp op x_pe-          letExp "contrib" <=< toExp $ pat_adj' ~*~ dx +        adj_shape <- askShape++        contrib <- letExp "unop_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $+          \(MkSolo pat_adj') ->+            toExp $ primExpFromSubExp t pat_adj' ~*~ pdUnOp op (primExpFromSubExp t x)+         updateSubExpAdj x contrib     --     BinOp op x y -> do@@ -82,10 +88,20 @@             (wrt_x, wrt_y) =               pdBinOp op (primExpFromSubExp t x) (primExpFromSubExp t y) -            pat_adj' = primExpFromSubExp t $ Var pat_adj+        adj_shape <- askShape -        adj_x <- letExp "binop_x_adj" <=< toExp $ pat_adj' ~*~ wrt_x-        adj_y <- letExp "binop_y_adj" <=< toExp $ pat_adj' ~*~ wrt_y+        adj_x <- letExp "binop_x_adj"+          <=< mapNest adj_shape (MkSolo (Var pat_adj))+          $ \(MkSolo pat_adj') ->+            let pat_adj'' = primExpFromSubExp t pat_adj'+             in toExp $ pat_adj'' ~*~ wrt_x++        adj_y <- letExp "binop_y_adj"+          <=< mapNest adj_shape (MkSolo (Var pat_adj))+          $ \(MkSolo pat_adj') ->+            let pat_adj'' = primExpFromSubExp t pat_adj'+             in toExp $ pat_adj'' ~*~ wrt_y+         updateSubExpAdj x adj_x         updateSubExpAdj y adj_y     --@@ -127,10 +143,10 @@     --     Rearrange arr perm -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      r <- shapeRank <$> askShape       returnSweepCode $-        void $-          updateAdj arr <=< letExp "adj_rearrange" . BasicOp $-            Rearrange pat_adj (rearrangeInverse perm)+        void . updateAdj arr <=< letExp "adj_rearrange" . BasicOp $+          Rearrange pat_adj ([0 .. r - 1] <> map (+ r) (rearrangeInverse perm))     --     Replicate (Shape []) (Var se) -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m@@ -189,34 +205,18 @@     Update safety arr slice v -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do-        v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj slice+        adj_shape <- askShape+        let adj_slice = Slice $ map sliceDim (shapeDims adj_shape) ++ unSlice slice+        v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj adj_slice         v_adj_copy <- copyIfArray v_adj         updateSubExpAdj v v_adj_copy-        zeroes <- letSubExp "update_zero" . zeroExp =<< subExpType v+        v_adj_t <- lookupType v_adj+        zeroes <- letSubExp "update_zero" $ zeroExp v_adj_t         void $           updateAdj arr-            =<< letExp "update_src_adj" (BasicOp $ Update safety pat_adj slice zeroes)-    -- See Note [Adjoints of accumulators]-    UpdateAcc safety _ is vs -> do-      addStm $ Let pat aux $ BasicOp e-      m-      pat_adjs <- mapM lookupAdjVal (patNames pat)-      returnSweepCode $ do-        forM_ (zip pat_adjs vs) $ \(adj, v) -> do-          adj_t <- lookupType adj-          let index_adj = pure $ BasicOp $ Index adj $ fullSlice adj_t $ map DimFix is-          adj_i <--            letExp "updateacc_val_adj" =<< case safety of-              Unsafe ->-                index_adj-              Safe ->-                -- The primal UpdateAcc may be out-of-bounds, in which case-                -- indexing the adjoint is dangerous.-                eIf-                  (eShapeInBounds (arrayShape adj_t) (map eSubExp is))-                  (eBody [index_adj])-                  (eBody [pure $ zeroExp $ stripArray (length is) adj_t])-          updateSubExpAdj v adj_i+            =<< letExp "update_src_adj" (BasicOp $ Update safety pat_adj adj_slice zeroes)+    UpdateAcc safety acc is vs ->+      diffUpdateAcc pat aux safety acc is vs m     --     UserParam {} ->       void $ commonBasicOp pat aux e m@@ -225,30 +225,10 @@ vjpOps =   VjpOps     { vjpLambda = diffLambda,-      vjpStm = diffStm+      vjpStm = diffStm,+      vjpBody = diffBody     } --- | Transform updates on accumulators matching the given certificates into--- updates that write provided zero values.-zeroOutUpdates :: [(VName, [SubExp])] -> Lambda SOACS -> Lambda SOACS-zeroOutUpdates certs_to_zeroes lam = lam {lambdaBody = onBody $ lambdaBody lam}-  where-    onExp = runIdentity . mapExpM mapper-      where-        mapper =-          (identityMapper :: (Monad m) => Mapper SOACS SOACS m)-            { mapOnOp = traverseSOACStms (\_ stms -> pure $ onStms stms),-              mapOnBody = \_ body -> pure $ onBody body-            }-    onStms = fmap onStm-    onStm (Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is _)))-      | Acc c _ _ _ <- patElemType pe,-        Just zero <- lookup c certs_to_zeroes =-          Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is zero))-    onStm (Let pat aux e) = Let pat aux $ onExp e--    onBody body = body {bodyStms = onStms $ bodyStms body}- diffStm :: Stm SOACS -> ADM () -> ADM () diffStm (Let pat aux (BasicOp e)) m =   diffBasicOp pat aux e m@@ -259,7 +239,6 @@        pat_adj <- lookupAdjVal =<< patName pat       let arg_pes = zipWith primExpFromSubExp argts (map fst args)-          pat_adj' = primExpFromSubExp ret (Var pat_adj)           convert ft tt             | ft == tt = id           convert (IntType ft) (IntType tt) = ConvOpExp (SExt ft tt)@@ -268,13 +247,17 @@           convert (FloatType ft) Bool = ConvOpExp (FToB ft)           convert ft tt = error $ "diffStm.convert: " ++ prettyString (f, ft, tt) +      adj_shape <- askShape+       contribs <-         case pdBuiltin f arg_pes of           Nothing ->             error $ "No partial derivative defined for builtin function: " ++ prettyString f           Just derivs ->             forM (zip derivs argts) $ \(deriv, argt) ->-              letExp "contrib" <=< toExp . convert ret argt $ pat_adj' ~*~ deriv+              letExp "apply_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $+                \(MkSolo pat_adj') ->+                  toExp $ convert ret argt $ primExpFromSubExp ret pat_adj' ~*~ deriv        zipWithM_ updateSubExpAdj (map fst args) contribs diffStm stm@(Let pat _ (Match ses cases defbody _)) m = do@@ -311,35 +294,8 @@ diffStm (Let pat aux loop@Loop {}) m =   diffLoop diffStms pat aux loop m -- See Note [Adjoints of accumulators]-diffStm stm@(Let pat _aux (WithAcc inputs lam)) m = do-  addStm stm-  m-  returnSweepCode $ do-    adjs <- mapM lookupAdj $ patNames pat-    lam' <- renameLambda lam-    free_vars <- filterM isActive $ namesToList $ freeIn lam'-    free_accs <- filterM (fmap isAcc . lookupType) free_vars-    let free_vars' = free_vars \\ free_accs-    lam'' <- diffLambda' adjs free_vars' lam'-    (inputs_zeroes, inputs') <--      unzip <$> zipWithM renameInputLambda (chunks lengths adjs) inputs-    let certs = map paramName $ take (length inputs) $ lambdaParams lam''-    free_adjs <- letTupExp "with_acc_contrib" $ WithAcc inputs' $ zeroOutUpdates (zip certs inputs_zeroes) lam''-    zipWithM_ insAdj (arrs <> free_vars') free_adjs-  where-    lengths = map (\(_, as, _) -> length as) inputs-    arrs = concatMap (\(_, as, _) -> as) inputs-    renameInputLambda as_adj (shape, as, _) = do-      nes_ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) as-      zeroes <- mapM (zeroArray mempty) nes_ts-      as' <- mapM adjVal as_adj-      pure (map Var zeroes, (shape, as', Nothing))-    diffLambda' res_adjs get_adjs_for (Lambda params ts body) = do-      localScope (scopeOfLParams params) $ do-        Body () stms res <- diffBody res_adjs get_adjs_for body-        let body' = Body () stms $ take (length inputs) res <> takeLast (length get_adjs_for) res-        ts' <- mapM lookupType get_adjs_for-        pure $ Lambda params (take (length inputs) ts <> ts') body'+diffStm (Let pat aux (WithAcc inputs lam)) m =+  diffWithAcc vjpOps pat aux inputs lam m diffStm stm _ = error $ "diffStm unhandled:\n" ++ prettyString stm  diffStms :: Stms SOACS -> ADM ()@@ -371,17 +327,24 @@  diffLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS) diffLambda res_adjs get_adjs_for (Lambda params _ body) =-  localScope (scopeOfLParams params) $ do-    Body () stms res <- diffBody res_adjs get_adjs_for body-    let body' = Body () stms $ takeLast (length get_adjs_for) res-    ts' <- mapM lookupType get_adjs_for-    pure $ Lambda params ts' body'+  mkLambda params $ do+    res <- bodyBind =<< diffBody res_adjs get_adjs_for body+    pure $ takeLast (length get_adjs_for) res -revVJP :: (MonadFreshNames m) => Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)-revVJP scope (Lambda params ts body) =-  runADM . localScope (scope <> scopeOfLParams params) $ do+revVJP ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Shape ->+  Attrs ->+  Lambda SOACS ->+  m (Lambda SOACS)+revVJP scope shape attrs (Lambda params ts body) = do+  runADM shape attrs . localScope (scope <> scopeOfLParams params) $ do+    adj_shape <- askShape     params_adj <- forM (zip (map resSubExp (bodyResult body)) ts) $ \(se, t) ->-      Param mempty <$> maybe (newVName "const_adj") adjVName (subExpVar se) <*> pure t+      Param mempty+        <$> maybe (newVName "const_res_adj") adjVName (subExpVar se)+        <*> pure (t `arrayOfShape` adj_shape)      body' <-       localScope (scopeOfLParams params_adj) $@@ -391,143 +354,3 @@           body      pure $ Lambda (params ++ params_adj) (ts <> map paramType params) body'---- Note [Adjoints of accumulators]------ The general case of taking adjoints of WithAcc is tricky.  We make--- some assumptions and lay down a basic design.------ First, we assume that any WithAccs that occur in the program are--- come from one of these sources:------ - A previous instance of VJP, which means we can rely on the operator having---   a constant adjoint (it's addition as appropriate to the type).------ - A scatter, meaning there is no operator.------ (These can actually be distinguished by the presence of an operator, although--- we do not currently bother.)------ Second, the adjoint of an accumulator is an array of the same type--- as the underlying array.  For example, the adjoint type of the--- primal type 'acc(c, [n], {f64})' is '[n]f64'.  In principle the--- adjoint of 'acc(c, [n], {f64,f32})' should be two arrays of type--- '[]f64', '[]f32'.  Our current design assumes that adjoints are--- single variables.  This is fixable.------ In the return sweep, when inserting the with_acc, we still compute the--- "original" accumulator result, but modified such that its initial value is--- the adjoint of the result of the accumulator. We also modify the update_accs--- of these accumulators to be with zero values. This means that the array that--- is produced will be equal to the adjoint of the result, except for those--- places that have been updated, where it will be zero. This is intuitively--- sensible - values that have been overwritten (and so do not contribute to the--- result) should obviously have zero sensitivity.------ # Adjoint of UpdateAcc------ Consider primal code------     update_acc(acc, i, v)------ Interpreted as an imperative statement, this means------     acc[i] ⊕= v------ for some '⊕'.  Normally all the compiler knows of '⊕' is that it--- is associative and commutative, but because we assume that all--- accumulators are the result of previous AD transformations, we--- can assume that '⊕' actually behaves like addition - that is, has--- unit partial derivatives.  So the return sweep is------     v_adj += acc_adj[i]------ Further, we modify the primal code so that it becomes------     update_acc(acc, i, 0)------ for some appropriate notion of zero.------ # Adjoint of Map------ Suppose we have primal code------   let acc' =---     map (...) acc------ where "acc : acc(c, [n], {f64})" and the width of the Map is "w".--- Our normal transformation for Map input arrays is to similarly map--- their adjoint, but clearly this doesn't work here because the--- semantics of mapping an adjoint is an "implicit replicate".  So--- when generating the return sweep we actually perform that--- replication:------   map (...) (replicate w acc_adj)------ But what about the contributions to "acc'"?  Those we also have to--- take special care of.  The result of the map itself is actually a--- multidimensional array:------   let acc_contribs =---     map (...) (replicate w acc'_adj)------ which we must then sum to add to the contribution.------   acc_adj += sum(acc_contribs)------ I'm slightly worried about the asymptotics of this, since my--- intuition of this is that the contributions might be rather sparse.--- (Maybe completely zero?  If so it will be simplified away--- entirely.)  Perhaps a better solution is to treat--- accumulator-inputs in the primal code as we do free variables, and--- create accumulators for them in the return sweep.------ # Consumption------ A minor problem is that our usual way of handling consumption (Note--- [Consumption]) is not viable, because accumulators are not--- copyable.  Fortunately, while the accumulators that are consumed in--- the forward sweep will also be present in the return sweep given--- our current translation rules, they will be dead code.  As long as--- we are careful to run dead code elimination after revVJP, we should--- be good.---- Note [Array Adjoints of Match]------ Some unusual, but sadly not completely contrived, contain Match--- expressions that return multiple arrays, and there the arrays--- returned by one branch have overlapping aliases with another--- branch, although in different places. As an example consider this:------   let (X,Y) = if c---               then (A, B)---               else (B, A)------ Because our aliasing representation cannot express mutually--- exclusive aliases, we will consider X and Y to be aliased to each--- other. In practice, this means it is unlikely for X or Y to be--- consumed, because it would also consume the other (although it's--- possible for carefully written code).------ When producing adjoints for this, it will be something like------   let (X_adj,Y_adj) = if c---                       then (A_adj, B_adj)---                       else (B_adj, A_adj)------ which completely reflects the primal code. However, while it is--- unlikely that any consumption takes place for the original primal--- variables, it is almost guaranteed that X_adj and Y_adj will be--- consumed (that is the main way we use adjoints after all), and due--- to the conservative aliasing, when one is consumed, so is the--- other! To avoid this tragic fate, we are forced to copy any--- array-typed adjoints returned by a Match. This can be quite costly.--- However:------ 1) Futhark has pretty OK copy removal, so maybe it can get rid of---    these by using information not available to the AD pass.------ 2) In many cases, arrays will have accumulator adjoints, which are---    not subject to this problem.------ Issue #2228 was caused by neglecting to do this.
+ src/Futhark/AD/Rev/Acc.hs view
@@ -0,0 +1,396 @@+-- | Differentiation related to accumulators in the input program.+module Futhark.AD.Rev.Acc+  ( diffWithAcc,+    diffUpdateAcc,+  )+where++-- Note [Adjoints of accumulators]+--+-- The general case of taking adjoints of WithAcc is tricky.  We make+-- some assumptions and lay down a basic design.+--+-- First, we assume that any WithAccs that occur in the program are+-- come from one of these sources:+--+-- - A previous instance of VJP, which means we can rely on the operator having+--   a constant adjoint (it's addition as appropriate to the type).+--+-- - A scatter, meaning there is no operator.+--+-- (These can actually be distinguished by the presence of an operator, although+-- we do not currently bother.)+--+-- Second, the adjoint of an accumulator is an array of the same type+-- as the underlying array.  For example, the adjoint type of the+-- primal type 'acc(c, [n], {f64})' is '[n]f64'.  In principle the+-- adjoint of 'acc(c, [n], {f64,f32})' should be two arrays of type+-- '[]f64', '[]f32'.  Our current design assumes that adjoints are+-- single variables.  This is fixable.+--+-- In the return sweep, when inserting the with_acc, we still compute the+-- "original" accumulator result, but modified such that its initial value is+-- the adjoint of the result of the accumulator. We also modify the update_accs+-- of these accumulators to be with zero values. This means that the array that+-- is produced will be equal to the adjoint of the result, except for those+-- places that have been updated, where it will be zero. This is intuitively+-- sensible - values that have been overwritten (and so do not contribute to the+-- result) should obviously have zero sensitivity.+--+-- # Adjoint of UpdateAcc+--+-- Consider primal code+--+--     update_acc(acc, i, v)+--+-- Interpreted as an imperative statement, this means+--+--     acc[i] ⊕= v+--+-- for some '⊕'.  Normally all the compiler knows of '⊕' is that it+-- is associative and commutative, but because we assume that all+-- accumulators are the result of previous AD transformations, we+-- can assume that '⊕' actually behaves like addition - that is, has+-- unit partial derivatives.  So the return sweep is+--+--     v_adj += acc_adj[i]+--+-- Further, we modify the primal code so that it becomes+--+--     update_acc(acc, i, 0)+--+-- for some appropriate notion of zero.+--+-- # Adjoint of Map+--+-- Suppose we have primal code+--+--   let acc' =+--     map (...) acc+--+-- where "acc : acc(c, [n], {f64})" and the width of the Map is "w".+-- Our normal transformation for Map input arrays is to similarly map+-- their adjoint, but clearly this doesn't work here because the+-- semantics of mapping an adjoint is an "implicit replicate".  So+-- when generating the return sweep we actually perform that+-- replication:+--+--   map (...) (replicate w acc_adj)+--+-- But what about the contributions to "acc'"?  Those we also have to+-- take special care of.  The result of the map itself is actually a+-- multidimensional array:+--+--   let acc_contribs =+--     map (...) (replicate w acc'_adj)+--+-- which we must then sum to add to the contribution.+--+--   acc_adj += sum(acc_contribs)+--+-- I'm slightly worried about the asymptotics of this, since my+-- intuition of this is that the contributions might be rather sparse.+-- (Maybe completely zero?  If so it will be simplified away+-- entirely.)  Perhaps a better solution is to treat+-- accumulator-inputs in the primal code as we do free variables, and+-- create accumulators for them in the return sweep.+--+-- # Vectorised WithAcc+--+-- When WithAcc occurs in vectorised AD, the accumulator element types gain+-- extra leading "vectorised" dimensions corresponding to the enclosing vector+-- shape. For example, if the primal type inside a map of width @w@ is @acc(c,+-- [n], {f64})@, the adjoint type is @[w][n]f64@ -- but the internal accumulator+-- layout expects shape @[n][w]f64@ (the accumulator shape comes first, then the+-- vectorised dimensions, then element dimensions).+--+-- This means we must transpose accumulator adjoints when entering and+-- leaving the return-sweep WithAcc:+--+--  * On entry: transpose result adjoints from @[vec...][shape...]elem@ to+--    @[shape...][vec...]elem@ so they can serve as initial values for the+--    accumulators.+--+--  * On exit: transpose the produced arrays back from @[shape...][vec...]elem@+--    to @[vec...][shape...]elem@ to match the expected adjoint layout.+--+-- This is actually quite similar to how other SOACs must be handled.+--+-- Additionally, the accumulator parameter types in the lambda (and any+-- Acc-typed pattern elements or inner lambda parameters referring to the same+-- certs) must be updated to reflect the vectorised element types *before*+-- differentiation. This ensures that 'lookupAdj' on accumulator variables+-- inside the lambda produces adjoints with the correct vectorised type.+--+-- The UpdateAcc case is simpler under vectorisation: because the accumulator+-- adjoint already has the vectorised dimensions folded into its element type, a+-- plain index into the adjoint at the update indices directly yields the+-- correctly-shaped contribution.+--+-- # Consumption+--+-- A minor problem is that our usual way of handling consumption (Note+-- [Consumption]) is not viable, because accumulators are not+-- copyable.  Fortunately, while the accumulators that are consumed in+-- the forward sweep will also be present in the return sweep given+-- our current translation rules, they will be dead code.  As long as+-- we are careful to run dead code elimination after revVJP, we should+-- be good.++-- Note [Array Adjoints of Match]+--+-- Some unusual, but sadly not completely contrived, contain Match+-- expressions that return multiple arrays, and there the arrays+-- returned by one branch have overlapping aliases with another+-- branch, although in different places. As an example consider this:+--+--   let (X,Y) = if c+--               then (A, B)+--               else (B, A)+--+-- Because our aliasing representation cannot express mutually+-- exclusive aliases, we will consider X and Y to be aliased to each+-- other. In practice, this means it is unlikely for X or Y to be+-- consumed, because it would also consume the other (although it's+-- possible for carefully written code).+--+-- When producing adjoints for this, it will be something like+--+--   let (X_adj,Y_adj) = if c+--                       then (A_adj, B_adj)+--                       else (B_adj, A_adj)+--+-- which completely reflects the primal code. However, while it is+-- unlikely that any consumption takes place for the original primal+-- variables, it is almost guaranteed that X_adj and Y_adj will be+-- consumed (that is the main way we use adjoints after all), and due+-- to the conservative aliasing, when one is consumed, so is the+-- other! To avoid this tragic fate, we are forced to copy any+-- array-typed adjoints returned by a Match. This can be quite costly.+-- However:+--+-- 1) Futhark has pretty OK copy removal, so maybe it can get rid of+--    these by using information not available to the AD pass.+--+-- 2) In many cases, arrays will have accumulator adjoints, which are+--    not subject to this problem.+--+-- Issue #2228 was caused by neglecting to do this.++import Control.Monad+import Control.Monad.Identity+import Data.List ((\\))+import Futhark.AD.Rev.Monad+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Util (chunks, takeLast)++-- | Transform updates on accumulators matching the given certificates into+-- updates that write provided zero values.+zeroOutUpdates :: [(VName, [SubExp])] -> Lambda SOACS -> Lambda SOACS+zeroOutUpdates certs_to_zeroes lam = lam {lambdaBody = onBody $ lambdaBody lam}+  where+    onExp = runIdentity . mapExpM mapper+      where+        mapper =+          (identityMapper :: (Monad m) => Mapper SOACS SOACS m)+            { mapOnOp = traverseSOACStms (\_ stms -> pure $ onStms stms),+              mapOnBody = \_ body -> pure $ onBody body+            }+    onStms = fmap onStm+    onStm (Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is _)))+      | Acc c _ _ _ <- patElemType pe,+        Just zero <- lookup c certs_to_zeroes =+          Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is zero))+    onStm (Let pat aux e) = Let pat aux $ onExp e++    onBody body = body {bodyStms = onStms $ bodyStms body}++-- Update accumulator parameter types in the lambda to include vectorised+-- element types. Also updates all Acc-typed pattern elements and inner+-- lambda parameters that reference the same accumulator certs.+updateAccParamTypes :: Int -> Shape -> Lambda SOACS -> Lambda SOACS+updateAccParamTypes n_inputs adj_sh lam+  | adj_sh == mempty = lam+  | otherwise =+      let (cert_ps, rest_ps) = splitAt n_inputs (lambdaParams lam)+          (acc_ps, other_ps) = splitAt n_inputs rest_ps+          acc_ps' = map (updateParam cert_names) acc_ps+          cert_names = map paramName cert_ps+          body' = updateBody cert_names (lambdaBody lam)+          ret' = map (updateAccType cert_names) (lambdaReturnType lam)+       in lam+            { lambdaParams = cert_ps ++ acc_ps' ++ other_ps,+              lambdaReturnType = ret',+              lambdaBody = body'+            }+  where+    updateParam :: [VName] -> Param Type -> Param Type+    updateParam certs p =+      p {paramDec = updateAccType certs (paramDec p)}++    updateAccType :: [VName] -> Type -> Type+    updateAccType certs (Acc cert acc_shape ts u)+      | cert `elem` certs =+          Acc cert acc_shape (map (`arrayOfShape` adj_sh) ts) u+    updateAccType _ t = t++    updateBody :: [VName] -> Body SOACS -> Body SOACS+    updateBody certs body =+      body {bodyStms = fmap (updateStm certs) (bodyStms body)}++    updateStm :: [VName] -> Stm SOACS -> Stm SOACS+    updateStm certs (Let pat aux e) =+      Let (updatePat certs pat) aux (updateExp certs e)++    updatePat :: [VName] -> Pat Type -> Pat Type+    updatePat certs (Pat pes) =+      Pat $ map (\pe -> pe {patElemDec = updateAccType certs (patElemDec pe)}) pes++    updateExp :: [VName] -> Exp SOACS -> Exp SOACS+    updateExp certs = runIdentity . mapExpM mapper+      where+        mapper =+          (identityMapper :: (Monad m) => Mapper SOACS SOACS m)+            { mapOnBody = \_ b -> pure $ updateBody certs b,+              mapOnOp = pure . updateSOAC certs+            }++    updateSOAC :: [VName] -> SOAC SOACS -> SOAC SOACS+    updateSOAC certs = runIdentity . mapSOACM mapper+      where+        mapper =+          identitySOACMapper+            { mapOnSOACLambda = pure . updateLambda certs+            }++    updateLambda :: [VName] -> Lambda SOACS -> Lambda SOACS+    updateLambda certs l =+      l+        { lambdaParams = map (updateParam certs) (lambdaParams l),+          lambdaReturnType = map (updateAccType certs) (lambdaReturnType l),+          lambdaBody = updateBody certs (lambdaBody l)+        }++diffWithAcc ::+  VjpOps ->+  Pat Type ->+  StmAux () ->+  [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] ->+  Lambda SOACS ->+  ADM () ->+  ADM ()+diffWithAcc ops pat aux inputs lam m = do+  addStm $ Let pat aux $ WithAcc inputs lam+  m+  returnSweepCode $ do+    adj_shape <- askShape+    adjs <- mapM lookupAdj $ patNames pat+    -- Transpose the accumulator result adjoints from [vec...][shape...]elem+    -- to [shape...][vec...]elem, matching the internal accumulator layout.+    adjs' <- transposeAdjs adj_shape adjs+    lam' <- renameLambda lam+    -- Update the lambda's accumulator parameter types to reflect vectorised+    -- element types BEFORE differentiation, so that lookupAdj on Acc variables+    -- inside the lambda gives the correct vectorised adjoint type.+    let lam'_vec = updateAccParamTypes n_inputs adj_shape lam'+    free_vars <- filterM isActive $ namesToList $ freeIn lam'_vec+    free_accs <- filterM (fmap isAcc . lookupType) free_vars+    let free_vars' = free_vars \\ free_accs+    lam'' <- diffLambda' adjs' free_vars' lam'_vec+    (inputs_zeroes, inputs') <-+      unzip <$> zipWithM (renameInputLambda adj_shape) (chunks lengths adjs) inputs+    let certs = map paramName $ take n_inputs $ lambdaParams lam''+    raw_adjs <-+      letTupExp "with_acc_contrib" . WithAcc inputs' $+        zeroOutUpdates (zip certs inputs_zeroes) lam''+    -- The accumulator results have shape [shape...][vec...]elem. Transpose+    -- back to [vec...][shape...]elem for the adjoint.+    let n_arrs = sum lengths+        (arr_adjs, free_adjs) = splitAt n_arrs raw_adjs+    arr_adjs' <- zipWithM (transposeAccResult adj_shape) (map (\(s, _, _) -> s) inputs) arr_adjs+    zipWithM_ insAdj arrs arr_adjs'+    zipWithM_ insAdj free_vars' free_adjs+  where+    n_inputs = length inputs+    lengths = map (\(_, as, _) -> length as) inputs+    arrs = concatMap (\(_, as, _) -> as) inputs++    -- Transpose the accumulator-related adjoints from [vec...][shape...]elem+    -- to [shape...][vec...]elem. Non-accumulator adjs are left unchanged.+    transposeAdjs :: Shape -> [Adj] -> ADM [Adj]+    transposeAdjs adj_sh adjs+      | adj_sh == mempty = pure adjs+      | otherwise = do+          let n_arrs = sum lengths+              (acc_adjs, other_adjs) = splitAt n_arrs adjs+          acc_adjs' <- mapM transposeAdj acc_adjs+          pure $ acc_adjs' ++ other_adjs++    transposeAdj :: Adj -> ADM Adj+    transposeAdj adj = do+      v <- adjVal adj+      v' <- vecToInner v+      pure $ AdjVal $ Var v'++    -- Transpose [shape...][vec...][elem...] to [vec...][shape...][elem...]+    transposeAccResult :: Shape -> Shape -> VName -> ADM VName+    transposeAccResult adj_sh shape v+      | adj_sh == mempty = pure v+      | otherwise = do+          v_t <- lookupType v+          let r = shapeRank adj_sh+              s = shapeRank shape+              total = arrayRank v_t+              perm = [s .. s + r - 1] ++ [0 .. s - 1] ++ [s + r .. total - 1]+          letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm++    renameInputLambda adj_sh as_adj (shape, as, _) = do+      -- Compute element types with vectorised dimensions included.+      orig_nes_ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) as+      let vec_nes_ts = map (`arrayOfShape` adj_sh) orig_nes_ts+      zeroes <- mapM (zeroArray mempty) vec_nes_ts+      -- Transpose adjoints from [vec...][shape...]elem to [shape...][vec...]elem+      -- so they match the accumulator layout.+      as' <- mapM adjVal as_adj+      as'' <- mapM vecToInner as'+      pure (map Var zeroes, (shape, as'', Nothing))++    diffLambda' res_adjs get_adjs_for (Lambda params ts body) = do+      localScope (scopeOfLParams params) $ do+        Body () stms res <- vjpBody ops res_adjs get_adjs_for body+        let body' = Body () stms $ take n_inputs res <> takeLast (length get_adjs_for) res+        ts' <- mapM lookupType get_adjs_for+        pure $ Lambda params (take n_inputs ts <> ts') body'++diffUpdateAcc ::+  Pat Type ->+  StmAux () ->+  Safety ->+  VName ->+  [SubExp] ->+  [SubExp] ->+  ADM () ->+  ADM ()+diffUpdateAcc pat aux safety acc is vs m = do+  addStm $ Let pat aux $ BasicOp $ UpdateAcc safety acc is vs+  m+  pat_adjs <- mapM lookupAdjVal (patNames pat)+  returnSweepCode $ do+    forM_ (zip pat_adjs vs) $ \(adj, v) -> do+      adj_t <- lookupType adj+      let index_adj = pure $ BasicOp $ Index adj $ fullSlice adj_t $ map DimFix is+      adj_i <-+        letExp "updateacc_val_adj" =<< case safety of+          Unsafe ->+            index_adj+          Safe ->+            -- The primal UpdateAcc may be out-of-bounds, in which case+            -- indexing the adjoint is dangerous.+            eIf+              (eShapeInBounds (arrayShape adj_t) (map eSubExp is))+              (eBody [index_adj])+              (eBody [pure $ zeroExp $ stripArray (length is) adj_t])+      updateSubExpAdj v adj_i
src/Futhark/AD/Rev/Hist.hs view
@@ -241,69 +241,70 @@    m -  x_bar <- lookupAdjVal x+  locallyNonvector (x, dst, vs) $ do+    x_bar <- lookupAdjVal x -  x_ind_dst <- newParam (baseName x <> "_ind_param") $ Prim int64-  x_bar_dst <- newParam (baseName x <> "_bar_param") $ Prim t-  dst_lam_inner <--    mkLambda [x_ind_dst, x_bar_dst] $-      fmap varsRes . letTupExp "dst_bar"-        =<< eIf-          (toExp $ le64 (paramName x_ind_dst) .==. -1)-          (eBody $ pure $ eParam x_bar_dst)-          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)-  dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner+    x_ind_dst <- newParam (baseName x <> "_ind_param") $ Prim int64+    x_bar_dst <- newParam (baseName x <> "_bar_param") $ Prim t+    dst_lam_inner <-+      mkLambda [x_ind_dst, x_bar_dst] $+        fmap varsRes . letTupExp "dst_bar"+          =<< eIf+            (toExp $ le64 (paramName x_ind_dst) .==. -1)+            (eBody $ pure $ eParam x_bar_dst)+            (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+    dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner -  dst_bar <--    letExp (baseName dst <> "_bar") . Op . Screma w [x_inds, x_bar]-      =<< mapSOAC dst_lam+    dst_bar <-+      letExp (baseName dst <> "_bar") . Op . Screma w [x_inds, x_bar]+        =<< mapSOAC dst_lam -  updateAdj dst dst_bar+    updateAdj dst dst_bar -  vs_bar <- lookupAdjVal vs+    vs_bar <- lookupAdjVal vs -  inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []-  let inds = x_inds : inds'+    inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []+    let inds = x_inds : inds' -  par_x_ind_vs <- replicateM nr_dims $ newParam (baseName x <> "_ind_param") $ Prim int64-  par_x_bar_vs <- newParam (baseName x <> "_bar_param") $ Prim t-  vs_lam_inner <--    mkLambda (par_x_bar_vs : par_x_ind_vs) $-      fmap varsRes . letTupExp "res"-        =<< eIf-          (toExp $ le64 (paramName $ head par_x_ind_vs) .==. -1)-          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)-          ( eBody $-              pure $ do-                vs_bar_i <--                  letSubExp (baseName vs_bar <> "_el") . BasicOp $-                    Index vs_bar . Slice $-                      fmap (DimFix . Var . paramName) par_x_ind_vs-                eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)-          )-  vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner+    par_x_ind_vs <- replicateM nr_dims $ newParam (baseName x <> "_ind_param") $ Prim int64+    par_x_bar_vs <- newParam (baseName x <> "_bar_param") $ Prim t+    vs_lam_inner <-+      mkLambda (par_x_bar_vs : par_x_ind_vs) $+        fmap varsRes . letTupExp "res"+          =<< eIf+            (toExp $ le64 (paramName $ head par_x_ind_vs) .==. -1)+            (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+            ( eBody $+                pure $ do+                  vs_bar_i <-+                    letSubExp (baseName vs_bar <> "_el") . BasicOp $+                      Index vs_bar . Slice $+                        fmap (DimFix . Var . paramName) par_x_ind_vs+                  eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)+            )+    vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner -  vs_bar_p <--    letExp (baseName vs <> "_partial") . Op . Screma w (x_bar : inds)-      =<< mapSOAC vs_lam+    vs_bar_p <-+      letExp (baseName vs <> "_partial") . Op . Screma w (x_bar : inds)+        =<< mapSOAC vs_lam -  q <--    letSubExp "q"-      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) dst_dims+    q <-+      letSubExp "q"+        =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) dst_dims -  scatter_inps <- do-    -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]-    -- ToDo: Cosmin asks: is the below the correct translation of the line above?-    forM (inds ++ [vs_bar_p]) $ \v -> do-      v_t <- lookupType v-      letExp "flat" . BasicOp . Reshape v $-        reshapeAll (arrayShape v_t) (Shape [q])+    scatter_inps <- do+      -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]+      -- ToDo: Cosmin asks: is the below the correct translation of the line above?+      forM (inds ++ [vs_bar_p]) $ \v -> do+        v_t <- lookupType v+        letExp "flat" . BasicOp . Reshape v $+          reshapeAll (arrayShape v_t) (Shape [q]) -  vs_bar' <--    fmap head $-      doScatter (baseName vs <> "_bar") nr_dims [vs_bar] scatter_inps $-        pure . map (Var . paramName)-  insAdj vs vs_bar'+    vs_bar' <-+      fmap head $+        doScatter (baseName vs <> "_bar") nr_dims [vs_bar] scatter_inps $+          pure . map (Var . paramName)+    insAdj vs vs_bar'   where     mk_indices :: [SubExp] -> [SubExp] -> ADM [VName]     mk_indices [] _ = pure []@@ -403,8 +404,8 @@       fmap varsRes . letTupExp "h_part"         =<< eIf           (toExp $ 0 .==. le64 (paramName c_param))-          (eBody $ pure $ eParam p_param)-          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+          (eBody [eParam p_param])+          (eBody [eSubExp $ Constant $ blankPrimValue t])   lam_h_part <- nestedmap dst_dims [vs_elm_type, int64] lam_h_part_inner   h_part_res <- eLambda lam_h_part $ map (eSubExp . Var) [nz_prods, zr_counts]   h_part' <- bindSubExpRes "h_part" h_part_res@@ -418,60 +419,61 @@    m -  x_bar <- lookupAdjVal x+  locallyNonvector (x, dst, vs) $ do+    x_bar <- lookupAdjVal x -  lam_mul'' <- renameLambda lam_mul'-  dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]-  dst_bar <- bindSubExpRes (baseName dst <> "_bar") dst_bar_res-  updateAdj dst $ head dst_bar+    lam_mul'' <- renameLambda lam_mul'+    dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]+    dst_bar <- bindSubExpRes (baseName dst <> "_bar") dst_bar_res+    updateAdj dst $ head dst_bar -  lam_mul''' <- renameLambda lam_mul'-  part_bar_res <- eLambda lam_mul''' $ map (eSubExp . Var) [dst, x_bar]-  part_bar' <- bindSubExpRes "part_bar" part_bar_res-  let [part_bar] = part_bar'+    lam_mul''' <- renameLambda lam_mul'+    part_bar_res <- eLambda lam_mul''' $ map (eSubExp . Var) [dst, x_bar]+    part_bar' <- bindSubExpRes "part_bar" part_bar_res+    let [part_bar] = part_bar' -  inner_params <- zipWithM newParam ["zr_cts", "pr_bar", "nz_prd", "a"] $ map Prim [int64, t, t, t]-  let [zr_cts, pr_bar, nz_prd, a_param] = inner_params-  lam_vsbar_inner <--    mkLambda inner_params $-      fmap varsRes . letTupExp "vs_bar" =<< do-        eIf-          (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 0) (eParam zr_cts))-          (eBody $ pure $ eBinOp mul (eParam pr_bar) $ eBinOp (getBinOpDiv t) (eParam nz_prd) $ eParam a_param)-          ( eBody $-              pure $-                eIf-                  ( eBinOp-                      LogAnd-                      (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 1) (eParam zr_cts))-                      (eCmpOp (CmpEq t) (eSubExp $ Constant $ blankPrimValue t) $ eParam a_param)-                  )-                  (eBody $ pure $ eBinOp mul (eParam nz_prd) (eParam pr_bar))-                  (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)-          )+    inner_params <- zipWithM newParam ["zr_cts", "pr_bar", "nz_prd", "a"] $ map Prim [int64, t, t, t]+    let [zr_cts, pr_bar, nz_prd, a_param] = inner_params+    lam_vsbar_inner <-+      mkLambda inner_params $+        fmap varsRes . letTupExp "vs_bar" =<< do+          eIf+            (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 0) (eParam zr_cts))+            (eBody [eBinOp mul (eParam pr_bar) $ eBinOp (getBinOpDiv t) (eParam nz_prd) $ eParam a_param])+            ( eBody+                [ eIf+                    ( eBinOp+                        LogAnd+                        (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 1) (eParam zr_cts))+                        (eCmpOp (CmpEq t) (eSubExp $ Constant $ blankPrimValue t) $ eParam a_param)+                    )+                    (eBody [eBinOp mul (eParam nz_prd) (eParam pr_bar)])+                    (eBody [eSubExp $ Constant $ blankPrimValue t])+                ]+            ) -  lam_vsbar_middle <- nestedmap inner_dims [int64, t, t, t] lam_vsbar_inner+    lam_vsbar_middle <- nestedmap inner_dims [int64, t, t, t] lam_vsbar_inner -  i_param <- newParam "i" $ Prim int64-  a_param' <- newParam "a" $ rowType vs_type-  lam_vsbar <--    mkLambda [i_param, a_param'] $-      fmap varsRes . letTupExp "vs_bar"-        =<< eIf-          (toExp $ withinBounds $ pure (w, paramName i_param))-          ( buildBody_ $ do-              let i = fullSlice vs_type [DimFix $ Var $ paramName i_param]-              names <- traverse newVName ["zr_cts", "pr_bar", "nz_prd"]-              zipWithM_ (\name -> letBindNames [name] . BasicOp . flip Index i) names [zr_counts, part_bar, nz_prods]-              eLambda lam_vsbar_middle $ map (eSubExp . Var) names <> [eParam a_param']-          )-          (eBody $ pure $ pure $ zeroExp $ rowType dst_type)+    i_param <- newParam "i" $ Prim int64+    a_param' <- newParam "a" $ rowType vs_type+    lam_vsbar <-+      mkLambda [i_param, a_param'] $+        fmap varsRes . letTupExp "vs_bar"+          =<< eIf+            (toExp $ withinBounds $ pure (w, paramName i_param))+            ( buildBody_ $ do+                let i = fullSlice vs_type [DimFix $ Var $ paramName i_param]+                names <- traverse newVName ["zr_cts", "pr_bar", "nz_prd"]+                zipWithM_ (\name -> letBindNames [name] . BasicOp . flip Index i) names [zr_counts, part_bar, nz_prods]+                eLambda lam_vsbar_middle $ map (eSubExp . Var) names <> [eParam a_param']+            )+            (eBody [pure $ zeroExp $ rowType dst_type]) -  vs_bar <--    letExp (baseName vs <> "_bar") . Op . Screma n [is, vs]-      =<< mapSOAC lam_vsbar+    vs_bar <-+      letExp (baseName vs <> "_bar") . Op . Screma n [is, vs]+        =<< mapSOAC lam_vsbar -  updateAdj vs vs_bar+    updateAdj vs vs_bar  -- -- special case of histogram with add as operator.@@ -502,25 +504,25 @@    m -  x_bar <- lookupAdjVal x+  locallyNonvector (x, dst, vs) $ do+    x_bar <- lookupAdjVal x -  updateAdj dst x_bar+    updateAdj dst x_bar -  x_type <- lookupType x-  i_param <- newParam (baseName vs <> "_i") $ Prim int64-  let i = paramName i_param-  lam_vsbar <--    mkLambda [i_param] $-      fmap varsRes . letTupExp "vs_bar"-        =<< eIf-          (toExp $ withinBounds $ pure (w, i))-          (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i])-          (eBody $ pure $ eSubExp ne)+    i_param <- newParam (baseName vs <> "_i") $ Prim int64+    let i = paramName i_param+    lam_vsbar <-+      mkLambda [i_param] $+        fmap varsRes . letTupExp "vs_bar"+          =<< eIf+            (toExp $ withinBounds $ pure (w, i))+            (eBody [eIndex x_bar [eVar i]])+            (eBody [eSubExp ne]) -  vs_bar <--    letExp (baseName vs <> "_bar") . Op . Screma n [is]-      =<< mapSOAC lam_vsbar-  updateAdj vs vs_bar+    vs_bar <-+      letExp (baseName vs <> "_bar") . Op . Screma n [is]+        =<< mapSOAC lam_vsbar+    updateAdj vs vs_bar  -- Special case for vectorised combining operator. Rewrite --   reduce_by_index dst (map2 op) nes is vss@@ -796,146 +798,145 @@    m -  xs_bar <- traverse lookupAdjVal xs+  locallyNonvector (xs, dst, lam0, as) $ do+    xs_bar <- traverse lookupAdjVal xs -  (dst_params, hp_params, f') <- mkF' lam0 dst_type $ head w-  f'_adj_dst <- vjpLambda ops (map adjFromVar xs_bar) dst_params f'-  f'_adj_hp <- vjpLambda ops (map adjFromVar xs_bar) hp_params f'+    (dst_params, hp_params, f') <- mkF' lam0 dst_type $ head w+    f'_adj_dst <- vjpLambda ops (map adjFromVar xs_bar) dst_params f'+    f'_adj_hp <- vjpLambda ops (map adjFromVar xs_bar) hp_params f' -  dst_bar' <- eLambda f'_adj_dst $ map (eSubExp . Var) $ dst <> h_part-  dst_bar <- bindSubExpRes "dst_bar" dst_bar'-  zipWithM_ updateAdj dst dst_bar+    dst_bar' <- eLambda f'_adj_dst $ map (eSubExp . Var) $ dst <> h_part+    dst_bar <- bindSubExpRes "dst_bar" dst_bar'+    zipWithM_ updateAdj dst dst_bar -  h_part_bar' <- eLambda f'_adj_hp $ map (eSubExp . Var) $ dst <> h_part-  h_part_bar <- bindSubExpRes "h_part_bar" h_part_bar'+    h_part_bar' <- eLambda f'_adj_hp $ map (eSubExp . Var) $ dst <> h_part+    h_part_bar <- bindSubExpRes "h_part_bar" h_part_bar' -  lam <- renameLambda lam0-  lam' <- renameLambda lam0+    lam <- renameLambda lam0+    lam' <- renameLambda lam0 -  -- is' <- mapout (head as) n (head w)-  -- sorted <- radixSort' (is' : tail as) n $ head w-  sorted <- radixSort' as n $ head w-  let siota = head sorted-  let sis = head $ tail sorted-  let sas = drop 2 sorted+    -- is' <- mapout (head as) n (head w)+    -- sorted <- radixSort' (is' : tail as) n $ head w+    sorted <- radixSort' as n $ head w+    let siota = head sorted+    let sis = head $ tail sorted+    let sas = drop 2 sorted -  iota_n <--    letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64+    iota_n <-+      letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64 -  par_i <- newParam "i" $ Prim int64-  flag_lam <- mkFlagLam par_i sis-  flag <- letExp "flag" . Op . Screma n [iota_n] =<< mapSOAC flag_lam+    par_i <- newParam "i" $ Prim int64+    flag_lam <- mkFlagLam par_i sis+    flag <- letExp "flag" . Op . Screma n [iota_n] =<< mapSOAC flag_lam -  -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n)-  par_i' <- newParam "i" $ Prim int64-  let i' = paramName par_i'-  g_lam <--    mkLambda [par_i'] $-      fmap subExpsRes . mapM (letSubExp "scan_inps") =<< do-        im1 <- letSubExp "i_1" =<< toExp (le64 i' - 1)-        nmi <- letSubExp "n_i" =<< toExp (pe64 n - le64 i')-        let s1 = [DimFix im1]-        let s2 = [DimFix nmi]+    -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n)+    par_i' <- newParam "i" $ Prim int64+    let i' = paramName par_i'+    g_lam <-+      mkLambda [par_i'] $+        fmap subExpsRes . mapM (letSubExp "scan_inps") =<< do+          im1 <- letSubExp "i_1" =<< toExp (le64 i' - 1)+          nmi <- letSubExp "n_i" =<< toExp (pe64 n - le64 i')+          let s1 = [DimFix im1]+          let s2 = [DimFix nmi] -        -- flag array for left scan-        f1 <- letSubExp "f1" $ BasicOp $ Index flag $ Slice [DimFix $ Var i']+          -- flag array for left scan+          f1 <- letSubExp "f1" $ BasicOp $ Index flag $ Slice [DimFix $ Var i'] -        -- array for left scan-        r1 <--          letTupExp' "r1"-            =<< eIf-              (eSubExp f1)-              (eBody $ fmap eSubExp ne)-              (eBody . fmap (eSubExp . Var) =<< multiIndex sas s1)+          -- array for left scan+          r1 <-+            letTupExp' "r1"+              =<< eIf+                (eSubExp f1)+                (eBody $ fmap eSubExp ne)+                (eBody . fmap (eSubExp . Var) =<< multiIndex sas s1) -        -- array for right scan inc flag-        r2 <--          letTupExp' "r2"-            =<< eIf-              (toExp $ le64 i' .==. 0)-              (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)-              ( eBody $-                  pure $ do-                    eIf-                      (pure $ BasicOp $ Index flag $ Slice s2)-                      (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)-                      ( eBody . fmap eSubExp . (Constant (blankPrimValue Bool) :) . fmap Var-                          =<< multiIndex sas s2-                      )-              )+          -- array for right scan inc flag+          r2 <-+            letTupExp' "r2"+              =<< eIf+                (toExp $ le64 i' .==. 0)+                (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)+                ( eBody $+                    pure $ do+                      eIf+                        (pure $ BasicOp $ Index flag $ Slice s2)+                        (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)+                        ( eBody . fmap eSubExp . (Constant (blankPrimValue Bool) :) . fmap Var+                            =<< multiIndex sas s2+                        )+                ) -        traverse eSubExp $ f1 : r1 ++ r2+          traverse eSubExp $ f1 : r1 ++ r2 -  -- scan (\(f1,v1) (f2,v2) ->-  --   let f = f1 || f2-  --   let v = if f2 then v2 else g v1 v2-  --   in (f,v) ) (false,ne) (zip flags vals)-  scan_lams <--    traverse-      ( \l -> do-          f1 <- newParam "f1" $ Prim Bool-          f2 <- newParam "f2" $ Prim Bool-          ps <- lambdaParams <$> renameLambda lam0-          let (p1, p2) = splitAt (length ne) ps+    -- scan (\(f1,v1) (f2,v2) ->+    --   let f = f1 || f2+    --   let v = if f2 then v2 else g v1 v2+    --   in (f,v) ) (false,ne) (zip flags vals)+    scan_lams <-+      traverse+        ( \l -> do+            f1 <- newParam "f1" $ Prim Bool+            f2 <- newParam "f2" $ Prim Bool+            ps <- lambdaParams <$> renameLambda lam0+            let (p1, p2) = splitAt (length ne) ps -          mkLambda (f1 : p1 ++ f2 : p2) $-            fmap varsRes . letTupExp "scan_res" =<< do-              let f = eBinOp LogOr (eParam f1) (eParam f2)-              eIf-                (eParam f2)-                (eBody $ f : fmap eParam p2)-                ( eBody . (f :) . fmap (eSubExp . Var)-                    =<< bindSubExpRes "gres"-                    =<< eLambda l (fmap eParam ps)-                )-      )-      [lam, lam']+            mkLambda (f1 : p1 ++ f2 : p2) $+              fmap varsRes . letTupExp "scan_res" =<< do+                let f = eBinOp LogOr (eParam f1) (eParam f2)+                eIf+                  (eParam f2)+                  (eBody $ f : fmap eParam p2)+                  ( eBody . (f :) . fmap (eSubExp . Var)+                      =<< bindSubExpRes "gres"+                      =<< eLambda l (fmap eParam ps)+                  )+        )+        [lam, lam'] -  let ne' = Constant (BoolValue False) : ne+    let ne' = Constant (BoolValue False) : ne -  scansres <--    letTupExp "adj_ctrb_scan" . Op . Screma n [iota_n]-      =<< scanomapSOAC (map (`Scan` ne') scan_lams) g_lam+    scansres <-+      letTupExp "adj_ctrb_scan" . Op . Screma n [iota_n]+        =<< scanomapSOAC (map (`Scan` ne') scan_lams) g_lam -  let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres+    let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres -  -- map (\i -> if i < w && -1 < w then (xs_bar[i], dst[i]) else (0,ne)) sis-  par_i'' <- newParam "i" $ Prim int64-  let i'' = paramName par_i''-  map_lam <--    mkLambda [par_i''] $-      fmap varsRes . letTupExp "scan_res"-        =<< eIf-          (toExp $ withinBounds $ pure (head w, i''))-          (eBody . fmap (eSubExp . Var) =<< multiIndex h_part_bar [DimFix $ Var i''])-          ( eBody $ do-              map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type-          )+    -- map (\i -> if i < w && -1 < w then (xs_bar[i], dst[i]) else (0,ne)) sis+    par_i'' <- newParam "i" $ Prim int64+    let i'' = paramName par_i''+    map_lam <-+      mkLambda [par_i''] $+        fmap varsRes . letTupExp "scan_res"+          =<< eIf+            (toExp $ withinBounds $ pure (head w, i''))+            (eBody . fmap (eSubExp . Var) =<< multiIndex h_part_bar [DimFix $ Var i''])+            ( eBody $ do+                map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type+            ) -  f_bar <- letTupExp "f_bar" . Op . Screma n [sis] =<< mapSOAC map_lam+    f_bar <- letTupExp "f_bar" . Op . Screma n [sis] =<< mapSOAC map_lam -  (as_params, f) <- mkF lam0 as_type n-  f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f+    (as_params, f) <- mkF lam0 as_type n+    f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f -  -- map (\i -> rs_arr_rev[n-i-1]) (iota n)-  par_i''' <- newParam "i" $ Prim int64-  let i''' = paramName par_i'''-  rev_lam <- mkLambda [par_i'''] $ do-    nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1)-    varsRes <$> multiIndex rs_arr_rev [DimFix nmim1]+    -- map (\i -> rs_arr_rev[n-i-1]) (iota n)+    par_i''' <- newParam "i" $ Prim int64+    let i''' = paramName par_i'''+    rev_lam <- mkLambda [par_i'''] $ do+      nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1)+      varsRes <$> multiIndex rs_arr_rev [DimFix nmim1] -  rs_arr <--    letTupExp "rs_arr" . Op . Screma n [iota_n]-      =<< mapSOAC rev_lam+    rs_arr <- letTupExp "rs_arr" . Op . Screma n [iota_n] =<< mapSOAC rev_lam -  sas_bar <--    bindSubExpRes "sas_bar"-      =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr)+    sas_bar <-+      bindSubExpRes "sas_bar"+        =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr) -  scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type-  as_bar <- multiScatter scatter_dst siota sas_bar+    scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type+    as_bar <- multiScatter scatter_dst siota sas_bar -  zipWithM_ updateAdj (tail as) as_bar+    zipWithM_ updateAdj (tail as) as_bar   where     -- map (\i -> if i == 0 then true else is[i] != is[i-1]) (iota n)     mkFlagLam :: LParam SOACS -> VName -> ADM (Lambda SOACS)
src/Futhark/AD/Rev/Loop.hs view
@@ -267,7 +267,7 @@      pure (M.singleton i i_rev, i_stms) --- | Pures a substitution which substitutes values in the reverse+-- | Returns a substitution which substitutes values in the reverse -- loop body with values from the tape. restore :: Stms SOACS -> [Param DeclType] -> VName -> ADM Substitutions restore stms_adj loop_params' i' =
src/Futhark/AD/Rev/Map.hs view
@@ -75,6 +75,32 @@     subAD $ mkLambda (cert_params ++ acc_params) $ m $ map paramName acc_params   letTupExp "withhacc_res" $ WithAcc inputs acc_lam +vecPerm :: Shape -> Type -> [Int]+vecPerm adj_shape t =+  [shapeRank adj_shape]+    ++ [0 .. shapeRank adj_shape - 1]+    ++ [shapeRank adj_shape + 1 .. arrayRank t - 1]++pushAdjShape :: VName -> ADM VName+pushAdjShape v = do+  adj_shape <- askShape+  v_t <- lookupType v+  if adj_shape == mempty || arrayShape v_t == adj_shape || isAcc v_t+    then pure v+    else do+      let perm = vecPerm adj_shape v_t+      letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm++popAdjShape :: VName -> ADM VName+popAdjShape v = do+  adj_shape <- askShape+  v_t <- lookupType v+  if adj_shape == mempty || arrayShape v_t == adj_shape || isAcc v_t+    then pure v+    else do+      let perm = rearrangeInverse $ vecPerm adj_shape v_t+      letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm+ -- | Perform VJP on a Map.  The 'Adj' list is the adjoints of the -- result of the map. vjpMap :: VjpOps -> [Adj] -> StmAux () -> SubExp -> Lambda SOACS -> [VName] -> ADM ()@@ -150,8 +176,8 @@        zipWithM_ forRes [0 ..] res_ivs   where-    isSparse (AdjSparse (Sparse shape _ ivs)) = do-      guard $ shapeDims shape == [w]+    isSparse (AdjSparse (Sparse shape _ vd ivs)) = do+      guard $ drop vd (shapeDims shape) == [w]       Just ivs     isSparse _ =       Nothing@@ -161,7 +187,8 @@   pat_adj_vals <- forM (zip pat_adj (lambdaReturnType map_lam)) $ \(adj, t) ->     case t of       Acc {} -> letExp "acc_adj_rep" . BasicOp . Replicate (Shape [w]) . Var =<< adjVal adj-      _ -> adjVal adj+      _ -> pushAdjShape =<< adjVal adj+   pat_adj_params <-     mapM (newParam "map_adj_p" . rowType <=< lookupType) pat_adj_vals @@ -195,8 +222,8 @@     let param_ts = map paramType (lambdaParams map_lam')     forM_ (zip3 param_ts as param_contribs) $ \(param_t, a, param_contrib) ->       case param_t of-        Acc {} -> freeContrib a param_contrib-        _ -> updateAdj a param_contrib+        Acc {} -> freeContrib a =<< popAdjShape param_contrib -- CHECKME+        _ -> updateAdj a =<< popAdjShape param_contrib   where     addIdxParams n lam = do       idxs <- replicateM n $ newParam "idx" $ Prim int64
src/Futhark/AD/Rev/Monad.hs view
@@ -9,10 +9,12 @@ module Futhark.AD.Rev.Monad   ( ADM,     RState (..),+    REnv,     runADM,     Adj (..),     InBounds (..),     Sparse (..),+    askShape,     adjFromParam,     adjFromVar,     lookupAdj,@@ -43,6 +45,7 @@     zeroArray,     unitAdjOfType,     addLambda,+    vecOpExp,     --     VjpOps (..),     --@@ -50,14 +53,19 @@     lookupLoopTape,     substLoopTape,     renameLoopTape,+    --+    locallyNonvector,+    vecToInner,   ) where  import Control.Monad+import Control.Monad.Reader import Control.Monad.State.Strict import Data.Bifunctor (second) import Data.Map qualified as M import Data.Maybe+import Futhark.AD.Shared import Futhark.Analysis.Alias qualified as Alias import Futhark.Analysis.PrimExp.Convert import Futhark.Builder@@ -104,10 +112,17 @@ -- | A symbolic representation of an array that is all zeroes, except -- at certain indexes. data Sparse = Sparse-  { -- | The shape of the array.+  { -- | The full shape of the array (including any vector dimensions, which are+    -- stored in sparseVecDims).     sparseShape :: Shape,     -- | Element type of the array.     sparseType :: PrimType,+    -- | Number of leading dimensions that are \"vector\" dimensions, due to+    -- vector AD. These are not indexed by the sparse index, but are present in+    -- the values. When zero, this is the ordinary non-vector case. This is+    -- equivalent to the rank of `askShape`, but it is convenient to store it+    -- here as well.+    sparseVecDims :: Int,     -- | Locations and values of nonzero values.  Indexes may be     -- negative, in which case the value is ignored (unless     -- 'AssumeBounds' is used).@@ -140,14 +155,15 @@           Replicate shape zero  sparseArray :: (MonadBuilder m, Rep m ~ SOACS) => Sparse -> m VName-sparseArray (Sparse shape t ivs) = do+sparseArray (Sparse shape t vec_dims ivs) = do   flip (foldM f) ivs =<< zeroArray shape (Prim t)   where     arr_t = Prim t `arrayOfShape` shape+    vec_slice = map sliceDim $ take vec_dims $ shapeDims shape     f arr (check, i, se) = do       let stm s =             letExp "sparse" . BasicOp $-              Update s arr (fullSlice arr_t [DimFix i]) se+              Update s arr (fullSlice arr_t (vec_slice ++ [DimFix i])) se       case check of         AssumeBounds -> stm Unsafe         CheckBounds _ -> stm Safe@@ -170,8 +186,8 @@ adjRep :: Adj -> ([SubExp], [SubExp] -> Adj) adjRep (AdjVal se) = ([se], \[se'] -> AdjVal se') adjRep (AdjZero shape pt) = ([], \[] -> AdjZero shape pt)-adjRep (AdjSparse (Sparse shape pt ivs)) =-  (concatMap ivRep ivs, AdjSparse . Sparse shape pt . repIvs ivs)+adjRep (AdjSparse (Sparse shape pt vd ivs)) =+  (concatMap ivRep ivs, AdjSparse . Sparse shape pt vd . repIvs ivs)   where     ivRep (_, i, v) = [i, v]     repIvs ((check, _, _) : ivs') (i : v : ses) =@@ -197,12 +213,18 @@     stateNameSource :: VNameSource   } -newtype ADM a = ADM (BuilderT SOACS (State RState) a)+data REnv = REnv+  { envAdjShape :: Shape,+    envAttrs :: Attrs+  }++newtype ADM a = ADM (BuilderT SOACS (ReaderT REnv (State RState)) a)   deriving     ( Functor,       Applicative,       Monad,       MonadState RState,+      MonadReader REnv,       MonadFreshNames,       HasScope SOACS,       LocalScope SOACS@@ -221,16 +243,21 @@   getNameSource = gets stateNameSource   putNameSource src = modify (\env -> env {stateNameSource = src}) -runADM :: (MonadFreshNames m) => ADM a -> m a-runADM (ADM m) =+askShape :: ADM Shape+askShape = ADM $ lift $ asks envAdjShape++runADM :: (MonadFreshNames m) => Shape -> Attrs -> ADM a -> m a+runADM shape attrs (ADM m) =   modifyNameSource $ \vn ->     second stateNameSource $       runState-        (fst <$> runBuilderT m mempty)+        ( runReaderT (fst <$> runBuilderT m mempty) $+            REnv shape attrs+        )         (RState mempty mempty mempty vn)  adjVal :: Adj -> ADM VName-adjVal (AdjVal se) = letExp "const_adj" $ BasicOp $ SubExp se+adjVal (AdjVal se) = letExp "const_val_adj" $ BasicOp $ SubExp se adjVal (AdjSparse sparse) = sparseArray sparse adjVal (AdjZero shape t) = zeroArray shape $ Prim t @@ -312,13 +339,12 @@   where     names' = namesToList names -addBinOp :: PrimType -> BinOp-addBinOp (IntType it) = Add it OverflowWrap-addBinOp (FloatType ft) = FAdd ft-addBinOp Bool = LogAnd-addBinOp Unit = LogAnd--tabNest :: Int -> [VName] -> ([VName] -> [VName] -> ADM [VName]) -> ADM [VName]+tabNest ::+  (MonadBuilder m, Rep m ~ SOACS) =>+  Int ->+  [VName] ->+  ([VName] -> [VName] -> m [VName]) ->+  m [VName] tabNest = tabNest' []   where     tabNest' is 0 vs f = f (reverse is) vs@@ -338,13 +364,14 @@       let lam = Lambda (iparam : params) ret (Body () stms res)       letTupExp "tab" . Op . Screma w (iota : vs) =<< mapSOAC lam --- | Construct a lambda for adding two values of the given type.-addLambda :: Type -> ADM (Lambda SOACS)-addLambda (Prim pt) = binOpLambda (addBinOp pt) pt-addLambda t@Array {} = do+-- | Construct a lambda for binop'ing two values of the given type,+-- which may be arrays.+vecOpLambda :: (PrimType -> BinOp) -> Type -> ADM (Lambda SOACS)+vecOpLambda bop (Prim pt) = binOpLambda (bop pt) pt+vecOpLambda bop t@Array {} = do   xs_p <- newParam "xs" t   ys_p <- newParam "ys" t-  lam <- addLambda $ rowType t+  lam <- vecOpLambda bop $ rowType t   body <- insertStmsM $ do     res <-       letSubExp "lam_map"@@ -358,10 +385,10 @@         lambdaReturnType = [t],         lambdaBody = body       }-addLambda t =-  error $ "addLambda: " ++ show t+vecOpLambda _ t =+  error $ "vecOpLambda: " ++ show t --- Construct an expression for adding the two variables.+-- | Construct an expression for adding the two variables. addExp :: VName -> VName -> ADM (Exp SOACS) addExp x y = do   x_t <- lookupType x@@ -374,9 +401,23 @@     _ ->       error $ "addExp: unexpected type: " ++ prettyString x_t +-- | Construct an expression for performing this binary operation on two variables.+vecOpExp :: (PrimType -> BinOp) -> VName -> VName -> ADM (Exp SOACS)+vecOpExp bop x y = do+  x_t <- lookupType x+  case x_t of+    Prim pt ->+      pure $ BasicOp $ BinOp (bop pt) (Var x) (Var y)+    Array {} -> do+      lam <- vecOpLambda bop $ rowType x_t+      Op . Screma (arraySize 0 x_t) [x, y] <$> mapSOAC lam+    _ ->+      error $ "vecOpExp: unexpected type: " ++ prettyString x_t+ lookupAdj :: VName -> ADM Adj lookupAdj v = do   maybeAdj <- gets $ M.lookup v . stateAdjs+  adj_shape <- askShape   case maybeAdj of     Nothing -> do       v_t <- lookupType v@@ -384,7 +425,7 @@         Acc _ shape [Prim t] _ -> pure $ AdjZero shape t         Acc _ shape [t] _ -> pure $ AdjZero (shape <> arrayShape t) (elemType t)         Acc {} -> error $ "lookupAdj: Non-singleton accumulator adjoint: " <> prettyString v_t-        _ -> pure $ AdjZero (arrayShape v_t) (elemType v_t)+        _ -> pure $ AdjZero (adj_shape <> arrayShape v_t) (elemType v_t)     Just v_adj -> pure v_adj  lookupAdjVal :: VName -> ADM VName@@ -394,27 +435,34 @@ updateAdjIndex v (check, i) se = do   maybeAdj <- gets $ M.lookup v . stateAdjs   t <- lookupType v+  adj_shape <- askShape   let iv = (check, i, se)+      vec_dims = shapeRank adj_shape+      full_shape = adj_shape <> arrayShape t   case maybeAdj of-    Nothing -> do-      setAdj v $ AdjSparse $ Sparse (arrayShape t) (elemType t) [iv]-    Just AdjZero {} ->-      setAdj v $ AdjSparse $ Sparse (arrayShape t) (elemType t) [iv]-    Just (AdjSparse (Sparse shape pt ivs)) ->-      setAdj v $ AdjSparse $ Sparse shape pt $ iv : ivs+    Nothing ->+      setAdj v $ AdjSparse $ Sparse full_shape (elemType t) vec_dims [iv]+    Just (AdjZero {}) ->+      setAdj v $ AdjSparse $ Sparse full_shape (elemType t) vec_dims [iv]+    Just (AdjSparse (Sparse shape pt vd ivs)) ->+      setAdj v $ AdjSparse $ Sparse shape pt vd $ iv : ivs     Just adj@AdjVal {} -> do       v_adj <- adjVal adj       v_adj_t <- lookupType v_adj       se_v <- letExp "se_v" $ BasicOp $ SubExp se+      vec_shape <- askShape       insAdj v         =<< case v_adj_t of           Acc {} -> do             let stms s = do+                  attrs <- asks envAttrs                   dims <- arrayDims <$> lookupType se_v                   ~[v_adj'] <--                    tabNest (length dims) [se_v, v_adj] $ \is [se_v', v_adj'] ->-                      letTupExp "acc" . BasicOp $-                        UpdateAcc s v_adj' (i : map Var is) [Var se_v']+                    attributing attrs $+                      tabNest (length dims) [se_v, v_adj] $ \is [se_v', v_adj'] -> do+                        let (vec_is, val_is) = splitAt (shapeRank vec_shape) $ map Var is+                        letTupExp "acc" . BasicOp $+                          UpdateAcc s v_adj' (vec_is ++ i : val_is) [Var se_v']                   pure v_adj'             case check of               CheckBounds _ -> stms Safe@@ -422,13 +470,15 @@               OutOfBounds -> pure v_adj           _ -> do             let stms s = do+                  let slice =+                        fullSlice v_adj_t $+                          map sliceDim (shapeDims vec_shape) ++ [DimFix i]                   v_adj_i <-                     letExp (baseName v_adj <> "_i") . BasicOp $-                      Index v_adj $-                        fullSlice v_adj_t [DimFix i]+                      Index v_adj slice                   se_update <- letSubExp "updated_adj_i" =<< addExp se_v v_adj_i                   letExp (baseName v_adj) . BasicOp $-                    Update s v_adj (fullSlice v_adj_t [DimFix i]) se_update+                    Update s v_adj slice se_update             case check of               CheckBounds _ -> stms Safe               AssumeBounds -> stms Unsafe@@ -518,7 +568,8 @@  data VjpOps = VjpOps   { vjpLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS),-    vjpStm :: Stm SOACS -> ADM () -> ADM ()+    vjpStm :: Stm SOACS -> ADM () -> ADM (),+    vjpBody :: [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)   }  -- | @setLoopTape v vs@ establishes @vs@ as the name of the array@@ -542,6 +593,50 @@ -- the names in the loop tape after a loop rename. renameLoopTape :: Substitutions -> ADM () renameLoopTape = mapM_ (uncurry substLoopTape) . M.toList++-- | Disable vector AD within the provided action. This results in a map that+-- computes each adjoint explicitly, then assembles the resulting adjoint+-- vectors. This is useful for constructs (such as scans) where vector AD is+-- impractical or inefficient.+locallyNonvector ::+  (FreeIn e) =>+  -- | Something that represents all the free variables used in the action.+  -- Usually just an expression or statement.+  e ->+  ADM () ->+  ADM ()+locallyNonvector e m = do+  adj_shape <- askShape+  if adj_shape == mempty+    then m+    else do+      -- We map over all adjoints of free variables in 'e'. To avoid clutter, we+      -- only consider those that actually have known nonzero adjoints.+      e_adjs <- filterM knownAdjoint e_free+      e_adjs_vals <- mapM lookupAdjVal e_adjs+      e_free_adjs <- mkMap "nonvec_adj" adj_shape e_adjs_vals $ \e_adjs_vals' -> do+        zipWithM_ insAdj e_adjs e_adjs_vals'+        local (\env -> env {envAdjShape = mempty}) m+        mapM lookupAdjVal e_free+      zipWithM_ insAdj e_free e_free_adjs+  where+    e_free = namesToList $ freeIn e+    knownAdjoint v = do+      v_adj <- lookupAdj v+      pure $ case v_adj of+        AdjZero {} -> False+        _ -> True++-- | If we are doing vector AD, apply 'vecPerm' to the array.+vecToInner :: VName -> ADM VName+vecToInner v = do+  adj_shape <- askShape+  if adj_shape == mempty+    then pure v+    else do+      v_t <- lookupType v+      letExp (baseName v <> "_tr") . BasicOp . Rearrange v $+        vecPerm adj_shape v_t  -- Note [Consumption] --
src/Futhark/AD/Rev/Reduce.hs view
@@ -9,7 +9,9 @@ where  import Control.Monad+import Data.Tuple import Futhark.AD.Rev.Monad+import Futhark.AD.Shared import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS@@ -78,10 +80,8 @@   | Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,     isAdd op = do       adj_rep <--        letExp (baseName adj <> "_rep") $-          BasicOp $-            Replicate (Shape [w]) $-              Var adj+        vecToInner <=< letExp (baseName adj <> "_rep") $+          BasicOp (Replicate (Shape [w]) (Var adj))       void $ updateAdj a adj_rep   where     isAdd FAdd {} = True@@ -118,10 +118,9 @@   f_adj <- vjpLambda ops (map adjFromVar pat_adj) as_params f    as_adj <--    letTupExp "adjs" . Op . Screma w (ls ++ as ++ rs)-      =<< mapSOAC f_adj+    letTupExp "red_contribs" . Op . Screma w (ls ++ as ++ rs) =<< mapSOAC f_adj -  zipWithM_ updateAdj as as_adj+  zipWithM_ updateAdj as =<< mapM vecToInner as_adj   where     renameRed (Reduce comm lam nes) =       Reduce comm <$> renameLambda lam <*> pure nes@@ -253,7 +252,8 @@   VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> ADM () -> ADM () diffMulReduce _ops x aux w mul ne as m = do   let t = binOpType mul-  let const_zero = eSubExp $ Constant $ blankPrimValue t+  let zero = Constant $ blankPrimValue t+      const_zero = eSubExp zero    a_param <- newParam "a" $ Prim t   map_lam <-@@ -292,42 +292,45 @@    x_adj <- lookupAdjVal x +  adj_shape <- askShape++  zero_contrib <- letExp "zero_contrib" $ BasicOp $ Replicate adj_shape zero+   a_param_rev <- newParam "a" $ Prim t   map_lam_rev <-     mkLambda [a_param_rev] $       fmap varsRes . letTupExp "adj_res"         =<< eIf           (toExp $ 0 .==. le64 zr_count)-          ( eBody $-              pure $-                eBinOp mul (eSubExp $ Var x_adj) $-                  eBinOp (getDiv t) (eSubExp $ Var nz_prods) $-                    eParam a_param_rev+          ( eBody+              [ mapNest adj_shape (MkSolo (Var x_adj)) $ \(MkSolo x_adj') ->+                  eBinOp mul (eSubExp x_adj') $+                    eBinOp (getDiv t) (eVar nz_prods) $+                      eParam a_param_rev+              ]           )-          ( eBody $-              pure $-                eIf+          ( eBody+              [ eIf                   (toExp $ 1 .==. le64 zr_count)-                  ( eBody $-                      pure $-                        eIf+                  ( eBody+                      [ eIf                           (eCmpOp (CmpEq t) (eParam a_param_rev) const_zero)-                          ( eBody $-                              pure $-                                eBinOp mul (eSubExp $ Var x_adj) $-                                  eSubExp $-                                    Var nz_prods+                          ( eBody+                              [ mapNest adj_shape (MkSolo (Var x_adj)) $+                                  \(MkSolo x_adj') ->+                                    eBinOp mul (eSubExp x_adj') $ eVar nz_prods+                              ]                           )-                          (eBody $ pure const_zero)+                          (eBody [eVar zero_contrib])+                      ]                   )-                  (eBody $ pure const_zero)+                  (eBody [eVar zero_contrib])+              ]           ) -  as_adjup <--    letExp "adjs" . Op . Screma w [as]-      =<< mapSOAC map_lam_rev+  as_adjup <- letExp "prod_contrib" . Op . Screma w [as] =<< mapSOAC map_lam_rev -  updateAdj as as_adjup+  updateAdj as =<< vecToInner as_adjup   where     getDiv :: PrimType -> BinOp     getDiv (IntType t) = SDiv t Unsafe
src/Futhark/AD/Rev/SOAC.hs view
@@ -177,14 +177,14 @@     Just [(op, _, _, _)] <- lamIsBinOp lam',     isAddOp op =       diffAddHist ops x aux n lam ne is vs w rf dst m-vjpSOAC ops pat aux (Hist n as [histop] f) m+vjpSOAC ops pat aux (Hist w as [histop] f) m   | isIdentityLambda f,-    HistOp (Shape w) rf dst ne lam <- histop = do-      diffHist ops (patNames pat) aux n lam ne as w rf dst m-vjpSOAC ops pat _aux (Hist n as histops f) m+    HistOp (Shape n) rf dst ne lam <- histop = do+      diffHist ops (patNames pat) aux w lam ne as n rf dst m+vjpSOAC ops pat _aux (Hist w as histops f) m   | not (isIdentityLambda f) = do       (mapstm, redstm) <--        histomapToMapAndHist pat (n, histops, f, as)+        histomapToMapAndHist pat (w, histops, f, as)       vjpStm ops mapstm $ vjpStm ops redstm m vjpSOAC ops pat aux (Stream w as accs lam) m = do   stms <- collectStms_ $ auxing aux $ sequentialStreamWholeArray pat w accs lam as@@ -194,13 +194,14 @@   forM_ (zip (patNames pat) lam_res) $ \(v, SubExpRes cs se) ->     certifying cs $ letBindNames [v] $ BasicOp $ SubExp se   m-  pat_adj <- mapM lookupAdjVal $ patNames pat-  contribs <--    eLambda lam_adj (map (eSubExp . resSubExp) lam_res ++ map (eSubExp . Var) pat_adj)-  forM_ (zip args contribs) $ \(arg, contrib) ->-    (updateSubExpAdj arg <=< letExp "contrib") $-      BasicOp . SubExp . resSubExp $-        contrib+  locallyNonvector (patNames pat, args) $ do+    pat_adj <- mapM lookupAdjVal $ patNames pat+    contribs <-+      eLambda lam_adj (map (eSubExp . resSubExp) lam_res ++ map (eSubExp . Var) pat_adj)+    forM_ (zip args contribs) $ \(arg, contrib) ->+      (updateSubExpAdj arg <=< letExp "contrib") $+        BasicOp . SubExp . resSubExp $+          contrib vjpSOAC _ _ _ soac _ =   error $ "vjpSOAC unhandled:\n" ++ prettyString soac
src/Futhark/AD/Rev/Scan.hs view
@@ -379,7 +379,7 @@     eLambda op_bar_2 $ toExp . Var . paramName <$> par_y_right ++ par_a  diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM ()-diffScan ops ys w as scan = do+diffScan ops ys w as scan = locallyNonvector (ys, scan, as) $ do   -- ys ~ results of scan, w ~ size of input array, as ~ (unzipped)   -- arrays, scan ~ scan: operator with ne   scan_case <- identifyCase ops $ scanLambda scan@@ -409,14 +409,12 @@       map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d       scans_lin_fun_o <- mkScanLinFunO (head as_ts) sc       scan_lams <- mkScans (specialScans sc) scans_lin_fun_o-      iota <--        letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+      iota <- letExp "iota" $ iota64 w       r_scan <-         letTupExp "adj_ctrb_scan" . Op . Screma w [iota]           =<< scanomapSOAC scan_lams map1_lam       mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d)   -- Goal: calculate as_contribs in new way-  -- zipWithM_ updateAdj as as_contribs -- as_bar += new adjoint   zipWithM_ updateAdj as as_contribs   where     mkScans :: Int -> Scan SOACS -> ADM [Scan SOACS]@@ -468,7 +466,7 @@   foldr (vjpStm ops) m stmts  diffScanAdd :: VjpOps -> VName -> SubExp -> Lambda SOACS -> SubExp -> VName -> ADM ()-diffScanAdd _ops ys n lam' ne as = do+diffScanAdd _ops ys n lam' ne as = locallyNonvector (ys, lam', as) $ do   lam <- renameLambda lam'   ys_bar <- lookupAdjVal ys
+ src/Futhark/AD/Shared.hs view
@@ -0,0 +1,62 @@+-- | Various definitions used for both forward and reverse mode.+module Futhark.AD.Shared+  ( vecPerm,+    asVName,+    mapNest,+    mkMap,+  )+where++import Control.Monad+import Data.Foldable+import Futhark.Construct+import Futhark.IR.SOACS++-- | A permutation for transposing the vector shape past the next dimension.+--+-- That is, converts @[vec...][d][elem...]@ to @[d][vec...][elem...]@.+vecPerm :: Shape -> Type -> [Int]+vecPerm vec_shape t =+  [shapeRank vec_shape]+    ++ [0 .. shapeRank vec_shape - 1]+    ++ [shapeRank vec_shape + 1 .. arrayRank t - 1]++asVName :: (MonadBuilder m) => SubExp -> m VName+asVName (Var v) = pure v+asVName (Constant x) = letExp "asv" $ BasicOp $ SubExp $ Constant x++mapNest ::+  (MonadBuilder m, Rep m ~ SOACS, Traversable f) =>+  Shape ->+  f SubExp ->+  (f SubExp -> m (Exp SOACS)) ->+  m (Exp SOACS)+mapNest shape x f = do+  if shape == mempty+    then f x+    else do+      let w = shapeSize 0 shape+      x_v <- traverse asVName x+      x_p <- traverse (newParam "xp" . rowType <=< lookupType) x_v+      lam <- mkLambda (toList x_p) $ do+        fmap (subExpsRes . pure) . letSubExp "mapnest_res"+          =<< f (fmap (Var . paramName) x_p)+      Op . Screma w (toList x_v) <$> mapSOAC lam++-- | Construct a map over the given arrays, which must have the provided outer+-- shape. The purpose of the 'Shape' argument is to handle the case where no+-- arrays are provided.+mkMap ::+  (MonadBuilder m, Rep m ~ SOACS, Traversable f) =>+  Name ->+  Shape ->+  f VName ->+  -- | Action for building the body, passed names+  -- corresponding to elements of the arrays.+  (f VName -> m [VName]) ->+  m [VName]+mkMap desc shape arrs f = do+  let w = shapeSize 0 shape+  x_p <- traverse (newParam "xp" . rowType <=< lookupType) arrs+  lam <- mkLambda (toList x_p) $ varsRes <$> f (fmap paramName x_p)+  letTupExp desc . Op . Screma w (toList arrs) =<< mapSOAC lam
src/Futhark/Analysis/SymbolTable.hs view
@@ -24,6 +24,7 @@     lookupExp,     lookupBasicOp,     lookupType,+    lookupSubExpType,     lookupSubExp,     lookupAliases,     lookupLoopVar,@@ -228,9 +229,11 @@   Just (BasicOp e, cs) -> Just (e, cs)   _ -> Nothing +-- | Look up the type of a name in the symbol table. lookupType :: (ASTRep rep) => VName -> SymbolTable rep -> Maybe Type lookupType name vtable = typeOf <$> lookup name vtable +-- | Look up the type of a 'SubExp'. lookupSubExpType :: (ASTRep rep) => SubExp -> SymbolTable rep -> Maybe Type lookupSubExpType (Var v) = lookupType v lookupSubExpType (Constant v) = const $ Just $ Prim $ primValueType v
src/Futhark/CLI/REPL.hs view
@@ -14,6 +14,7 @@ import Data.Map qualified as M import Data.Maybe import Data.Text qualified as T+import Data.Text.Encoding qualified as T import Data.Text.IO qualified as T import Data.Version import Futhark.Compiler@@ -436,11 +437,29 @@         Right parts' -> do           parts'' <- mapM sequenceA <$> mapM (traverse onExp) parts'           case parts'' of-            Left err -> liftIO $ putDoc err+            Left err -> liftIO $ putDocLn err             Right parts''' ->               liftIO . T.putStrLn . mconcat $                 map (either id (docText . I.prettyValue)) parts''' +stringCommand :: Command+stringCommand input = do+  prompt <- getPrompt+  case parseExp prompt input of+    Left (SyntaxError _ err) ->+      liftIO $ T.putStr err+    Right e -> do+      e' <- onExp e+      liftIO $ case e' of+        Left err -> putDocLn err+        Right v+          | Just bytes <- I.asByteString v ->+              T.putStrLn $ T.decodeUtf8 bytes+          | otherwise ->+              T.putStrLn $+                "Cannot print this value as string:\n"+                  <> docText (I.prettyValue v)+ unbreakCommand :: Command unbreakCommand _ = do   top <- gets $ fmap (NE.head . breakingStack) . futharkiBreaking@@ -529,6 +548,15 @@    > :format The value of foo: {foo}. The value of 2+2={2+2} |]+      )+    ),+    ( "string",+      ( stringCommand,+        [text|+Evaluate the given expression, which must have type []u8,+and print the result as a string. It is assumes that the+byte array contains valid UTF-8.+           |]       )     ),     ( "type",
src/Futhark/CodeGen/Backends/GenericC/CLI.hs view
@@ -191,7 +191,7 @@             [C.cstm|;|],             [C.cexp|$id:dest|]           )-    Just (TypeOpaque desc _ _) ->+    Just (TypeOpaque desc _ _ _) ->       ( [C.citems|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:(T.unpack desc));|],         [C.cstm|;|],         [C.cstm|;|],@@ -257,7 +257,7 @@             [C.cexp|$id:result|],             [C.cstm|assert($id:(arrayFree ops)(ctx, $id:result) == 0);|]           )-        Just (TypeOpaque t ops _) ->+        Just (TypeOpaque t ops _ _) ->           ( [C.citem|typename $id:t $id:result;|],             [C.cexp|$id:result|],             [C.cstm|assert($id:(opaqueFree ops)(ctx, $id:result) == 0);|]@@ -268,7 +268,7 @@   case M.lookup t $ manifestTypes manifest of     Nothing -> uncurry primAPIType $ scalarToPrim t     Just (TypeArray tname _ _ _) -> [C.cty|typename $id:tname|]-    Just (TypeOpaque tname _ _) -> [C.cty|typename $id:tname|]+    Just (TypeOpaque tname _ _ _) -> [C.cty|typename $id:tname|]   where     t = recordFieldType field @@ -279,23 +279,24 @@     Nothing ->       let info = tname <> "_info"        in [C.cstm|write_scalar(stdout, binary_output, &$id:info, &$exp:e);|]-    Just (TypeOpaque _ _ (Just (OpaqueRecord record)))+    Just (TypeOpaque _ _ (Just (OpaqueRecord record)) _)       | map recordFieldName fields == take (length fields) (map showText [0 :: Int ..]) ->           [C.cstm|{$stms:(intersperse newline (map getField fields))}|]       where         fields = recordFields record         printField field =-          printStm manifest (recordFieldType field) [C.cexp|field|]+          printStm manifest (recordFieldType field) [C.cexp|field_val|]         newline = [C.cstm|puts("");|]         getField field =           [C.cstm|{$ty:(recordFieldCType manifest field) field;                    if ($id:(recordFieldProject field)(ctx, &field, $exp:e) != FUTHARK_SUCCESS) {                      futhark_panic(1, "Failed to project field %s from result\n", $string:(T.unpack (recordFieldName field)));                    } else {+                     $ty:(recordFieldCType manifest field) field_val = field;                      $stm:(printField field)                    }                    }|]-    Just (TypeOpaque desc _ _) ->+    Just (TypeOpaque desc _ _ _) ->       [C.cstm|{          fprintf(stderr, "Values of type \"%s\" have no external representation.\n", $string:(T.unpack desc));          retval = 1;@@ -325,7 +326,7 @@  cliEntryPoint ::   Manifest -> T.Text -> EntryPoint -> (C.Definition, C.Initializer)-cliEntryPoint manifest entry_point_name (EntryPoint cfun _tuning_params output inputs _attrs) =+cliEntryPoint manifest entry_point_name (EntryPoint cfun _tuning_params output inputs _attrs _doc) =   let (input_items, pack_input, free_input, free_parsed, input_args) =         unzip5 $ readInputs manifest $ map inputType inputs
src/Futhark/CodeGen/Backends/GenericC/Code.hs view
@@ -76,6 +76,8 @@   y' <- compilePrimExp f y   pure $ case cmp of     CmpEq {} -> [C.cexp|$exp:x' == $exp:y'|]+    FCmpLt Float16 -> [C.cexp|cmplt16($exp:x', $exp:y')|]+    FCmpLe Float16 -> [C.cexp|cmple16($exp:x', $exp:y')|]     FCmpLt {} -> [C.cexp|$exp:x' < $exp:y'|]     FCmpLe {} -> [C.cexp|$exp:x' <= $exp:y'|]     CmpLlt {} -> [C.cexp|$exp:x' < $exp:y'|]
src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs view
@@ -137,7 +137,7 @@   Function op ->   CompilerM op s (Maybe (C.Definition, (T.Text, Manifest.EntryPoint))) onEntryPoint _ _ _ (Function Nothing _ _ _ _) = pure Nothing-onEntryPoint get_consts relevant_params fname (Function (Just (EntryPoint ename results args)) outputs inputs attrs _) = inNewFunction $ do+onEntryPoint get_consts relevant_params fname (Function (Just (EntryPoint ename results args doc)) outputs inputs attrs _) = inNewFunction $ do   let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs       in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs @@ -216,7 +216,8 @@             -- manifest and ImpCode.             Manifest.entryPointOutput = outputManifest results,             Manifest.entryPointInputs = map inputManifest args,-            Manifest.entryPointAttrs = map prettyText (S.toList (unAttrs attrs))+            Manifest.entryPointAttrs = map prettyText (S.toList (unAttrs attrs)),+            Manifest.entryPointDoc = doc           }    pure $ Just (cdef, (nameToText ename, manifest))
src/Futhark/CodeGen/Backends/GenericC/Server.hs view
@@ -121,7 +121,7 @@ cType manifest tname =   case M.lookup tname $ manifestTypes manifest of     Just (TypeArray ctype _ _ _) -> [C.cty|typename $id:(T.unpack ctype)|]-    Just (TypeOpaque ctype _ _) -> [C.cty|typename $id:(T.unpack ctype)|]+    Just (TypeOpaque ctype _ _ _) -> [C.cty|typename $id:(T.unpack ctype)|]     Nothing -> uncurry primAPIType $ scalarToPrim tname  data Kind@@ -208,7 +208,7 @@                 .info = &$id:array_name               };|]       )-typeBoilerplate manifest (tname, TypeOpaque c_type_name ops extra_ops) =+typeBoilerplate manifest (tname, TypeOpaque c_type_name ops extra_ops _) =   let type_name = typeStructName tname       aux_name = type_name <> "_aux"       (transparent_edecls, transparent_init, kind) = transparentDefs type_name extra_ops@@ -450,7 +450,7 @@     manifest  oneEntryBoilerplate :: Manifest -> (T.Text, EntryPoint) -> ([C.Definition], C.Initializer)-oneEntryBoilerplate manifest (name, EntryPoint cfun tuning_params output inputs attrs) =+oneEntryBoilerplate manifest (name, EntryPoint cfun tuning_params output inputs attrs _doc) =   let call_f = "call_" <> nameFromText name       out_type = outputType output       in_types = map inputType inputs
src/Futhark/CodeGen/Backends/GenericC/Types.hs view
@@ -252,13 +252,13 @@ lookupOpaqueType :: Name -> OpaqueTypes -> OpaqueType lookupOpaqueType v (OpaqueTypes types) =   case lookup v types of-    Just t -> t+    Just (t, _) -> t     Nothing -> error $ "Unknown opaque type: " ++ show v  opaquePayload :: OpaqueTypes -> OpaqueType -> [ValueType]-opaquePayload _ (OpaqueType ts) = ts opaquePayload _ (OpaqueSum ts _) = ts opaquePayload _ (OpaqueArray _ _ ts) = ts+opaquePayload _ (OpaqueRecord []) = [ValueType Signed (Rank 0) Unit] opaquePayload types (OpaqueRecord fs) = concatMap f fs   where     f (_, TypeOpaque s) = opaquePayload types $ lookupOpaqueType s types@@ -346,7 +346,7 @@           [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj);|]         libDecl           [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj) {-                      (void)ctx;+                      (void)ctx; (void)obj;                       $ty:et_ty v;                       $items:project_items                       *out = v;@@ -413,7 +413,9 @@             ( offset + length f_vts,               ( param_name,                 [C.cparam|const $ty:ct* $id:param_name|],-                [C.citem|{$stms:(zipWith3 setFieldField [offset ..] param_fields f_vts)}|]+                if null f_vts+                  then [C.citem|(void)$id:param_name;|]+                  else [C.citem|{$stms:(zipWith3 setFieldField [offset ..] param_fields f_vts)}|]               )             ) @@ -455,14 +457,14 @@   CompilerM op s [Manifest.RecordField] recordArrayProjectFunctions = recordProjectFunctions -recordArrayZipFunctions ::+recordArrayZipFunction ::   OpaqueTypes ->   Name ->   [(Name, EntryPointType)] ->   [ValueType] ->   Int ->   CompilerM op s Manifest.CFuncName-recordArrayZipFunctions types desc fs vds rank = do+recordArrayZipFunction types desc fs vds rank = do   opaque_type <- opaqueToCType desc   ctx_ty <- contextType   ops <- asks envOperations@@ -496,8 +498,9 @@  indexingDefs ::   Int ->+  C.Exp ->   ([C.Param], PrimType -> Int -> C.Exp -> C.Exp, C.Exp)-indexingDefs rank =+indexingDefs rank outer_shape =   (index_params, indexExp, in_bounds)   where     index_names = ["i" <> prettyText i | i <- [0 .. rank - 1]]@@ -514,7 +517,7 @@      in_bounds =       allTrue-        [ [C.cexp|$id:p >= 0 && $id:p < arr->$id:(tupleField 0)->shape[$int:i]|]+        [ [C.cexp|$id:p >= 0 && $id:p < $exp:outer_shape[$int:i]|]         | (p, i) <- zip index_names [0 .. rank - 1]         ] @@ -542,6 +545,7 @@   libDecl     [C.cedecl|int $id:index_f($ty:ctx_ty *ctx, $ty:obj_ct **out, $ty:array_ct *arr,                               $params:index_params) {+                (void)arr;                 int err = 0;                 if ($exp:in_bounds) {                   $ty:obj_ct* v = malloc(sizeof($ty:obj_ct));@@ -558,9 +562,14 @@    pure index_f   where-    (index_params, indexExp, in_bounds) = indexingDefs rank+    (index_params, indexExp, in_bounds) =+      indexingDefs rank $+        if null vds+          then [C.cexp|arr->shape|]+          else [C.cexp|arr->$id:(tupleField 0)->shape|]      setField copy j (ValueType _ (Rank r) pt)+      | pt == Unit = pure ()       | r == rank =           -- Easy case: just copy the scalar from the array into the           -- variable.@@ -619,6 +628,7 @@                               $ty:array_ct *arr,                               $ty:obj_ct *v,                               $params:index_params) {+                (void)arr; (void)v;                 int err = 0;                 if (!$exp:in_bounds) {                   err = 1;@@ -634,7 +644,11 @@    pure index_f   where-    (index_params, indexExp, in_bounds) = indexingDefs rank+    (index_params, indexExp, in_bounds) =+      indexingDefs rank $+        if null vds+          then [C.cexp|arr->shape|]+          else [C.cexp|arr->$id:(tupleField 0)->shape|]      same_shape = allTrue $ do       (j, ValueType _ (Rank r) _) <- zip [0 ..] vds@@ -645,6 +659,7 @@                        $int:(r-rank) * sizeof(int64_t)) == 0|]      setField copy j (ValueType _ (Rank r) pt)+      | pt == Unit = pure ()       | r == rank =           copy             CopyBarrier@@ -669,22 +684,31 @@             space             $ cproduct ([C.cexp|$int:(primByteSize pt::Int)|] : shape) -recordArrayShapeFunction :: Name -> CompilerM op s Manifest.CFuncName-recordArrayShapeFunction desc = do+recordArrayShapeFunction ::+  Name ->+  [ValueType] ->+  CompilerM op s Manifest.CFuncName+recordArrayShapeFunction desc vds = do   shape_f <- publicName $ "shape_" <> opaqueName desc   ctx_ty <- contextType   array_ct <- opaqueToCType desc -  -- We know that the opaque value consists of arrays of at least the-  -- expected rank, and which have the same outer shape, so we just-  -- return the shape of the first one.+  -- We know that the opaque value consists of arrays of at least the expected+  -- rank, and which have the same outer shape, so we just return the shape of+  -- the first one. However, in the special case of an array of units, there are+  -- actually no embedded arrays, so we return the special shape field instead.++  let shape = case vds of+        [] -> [C.cexp|arr->shape|]+        _ -> [C.cexp|arr->$id:(tupleField 0)->shape|]+   headerDecl     (OpaqueDecl desc)     [C.cedecl|const typename int64_t* $id:shape_f($ty:ctx_ty *ctx, $ty:array_ct *arr);|]   libDecl     [C.cedecl|const typename int64_t* $id:shape_f($ty:ctx_ty *ctx, $ty:array_ct *arr) {                 (void)ctx;-                return arr->$id:(tupleField 0)->shape;+                return $exp:shape;               }|]    pure shape_f@@ -720,6 +744,7 @@   libDecl     [C.cedecl|int $id:new_f($ty:ctx_ty *ctx, $ty:array_ct **out,                             $ty:obj_ct **elems, $params:shape_params) {+                (void)elems;                 int err = 0;                 typename int64_t n = $exp:(cproduct outer_shape);                 $items:check_items@@ -756,6 +781,8 @@                     goto end;                   }|] +    handleField _ _ (ValueType _ _ Unit) =+      pure ()     handleField copy i (ValueType _ (Rank r) pt) = do       let elem_size =             cproduct $@@ -805,39 +832,6 @@         [C.cstm|for (typename int64_t i = 1; i < n; i++)                 { $items:check_one }|] -opaqueArrayIndexFunction ::-  Space ->-  OpaqueTypes ->-  Name ->-  Int ->-  Name ->-  [ValueType] ->-  CompilerM op s Manifest.CFuncName-opaqueArrayIndexFunction = recordArrayIndexFunction--opaqueArrayShapeFunction :: Name -> CompilerM op s Manifest.CFuncName-opaqueArrayShapeFunction = recordArrayShapeFunction--opaqueArrayNewFunction ::-  Space ->-  OpaqueTypes ->-  Name ->-  Int ->-  Name ->-  [ValueType] ->-  CompilerM op s Manifest.CFuncName-opaqueArrayNewFunction = recordArrayNewFunction--opaqueArraySetFunction ::-  Space ->-  OpaqueTypes ->-  Name ->-  Int ->-  Name ->-  [ValueType] ->-  CompilerM op s Manifest.CFuncName-opaqueArraySetFunction = recordArraySetFunction- sumVariants ::   Name ->   [(Name, [(EntryPointType, [Int])])] ->@@ -1004,15 +998,13 @@   _   types   "()"-  (OpaqueType [ValueType Signed (Rank 0) Unit])+  (OpaqueRecord [])   [ValueType Signed (Rank 0) Unit] =     Just . Manifest.OpaqueRecord       <$> ( Manifest.RecordOps               <$> recordProjectFunctions types "()" [] []               <*> recordNewFunctions types "()" [] []           )-opaqueExtraOps _ _ _ (OpaqueType _) _ =-  pure Nothing opaqueExtraOps _ _types desc (OpaqueSum _ cs) vds =   Just . Manifest.OpaqueSum     <$> ( Manifest.SumOps@@ -1029,19 +1021,19 @@   Just . Manifest.OpaqueRecordArray     <$> ( Manifest.RecordArrayOps rank (nameToText elemtype)             <$> recordArrayProjectFunctions types desc fs vds-            <*> recordArrayZipFunctions types desc fs vds rank+            <*> recordArrayZipFunction types desc fs vds rank             <*> recordArrayIndexFunction space types desc rank elemtype vds-            <*> recordArrayShapeFunction desc+            <*> recordArrayShapeFunction desc vds             <*> recordArrayNewFunction space types desc rank elemtype vds             <*> recordArraySetFunction space types desc rank elemtype vds         ) opaqueExtraOps space types desc (OpaqueArray rank elemtype _) vds =   Just . Manifest.OpaqueArray     <$> ( Manifest.OpaqueArrayOps rank (nameToText elemtype)-            <$> opaqueArrayIndexFunction space types desc rank elemtype vds-            <*> opaqueArrayShapeFunction desc-            <*> opaqueArrayNewFunction space types desc rank elemtype vds-            <*> opaqueArraySetFunction space types desc rank elemtype vds+            <$> recordArrayIndexFunction space types desc rank elemtype vds+            <*> recordArrayShapeFunction desc vds+            <*> recordArrayNewFunction space types desc rank elemtype vds+            <*> recordArraySetFunction space types desc rank elemtype vds         )  opaqueLibraryFunctions ::@@ -1167,7 +1159,7 @@            int $id:store_opaque($ty:ctx_ty *ctx,                                const $ty:opaque_type *obj, void **p, size_t *n) {-            (void)ctx;+            (void)ctx; (void)obj;             int ret = 0;             $items:store_body             return ret;@@ -1224,17 +1216,23 @@ generateOpaque ::   Space ->   OpaqueTypes ->-  (Name, OpaqueType) ->+  (Name, (OpaqueType, Maybe T.Text)) ->   CompilerM op s (T.Text, Manifest.Type)-generateOpaque space types (desc, ot) = do+generateOpaque space types (desc, (ot, doc)) = do   name <- publicName $ opaqueName desc-  members <- zipWithM field (opaquePayload types ot) [(0 :: Int) ..]+  members <- case ot of+    -- We need to treat arrays of unit specially, because otherwise they would+    -- have no members.+    OpaqueRecordArray rank _ [] ->+      pure [[C.csdecl|typename int64_t shape[$int:rank];|]]+    _ ->+      dummyIfNone <$> zipWithM field (opaquePayload types ot) [(0 :: Int) ..]   libDecl [C.cedecl|struct $id:name { $sdecls:members };|]   (ops, extra_ops) <- opaqueLibraryFunctions space types desc ot   let opaque_type = [C.cty|struct $id:name*|]   pure     ( nameToText desc,-      Manifest.TypeOpaque (typeText opaque_type) ops extra_ops+      Manifest.TypeOpaque (typeText opaque_type) ops extra_ops doc     )   where     field vt@(ValueType _ (Rank r) _) i = do@@ -1244,12 +1242,17 @@           then [C.csdecl|$ty:ct $id:(tupleField i);|]           else [C.csdecl|$ty:ct *$id:(tupleField i);|] +    -- C compilers tend to warn about empty structs, so put in a dummy field if+    -- we have no real fields.+    dummyIfNone [] = [[C.csdecl|char dummy;|]]+    dummyIfNone members = members+ generateAPITypes ::   Space ->   OpaqueTypes ->   CompilerM op s (M.Map T.Text Manifest.Type) generateAPITypes arr_space types@(OpaqueTypes opaques) = do-  mapM_ (findNecessaryArrays . snd) opaques+  mapM_ (findNecessaryArrays . fst . snd) opaques   array_ts <- mapM (generateArray arr_space) . M.toList =<< gets compArrayTypes   opaque_ts <- mapM (generateOpaque arr_space types) opaques   pure $ M.fromList $ catMaybes array_ts <> opaque_ts@@ -1258,8 +1261,6 @@     -- types that allow projection of them.  This is because the     -- projection functions somewhat uglily directly poke around in     -- the innards to increment reference counts.-    findNecessaryArrays (OpaqueType _) =-      pure ()     findNecessaryArrays (OpaqueArray _ _ vts) =       mapM_ (valueTypeToCType Private) vts     findNecessaryArrays (OpaqueRecordArray _ _ fs) =
src/Futhark/CodeGen/Backends/GenericPython.hs view
@@ -511,7 +511,7 @@     opaqueTupleElems opaque_name =       case opaques of         Imp.OpaqueTypes m-          | Just (Imp.OpaqueRecord ts) <- lookup (nameFromText opaque_name) m ->+          | Just (Imp.OpaqueRecord ts, _) <- lookup (nameFromText opaque_name) m ->               -- XXX: might not be tuple.               Tuple $ map (String . p . snd) ts           where@@ -867,7 +867,7 @@       [PyStmt],       (Imp.ExternalValue, PyExp)     )-prepareEntry (Imp.EntryPoint _ result args) (fname, Imp.Function _ outputs inputs _ _) = do+prepareEntry (Imp.EntryPoint _ result args _doc) (fname, Imp.Function _ outputs inputs _ _) = do   let output_paramNames = map (compileName . Imp.paramName) outputs       funTuple = tupleOrSingle $ fmap Var output_paramNames @@ -961,7 +961,7 @@   | otherwise = pure Nothing  entryTypes :: Imp.EntryPoint -> ([T.Text], T.Text)-entryTypes (Imp.EntryPoint _ res args) =+entryTypes (Imp.EntryPoint _ res args _doc) =   (map descArg args, desc res)   where     descArg ((_, u), d) = desc (u, d)@@ -976,7 +976,7 @@   CompilerM op s (Maybe (PyFunDef, T.Text, PyExp)) callEntryFun _ (_, Imp.Function Nothing _ _ _ _) = pure Nothing callEntryFun pre_timing fun@(fname, Imp.Function (Just entry) _ _ _ _) = do-  let Imp.EntryPoint ename _ decl_args = entry+  let Imp.EntryPoint ename _ decl_args _doc = entry   (_, prepare_in, body_bin, _, res) <- prepareEntry entry fun    let str_input = map (readInput . snd) decl_args
src/Futhark/CodeGen/Backends/GenericWASM.hs view
@@ -77,7 +77,7 @@ opaqueToJS :: Imp.OpaqueTypes -> [(String, JSOpaqueType)] opaqueToJS (Imp.OpaqueTypes types) = map convertOne types   where-    convertOne (desc, Imp.OpaqueRecord fields) =+    convertOne (desc, (Imp.OpaqueRecord fields, _)) =       (T.unpack (opaqueName desc), JSOpaqueRecord (map (convertField desc) fields))     convertOne (desc, _) =       (T.unpack (opaqueName desc), JSOpaqueOther)
src/Futhark/CodeGen/Backends/MulticoreWASM.hs view
@@ -67,7 +67,7 @@ fRepMyRep prog =   let Imp.Functions fs = Imp.defFuns prog       function fun = do-        Imp.EntryPoint n res args <- Imp.functionEntry fun+        Imp.EntryPoint n res args _ <- Imp.functionEntry fun         Just $           JSEntryPoint             { name = nameToString n,
src/Futhark/CodeGen/Backends/SequentialWASM.hs view
@@ -61,7 +61,7 @@ fRepMyRep prog =   let Imp.Functions fs = Imp.defFuns prog       function (Imp.Function entry _ _ _ _) = do-        Imp.EntryPoint n res args <- entry+        Imp.EntryPoint n res args _doc <- entry         Just $           JSEntryPoint             { name = nameToString n,
src/Futhark/CodeGen/ImpCode.hs view
@@ -217,7 +217,8 @@ data EntryPoint = EntryPoint   { entryPointName :: Name,     entryPointResults :: (Uniqueness, ExternalValue),-    entryPointArgs :: [((Name, Uniqueness), ExternalValue)]+    entryPointArgs :: [((Name, Uniqueness), ExternalValue)],+    entryPointDocs :: Maybe T.Text   }   deriving (Show) @@ -515,7 +516,7 @@       <+> nestedBlock (pretty code)  instance Pretty EntryPoint where-  pretty (EntryPoint name result args) =+  pretty (EntryPoint name result args _doc) =     stack       [ "name" <+> nestedBlock (dquotes (pretty name)),         "arguments" <+> nestedBlock (stack $ map ppArg args),@@ -777,7 +778,7 @@ declaredIn _ = mempty  instance FreeIn EntryPoint where-  freeIn' (EntryPoint _ res args) =+  freeIn' (EntryPoint _ res args _) =     freeIn' (snd res) <> freeIn' (map snd args)  instance (FreeIn a) => FreeIn (Functions a) where
src/Futhark/CodeGen/ImpGen.hs view
@@ -501,7 +501,7 @@ lookupOpaqueType :: Name -> OpaqueTypes -> OpaqueType lookupOpaqueType v (OpaqueTypes types) =   case lookup v types of-    Just t -> t+    Just (t, _) -> t     Nothing -> error $ "Unknown opaque type: " ++ show v  valueTypeSign :: ValueType -> Signedness@@ -511,9 +511,9 @@ entryPointSignedness _ (TypeTransparent vt) = [valueTypeSign vt] entryPointSignedness types (TypeOpaque desc) =   case lookupOpaqueType desc types of-    OpaqueType vts -> map valueTypeSign vts     OpaqueArray _ _ vts -> map valueTypeSign vts     OpaqueRecordArray _ _ fs -> foldMap (entryPointSignedness types . snd) fs+    OpaqueRecord [] -> [Signed]     OpaqueRecord fs -> foldMap (entryPointSignedness types . snd) fs     OpaqueSum vts _ -> map valueTypeSign vts @@ -525,9 +525,9 @@ entryPointSize _ (TypeTransparent _) = 1 entryPointSize types (TypeOpaque desc) =   case lookupOpaqueType desc types of-    OpaqueType vts -> length vts     OpaqueArray _ _ vts -> length vts     OpaqueRecordArray _ _ fs -> sum $ map (entryPointSize types . snd) fs+    OpaqueRecord [] -> 1     OpaqueRecord fs -> sum $ map (entryPointSize types . snd) fs     OpaqueSum vts _ -> length vts @@ -692,16 +692,16 @@ compileFunDef types (FunDef entry attrs fname rettype params body) =   local (\env -> env {envFunction = name_entry `mplus` Just fname}) $ do     ((outparams, inparams, results, args), body') <- collect' compile-    let entry' = case (name_entry, results, args) of-          (Just name_entry', Just results', Just args') ->-            Just $ Imp.EntryPoint name_entry' results' args'+    let entry' = case (name_entry, results, args, doc_entry) of+          (Just name_entry', Just results', Just args', Just docs') ->+            Just $ Imp.EntryPoint name_entry' results' args' docs'           _ ->             Nothing     emitFunction fname $ Imp.Function entry' outparams inparams attrs body'   where-    (name_entry, params_entry, ret_entry) = case entry of-      Nothing -> (Nothing, Nothing, Nothing)-      Just (x, y, z) -> (Just x, Just y, Just z)+    (name_entry, params_entry, ret_entry, doc_entry) = case entry of+      Nothing -> (Nothing, Nothing, Nothing, Nothing)+      Just (x, y, z, v) -> (Just x, Just y, Just z, Just v)     compile = do       (inparams, arrayds, args) <- compileInParams types params params_entry       (results, outparams, dests) <- compileOutParams types (map fst rettype) ret_entry@@ -1864,6 +1864,17 @@ sArray :: Name -> PrimType -> ShapeBase SubExp -> VName -> LMAD -> ImpM rep r op VName sArray name bt shape mem lmad = do   name' <- newVName name+  when (LMAD.rank lmad /= shapeRank shape) $+    error $+      unlines+        [ "sArray: array "+            <> prettyString name'+            <> " of rank "+            <> show (shapeRank shape)+            <> " with LMAD of rank "+            <> show (LMAD.rank lmad)+            <> "."+        ]   dArray name' bt shape mem lmad   pure name'
src/Futhark/CodeGen/ImpGen/GPU/Base.hs view
@@ -380,20 +380,16 @@     in_block_id = ltid32 - block_id * block_size     ltid32 = kernelLocalThreadId constants     ltid = sExt64 ltid32-    gtid = sExt64 $ kernelGlobalThreadId constants     array_scan = not $ all primType $ lambdaReturnType scan_lam -    readInitial p arr-      | isPrimParam p =-          copyDWIMFix (paramName p) [] (Var arr) [ltid]-      | otherwise =-          copyDWIMFix (paramName p) [] (Var arr) [gtid]+    readInitial p arr =+      copyDWIMFix (paramName p) [] (Var arr) [ltid]      readParam behind p arr       | isPrimParam p =           copyDWIMFix (paramName p) [] (Var arr) [ltid - behind]       | otherwise =-          copyDWIMFix (paramName p) [] (Var arr) [gtid - behind + arrs_full_size]+          copyDWIMFix (paramName p) [] (Var arr) [ltid - behind + arrs_full_size]      writeResult x y arr = do       when (isPrimParam x) $@@ -457,19 +453,11 @@         | otherwise =             sOp $ Imp.ErrorSync Imp.FenceLocal -      block_offset = sExt64 (kernelBlockId constants) * kernelBlockSize constants--      writeBlockResult p arr-        | isPrimParam p =-            copyDWIMFix arr [sExt64 chunk_id] (Var $ paramName p) []-        | otherwise =-            copyDWIMFix arr [block_offset + sExt64 chunk_id] (Var $ paramName p) []+      writeBlockResult p arr =+        copyDWIMFix arr [sExt64 chunk_id] (Var $ paramName p) [] -      readPrevBlockResult p arr-        | isPrimParam p =-            copyDWIMFix (paramName p) [] (Var arr) [sExt64 chunk_id - 1]-        | otherwise =-            copyDWIMFix (paramName p) [] (Var arr) [block_offset + sExt64 chunk_id - 1]+      readPrevBlockResult p arr =+        copyDWIMFix (paramName p) [] (Var arr) [sExt64 chunk_id - 1]    doInChunkScan seg_flag ltid_in_bounds lam   barrier@@ -480,7 +468,7 @@       sWhen is_first_block $         forM_ (zip x_params arrs) $ \(x, arr) ->           unless (isPrimParam x) $-            copyDWIMFix arr [arrs_full_size + block_offset + sExt64 chunk_size + ltid] (Var $ paramName x) []+            copyDWIMFix arr [arrs_full_size + sExt64 chunk_size + ltid] (Var $ paramName x) []      barrier @@ -509,9 +497,9 @@           unless (isPrimParam x) $             copyDWIMFix               arr-              [arrs_full_size + block_offset + ltid]+              [arrs_full_size + ltid]               (Var arr)-              [arrs_full_size + block_offset + sExt64 chunk_size + ltid]+              [arrs_full_size + sExt64 chunk_size + ltid]      barrier @@ -553,7 +541,7 @@       forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) ->         if isPrimParam y           then copyDWIMFix arr [ltid] (Var $ paramName y) []-          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + block_offset + ltid]+          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + ltid]    barrier
src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs view
@@ -218,7 +218,7 @@     all isPrimSegBinOp segbinops =       noncommPrimSegRedInterms   | otherwise =-      generalSegRedInterms False tblock_id tblock_size segbinops+      generalSegRedInterms tblock_id tblock_size segbinops   where     params = map paramOf segbinops @@ -272,24 +272,18 @@     forAccumLM2D acc ls f = mapAccumLM (mapAccumLM f) acc ls  generalSegRedInterms ::-  Bool ->   Imp.TExp Int64 ->   SubExp ->   [SegBinOp GPUMem] ->   InKernelGen [SegRedIntermediateArrays]-generalSegRedInterms segmented tblock_id tblock_size segbinops =+generalSegRedInterms tblock_id tblock_size segbinops =   fmap (map GeneralSegRedInterms) . forM (map paramOf segbinops) . mapM $ \p ->     case paramDec p of-      MemArray pt shape _ (ArrayIn mem ixfun) -> do+      MemArray pt shape _ (ArrayIn mem _) -> do         let shape' = Shape [tblock_size] <> shape         let shape_E = map pe64 $ shapeDims shape'         sArray ("red_arr_" <> nameFromText (prettyText pt)) pt shape' mem $-          -- This 'segmented' thing here is a hack, related to #2227.-          -- There absolutely must be some unifying principle we are-          -- missing.-          if segmented-            then ixfun-            else LMAD.iota (tblock_id * product shape_E) shape_E+          LMAD.iota (tblock_id * product shape_E) shape_E       _ -> do         let pt = elemType $ paramType p             shape = Shape [tblock_size]@@ -426,10 +420,10 @@    sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do     constants <- kernelConstants <$> askEnv-    let tblock_id = kernelBlockSize constants+    let tblock_id = kernelBlockId constants         ltid = sExt64 $ kernelLocalThreadId constants -    interms <- generalSegRedInterms True tblock_id tblock_size_se segbinops+    interms <- generalSegRedInterms (sExt64 tblock_id) tblock_size_se segbinops     let reds_arrs = map blockRedArrs interms      -- We probably do not have enough actual threadblocks to cover the@@ -449,10 +443,9 @@        let in_bounds =             map_body_cont $ \red_res ->-              sComment "save results to be reduced" $ do-                let red_dests = map (,[ltid]) (concat reds_arrs)-                forM2_ red_dests red_res $ \(d, d_is) (res, res_is) ->-                  copyDWIMFix d d_is res res_is+              sComment "save results to be reduced" $+                forM2_ (concat reds_arrs) red_res $ \d (res, res_is) ->+                  copyDWIMFix d [ltid] res res_is           out_of_bounds =             forM2_ segbinops reds_arrs $ \(SegBinOp _ _ nes _) red_arrs ->               forM2_ red_arrs nes $ \arr ne ->
src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs view
@@ -20,11 +20,11 @@ -- Aggressively try to reuse memory for different SegBinOps, because -- we will run them sequentially after another. makeLocalArrays ::+  Imp.TExp Int64 ->   Count BlockSize SubExp ->-  SubExp ->   [SegBinOp GPUMem] ->   InKernelGen [[VName]]-makeLocalArrays (Count tblock_size) num_threads scans = do+makeLocalArrays tblock_id (Count tblock_size) scans = do   (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty   let maxSize sizes = Imp.bytes $ L.foldl' sMax64 1 $ map Imp.unCount sizes   forM_ mems_and_sizes $ \(sizes, mem) ->@@ -34,21 +34,21 @@     onScan (SegBinOp _ scan_op nes _) = do       let (scan_x_params, _scan_y_params) =             splitAt (length nes) $ lambdaParams scan_op-      (arrs, used_mems) <- fmap unzip $-        forM scan_x_params $ \p ->-          case paramDec p of-            MemArray pt shape _ (ArrayIn mem _) -> do-              let shape' = Shape [num_threads] <> shape-              arr <--                lift . sArray "scan_arr" pt shape' mem $-                  LMAD.iota 0 (map pe64 $ shapeDims shape')-              pure (arr, [])-            _ -> do-              let pt = elemType $ paramType p-                  shape = Shape [tblock_size]-              (sizes, mem') <- getMem pt shape-              arr <- lift $ sArrayInMem "scan_arr" pt shape mem'-              pure (arr, [(sizes, mem')])+      (arrs, used_mems) <- fmap unzip . forM scan_x_params $ \p ->+        case paramDec p of+          MemArray pt shape _ (ArrayIn mem _) -> do+            let shape' = Shape [tblock_size] <> shape+            let shape_E = map pe64 $ shapeDims shape'+            arr <-+              lift . sArray "scan_arr" pt shape' mem $+                LMAD.iota (tblock_id * product shape_E) shape_E+            pure (arr, [])+          _ -> do+            let pt = elemType $ paramType p+                shape = Shape [tblock_size]+            (sizes, mem') <- getMem pt shape+            arr <- lift $ sArrayInMem "scan_arr" pt shape mem'+            pure (arr, [(sizes, mem')])       modify (<> concat used_mems)       pure arrs @@ -68,11 +68,8 @@  type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) -localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64-localArrayIndex constants t =-  if primType t-    then sExt64 (kernelLocalThreadId constants)-    else sExt64 (kernelGlobalThreadId constants)+localArrayIndex :: KernelConstants -> Imp.TExp Int64+localArrayIndex constants = sExt64 (kernelLocalThreadId constants)  barrierFor :: Lambda GPUMem -> (Bool, Imp.Fence, InKernelGen ()) barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence)@@ -174,7 +171,8 @@    sKernelThread "scan_stage1" (segFlat space) (defKernelAttrs num_tblocks tblock_size) $ do     constants <- kernelConstants <$> askEnv-    all_local_arrs <- makeLocalArrays tblock_size (tvSize num_threads) scans+    let tblock_id = kernelBlockId constants+    all_local_arrs <- makeLocalArrays (sExt64 tblock_id) tblock_size scans      -- The variables from scan_op will be used for the carry and such     -- in the big chunking loop.@@ -229,8 +227,7 @@       forM_ (zip3 per_scan_pes scans all_local_arrs) $         \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) ->           sComment "do one intra-group scan operation" $ do-            let rets = lambdaReturnType scan_op-                scan_x_params = xParams scan+            let scan_x_params = xParams scan                 (array_scan, fence, barrier) = barrierFor scan_op              when array_scan barrier@@ -249,9 +246,9 @@                sComment "combine with carry and write to shared memory" $                 compileStms mempty (bodyStms $ lambdaBody scan_op) $-                  forM_ (zip3 rets local_arrs $ map resSubExp $ bodyResult $ lambdaBody scan_op) $-                    \(t, arr, se) ->-                      copyDWIMFix arr [localArrayIndex constants t] se []+                  forM_ (zip local_arrs $ map resSubExp $ bodyResult $ lambdaBody scan_op) $+                    \(arr, se) ->+                      copyDWIMFix arr [localArrayIndex constants] se []                let crossesSegment' = do                     f <- crossesSegment@@ -273,12 +270,12 @@                sComment "threads in bounds write partial scan result" $                 sWhen in_bounds $-                  forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->+                  forM_ (zip pes local_arrs) $ \(pe, arr) ->                     copyDWIMFix                       pe                       (map Imp.le64 gtids ++ vec_is)                       (Var arr)-                      [localArrayIndex constants t]+                      [localArrayIndex constants]                barrier @@ -288,12 +285,7 @@                         (paramName p)                         []                         (Var arr)-                        [ if primType $ paramType p-                            then sExt64 (kernelBlockSize constants) - 1-                            else-                              (sExt64 (kernelBlockId constants) + 1)-                                * sExt64 (kernelBlockSize constants)-                                - 1+                        [ sExt64 (kernelBlockSize constants) - 1                         ]                   load_neutral =                     forM_ (zip nes scan_x_params) $ \(ne, p) ->@@ -346,9 +338,9 @@    sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) stage2_tblock_size) $ do     constants <- kernelConstants <$> askEnv-    per_scan_local_arrs <- makeLocalArrays stage2_tblock_size (tvSize stage1_num_threads) scans-    let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans-        per_scan_pes = segBinOpChunks scans scan_out+    let tblock_id = kernelBlockId constants+    per_scan_local_arrs <- makeLocalArrays (sExt64 tblock_id) stage2_tblock_size scans+    let per_scan_pes = segBinOpChunks scans scan_out      -- Declare lambda params and initialise carries (xParams) to the     -- neutral element.  For scalar scans these persist across chunk@@ -368,8 +360,8 @@       -- Construct segment indices.       zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx -      forM_ (L.zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $-        \(scan@(SegBinOp _ scan_op nes vec_shape), local_arrs, rets, pes) ->+      forM_ (L.zip3 scans per_scan_local_arrs per_scan_pes) $+        \(scan@(SegBinOp _ scan_op nes vec_shape), local_arrs, pes) ->           sComment "do one stage-2 scan chunk" $ do             let (array_scan, fence, barrier) = barrierFor scan_op                 scan_x_params = xParams scan@@ -429,9 +421,9 @@               -- incorporates the inter-chunk carry; other threads have               -- neutral in xParams, so the op is a no-op for them.               compileStms mempty (bodyStms $ lambdaBody scan_op) $-                forM_ (zip3 rets local_arrs $ map resSubExp $ bodyResult $ lambdaBody scan_op) $-                  \(t, arr, se) ->-                    copyDWIMFix arr [localArrayIndex constants t] se []+                forM_ (zip local_arrs $ map resSubExp $ bodyResult $ lambdaBody scan_op) $+                  \(arr, se) ->+                    copyDWIMFix arr [localArrayIndex constants] se []                sOp $ Imp.ErrorSync fence @@ -454,12 +446,12 @@                sComment "threads in bounds write scanned carries" $                 sWhen in_bounds $-                  forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->+                  forM_ (zip pes local_arrs) $ \(pe, arr) ->                     copyDWIMFix                       pe                       glob_is                       (Var arr)-                      [localArrayIndex constants t]+                      [localArrayIndex constants]                barrier
src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs view
@@ -107,6 +107,8 @@   KernelBody MCMem ->   MulticoreGen () atomicHistogram pat space histops kbody = do+  emit $ Imp.DebugPrint "$ SegHist atomicHistogram" Nothing+   let (is, ns) = unzip $ unSegSpace space       ns_64 = map pe64 ns   let num_red_res = length histops + sum (map (length . histNeutral) histops)@@ -183,7 +185,7 @@   KernelBody MCMem ->   MulticoreGen () subHistogram pat space histops num_histos kbody = do-  emit $ Imp.DebugPrint "subHistogram segHist" Nothing+  emit $ Imp.DebugPrint "$ SegHist subHistogram" Nothing    let (is, ns) = unzip $ unSegSpace space       ns_64 = map pe64 ns@@ -226,53 +228,55 @@        pure op_local_subhistograms -    inISPC $-      generateChunkLoop "SegRed" Vectorized $ \i -> do-        zipWithM_ dPrimV_ is $ unflattenIndex ns_64 i-        compileStms mempty (bodyStms kbody) $ do-          let (red_res, map_res) =-                splitFromEnd (length map_pes) $-                  map kernelResultSubExp $-                    bodyResult kbody+    inISPC . generateChunkLoop "SegRed" Vectorized $ \i -> do+      zipWithM_ dPrimV_ is $ unflattenIndex ns_64 i+      compileStms mempty (bodyStms kbody) $ do+        let (red_res, map_res) =+              splitFromEnd (length map_pes) $+                map kernelResultSubExp $+                  bodyResult kbody -          sComment "save map-out results" $-            forM_ (zip map_pes map_res) $ \(pe, res) ->-              copyDWIMFix (patElemName pe) (map Imp.le64 is) res []+        sComment "save map-out results" $+          forM_ (zip map_pes map_res) $ \(pe, res) ->+            copyDWIMFix (patElemName pe) (map Imp.le64 is) res [] -          forM_ (zip3 histops local_subhistograms (splitHistResults histops red_res)) $-            \( histop@(HistOp dest_shape _ _ _ shape _),-               histop_subhistograms,-               (bucket, vs')-               ) -> do-                histop' <- renameHistop histop+        forM_ (zip3 histops local_subhistograms (splitHistResults histops red_res)) $+          \( histop@(HistOp dest_shape _ _ _ shape _),+             histop_subhistograms,+             (bucket, vs')+             ) -> do+              histop' <- renameHistop histop -                let bucket' = map pe64 bucket-                    dest_shape' = map pe64 $ shapeDims dest_shape-                    acc_params' = (lambdaParams . histOp) histop'-                    vs_params' = takeLast (length vs') $ lambdaParams $ histOp histop'+              let bucket' = map pe64 bucket+                  dest_shape' = map pe64 $ shapeDims dest_shape+                  acc_params' = (lambdaParams . histOp) histop'+                  vs_params' = takeLast (length vs') $ lambdaParams $ histOp histop' -                generateUniformizeLoop $ \j ->-                  sComment "perform updates" $ do-                    -- Create new set of uniform buckets-                    -- That is extract each bucket from a SIMD vector lane-                    extract_buckets <- mapM (dPrimSV "extract_bucket" . (primExpType . untyped)) bucket'-                    forM_ (zip extract_buckets bucket') $ \(x, y) ->-                      emit $ Imp.Op $ Imp.ExtractLane (tvVar x) (untyped y) (untyped j)-                    let bucket'' = map tvExp extract_buckets-                        bucket_in_bounds =-                          inBounds (Slice (map DimFix bucket'')) dest_shape'-                    sWhen bucket_in_bounds $ do-                      genHistOpParams histop'-                      sLoopNest shape $ \is' -> do-                        -- read values vs and perform lambda writing result back to is-                        forM_ (zip vs_params' vs') $ \(p, res) ->-                          ifPrimType (paramType p) $ \pt -> do+              generateUniformizeLoop $ \j ->+                sComment "perform updates" $ do+                  -- Create new set of uniform buckets+                  -- That is extract each bucket from a SIMD vector lane+                  extract_buckets <- mapM (dPrimSV "extract_bucket" . (primExpType . untyped)) bucket'+                  forM_ (zip extract_buckets bucket') $ \(x, y) ->+                    emit $ Imp.Op $ Imp.ExtractLane (tvVar x) (untyped y) (untyped j)+                  let bucket'' = map tvExp extract_buckets+                      bucket_in_bounds =+                        inBounds (Slice (map DimFix bucket'')) dest_shape'+                  sWhen bucket_in_bounds $ do+                    genHistOpParams histop'+                    sLoopNest shape $ \is' -> do+                      -- read values vs and perform lambda writing result back to is+                      forM_ (zip vs_params' vs') $ \(p, res) ->+                        case paramType p of+                          Prim pt -> do                             -- Hack to copy varying load into uniform result variable                             tmp <- dPrimS "tmp" pt                             copyDWIMFix tmp [] res is'                             extractVectorLane j . pure $                               Imp.SetScalar (paramName p) (toExp' pt tmp)-                        updateHisto histop' histop_subhistograms (bucket'' ++ is') j acc_params'+                          _ ->+                            copyDWIMFix (paramName p) [] res is'+                      updateHisto histop' histop_subhistograms (bucket'' ++ is') j acc_params'      -- Copy the task-local subhistograms to the global subhistograms,     -- where they will be combined.@@ -316,8 +320,6 @@     emit $ Imp.Op $ Imp.SegOp "seghist_red" free_params_red red_task Nothing mempty scheduler_info   where     segment_dims = init $ unSegSpace space-    ifPrimType (Prim pt) f = f pt-    ifPrimType _ _ = pure ()  -- Note: This isn't currently used anywhere. -- This implementation for a Segmented Hist only
src/Futhark/Construct.hs view
@@ -68,6 +68,7 @@      -- * Monadic expression builders     eSubExp,+    eVar,     eParam,     eMatch',     eMatch,@@ -108,6 +109,7 @@     fullSliceNum,     isFullSlice,     sliceAt,+    iota64,      -- * Result types     instantiateShapes,@@ -204,6 +206,13 @@   m (Exp (Rep m)) eSubExp = pure . BasicOp . SubExp +-- | Turn a variable into a monad expression, through 'eSubExp'.+eVar ::+  (MonadBuilder m) =>+  VName ->+  m (Exp (Rep m))+eVar = eSubExp . Var+ -- | Treat a parameter as a monadic expression. eParam ::   (MonadBuilder m) =>@@ -574,6 +583,11 @@ sliceAt :: Type -> Int -> [DimIndex SubExp] -> Slice SubExp sliceAt t n slice =   fullSlice t $ map sliceDim (take n $ arrayDims t) ++ slice++-- | Produce a straightforward `Int64` `Iota` of the given length with offset 0+-- and stride 1.+iota64 :: SubExp -> Exp rep+iota64 n = BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64  -- | Like 'fullSlice', but the dimensions are simply numeric. fullSliceNum :: (Num d) => [d] -> [DimIndex d] -> Slice d
src/Futhark/IR/MC/Op.hs view
@@ -90,7 +90,7 @@   opAliases (ParOp _ op) = opAliases op   opAliases (OtherOp op) = opAliases op -  consumedInOp (ParOp _ op) = consumedInOp op+  consumedInOp (ParOp seq_op op) = maybe mempty consumedInOp seq_op <> consumedInOp op   consumedInOp (OtherOp op) = consumedInOp op  instance (CanBeAliased op) => CanBeAliased (MCOp op) where
src/Futhark/IR/Mem/LMAD.hs view
@@ -17,6 +17,7 @@     coerce,     permute,     shape,+    rank,     substitute,     iota,     equivalent,@@ -276,6 +277,10 @@ -- | Shape of an LMAD. shape :: LMAD num -> Shape num shape = map ldShape . dims++-- | Rank of an LMAD.+rank :: LMAD num -> Int+rank = length . shape  iotaStrided ::   (IntegralExp num) =>
src/Futhark/IR/Parse.hs view
@@ -659,12 +659,16 @@ pEntry :: Parser EntryPoint pEntry =   parens $-    (,,)+    (,,,)       <$> (nameFromText <$> pStringLiteral)       <* pComma       <*> pEntryPointInputs       <* pComma       <*> pEntryPointResult+      <*> choice+        [ pComma *> (Just <$> pStringLiteral),+          pure Nothing+        ]   where     pEntryPointInputs = braces (pEntryPointInput `sepBy` pComma)     pEntryPointInput =@@ -686,11 +690,11 @@   FunDef entry attrs fname ret fparams     <$> (pEqual *> braces (pBody pr)) -pOpaqueType :: Parser (Name, OpaqueType)+pOpaqueType :: Parser (Name, (OpaqueType, Maybe T.Text)) pOpaqueType =   (,)     <$> (keyword "type" *> (nameFromText <$> pStringLiteral) <* pEqual)-    <*> choice [pRecord, pSum, pOpaque, pRecordArray, pOpaqueArray]+    <*> ((,Nothing) <$> choice [pRecord, pSum, pRecordArray, pOpaqueArray])   where     pFieldName = choice [pName, nameFromString . show <$> pInt]     pField = (,) <$> pFieldName <* pColon <*> pEntryPointType@@ -711,8 +715,6 @@               <*> many pVariant           ) -    pOpaque = keyword "opaque" $> OpaqueType <*> braces (many pValueType)-     pRecordArray =       keyword "record_array"         $> OpaqueRecordArray@@ -841,15 +843,19 @@     pVJP =       parens $         SOAC.VJP-          <$> braces (pSubExp `sepBy` pComma)+          <$> pShape           <* pComma           <*> braces (pSubExp `sepBy` pComma)           <* pComma+          <*> braces (pSubExp `sepBy` pComma)+          <* pComma           <*> pLambda pr     pJVP =       parens $         SOAC.JVP-          <$> braces (pSubExp `sepBy` pComma)+          <$> pShape+          <* pComma+          <*> braces (pSubExp `sepBy` pComma)           <* pComma           <*> braces (pSubExp `sepBy` pComma)           <* pComma
src/Futhark/IR/Pretty.hs view
@@ -14,6 +14,7 @@ import Data.Foldable (toList) import Data.List.NonEmpty (NonEmpty (..)) import Data.Maybe+import Data.Text qualified as T import Futhark.IR.Syntax import Futhark.Util.Pretty @@ -413,21 +414,18 @@     where       fun = case entry of         Nothing -> "fun"-        Just (p_name, p_entry, ret_entry) ->+        Just (p_name, p_entry, ret_entry, doc_entry) ->           "entry"             <> (parens . align)               ( "\""                   <> pretty p_name                   <> "\""-                  <> comma-                    </> ppTupleLines' (map pretty p_entry)-                  <> comma-                    </> pretty ret_entry+                  <> comma </> ppTupleLines' (map pretty p_entry)+                  <> comma </> pretty ret_entry+                  <> maybe mempty ((comma </>) . pretty . show) doc_entry               )  instance Pretty OpaqueType where-  pretty (OpaqueType ts) =-    "opaque" <+> nestedBlock (stack $ map pretty ts)   pretty (OpaqueRecord fs) =     "record" <+> nestedBlock (stack $ map p fs)     where@@ -449,7 +447,14 @@ instance Pretty OpaqueTypes where   pretty (OpaqueTypes ts) = "types" <+> nestedBlock (stack $ map p ts)     where-      p (name, t) = "type" <+> dquotes (pretty name) <+> equals <+> pretty t+      p (name, (t, doc)) =+        ( case doc of+            Nothing -> mempty+            Just doc' -> mconcat $ map wrap (T.lines doc')+        )+          <> "type" <+> dquotes (pretty name) <+> equals <+> pretty t+        where+          wrap x = "-- " <> pretty x <> line  instance (PrettyRep rep) => Pretty (Prog rep) where   pretty (Prog types consts funs) =
src/Futhark/IR/Prop.hs view
@@ -122,6 +122,7 @@     safeBasicOp Manifest {} = True     safeBasicOp Iota {} = True     safeBasicOp Replicate {} = True+    safeBasicOp (Index _ slice) = sliceShape slice /= mempty     safeBasicOp _ = False safeExp (Loop _ _ body) = safeBody body safeExp (Apply fname _ _ _) =
src/Futhark/IR/SOACS/SOAC.hs view
@@ -79,9 +79,9 @@     -- The final lambda produces indexes and values for the 'HistOp's.     Hist SubExp [VName] [HistOp rep] (Lambda rep)   | -- FIXME: this should not be here-    JVP [SubExp] [SubExp] (Lambda rep)+    JVP Shape [SubExp] [SubExp] (Lambda rep)   | -- FIXME: this should not be here-    VJP [SubExp] [SubExp] (Lambda rep)+    VJP Shape [SubExp] [SubExp] (Lambda rep)   | -- FIXME: this should not be here     WithVJP [SubExp] (Lambda rep) (Lambda rep)   | -- | A combination of scan, reduction, and map.  The first@@ -401,14 +401,16 @@   SOACMapper frep trep m ->   SOAC frep ->   m (SOAC trep)-mapSOACM tv (JVP args vec lam) =+mapSOACM tv (JVP shape args vec lam) =   JVP-    <$> mapM (mapOnSOACSubExp tv) args+    <$> mapM (mapOnSOACSubExp tv) shape+    <*> mapM (mapOnSOACSubExp tv) args     <*> mapM (mapOnSOACSubExp tv) vec     <*> mapOnSOACLambda tv lam-mapSOACM tv (VJP args vec lam) =+mapSOACM tv (VJP shape args vec lam) =   VJP-    <$> mapM (mapOnSOACSubExp tv) args+    <$> mapM (mapOnSOACSubExp tv) shape+    <*> mapM (mapOnSOACSubExp tv) args     <*> mapM (mapOnSOACSubExp tv) vec     <*> mapOnSOACLambda tv lam mapSOACM tv (WithVJP args lam0 lam1) =@@ -509,10 +511,10 @@  -- | The type of a SOAC. soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [Type]-soacType (JVP _ _ lam) =-  lambdaReturnType lam ++ lambdaReturnType lam-soacType (VJP _ _ lam) =-  lambdaReturnType lam ++ map paramType (lambdaParams lam)+soacType (JVP shape _ _ lam) =+  lambdaReturnType lam ++ map (`arrayOfShape` shape) (lambdaReturnType lam)+soacType (VJP shape _ _ lam) =+  lambdaReturnType lam ++ map ((`arrayOfShape` shape) . paramType) (lambdaParams lam) soacType (WithVJP _ lam _) =   lambdaReturnType lam soacType (Stream outersize _ accs lam) =@@ -561,10 +563,10 @@   HistOp w rf dests nes $ f lam  instance CanBeAliased SOAC where-  addOpAliases aliases (JVP args vec lam) =-    JVP args vec (Alias.analyseLambda aliases lam)-  addOpAliases aliases (VJP args vec lam) =-    VJP args vec (Alias.analyseLambda aliases lam)+  addOpAliases aliases (JVP shape args vec lam) =+    JVP shape args vec (Alias.analyseLambda aliases lam)+  addOpAliases aliases (VJP shape args vec lam) =+    VJP shape args vec (Alias.analyseLambda aliases lam)   addOpAliases aliases (WithVJP args lam lam_adj) =     WithVJP       args@@ -618,12 +620,12 @@       concatIndicesToEachValue is vs =         let is_flat = mconcat is          in map (is_flat <>) vs-  opDependencies (JVP args vec lam) =+  opDependencies (JVP _ args vec lam) =     mconcat $       replicate 2 $         lambdaDependencies mempty lam $           zipWith (<>) (map depsOf' args) (map depsOf' vec)-  opDependencies (VJP args vec lam) =+  opDependencies (VJP _ args vec lam) =     lambdaDependencies       mempty       lam@@ -708,26 +710,32 @@  -- | Type-check a SOAC. typeCheckSOAC :: (TC.Checkable rep) => SOAC (Aliases rep) -> TC.TypeM rep ()-typeCheckSOAC (VJP args vec lam) = do+typeCheckSOAC (VJP shape args vec lam) = do+  mapM_ (TC.require $ Prim int64) shape   args' <- mapM TC.checkArg args   TC.checkLambda lam $ map TC.noArgAliases args'   vec_ts <- mapM TC.checkSubExp vec-  unless (vec_ts == lambdaReturnType lam) $+  unless (vec_ts == map (`arrayOfShape` shape) (lambdaReturnType lam)) $     TC.bad . TC.TypeError . docText $       "Return type"         </> PP.indent 2 (pretty (lambdaReturnType lam))-        </> "does not match type of seed vector"+        </> "inconsistent with type of seed vector"         </> PP.indent 2 (pretty vec_ts)-typeCheckSOAC (JVP args vec lam) = do+        </> "with shape"+        </> PP.indent 2 (pretty shape)+typeCheckSOAC (JVP shape args vec lam) = do+  mapM_ (TC.require $ Prim int64) shape   args' <- mapM TC.checkArg args   TC.checkLambda lam $ map TC.noArgAliases args'   vec_ts <- mapM TC.checkSubExp vec-  unless (vec_ts == map TC.argType args') $+  unless (vec_ts == map ((`arrayOfShape` shape) . TC.argType) args') $     TC.bad . TC.TypeError . docText $       "Parameter type"         </> PP.indent 2 (pretty $ map TC.argType args')         </> "does not match type of seed vector"         </> PP.indent 2 (pretty vec_ts)+        </> "with shape"+        </> PP.indent 2 (pretty shape) typeCheckSOAC (WithVJP args lam lam_adj) = do   args' <- mapM TC.checkArg args   TC.checkLambda lam $ map TC.noArgAliases args'@@ -748,7 +756,6 @@     chunk : _ -> pure chunk     [] -> TC.bad $ TC.TypeError "Stream lambda without parameters."   let asArg t = (t, mempty)-      inttp = Prim int64       lamarrs' = map (`setOuterSize` Var (paramName chunk)) arrargs       acc_len = length accexps       lamrtp = take acc_len $ lambdaReturnType lam@@ -758,7 +765,7 @@   -- just get the dflow of lambda on the fakearg, which does not alias   -- arr, so we can later check that aliases of arr are not used inside lam.   let fake_lamarrs' = map asArg lamarrs'-  TC.checkLambda lam $ asArg inttp : accargs ++ fake_lamarrs'+  TC.checkLambda lam $ asArg (Prim int64) : accargs ++ fake_lamarrs' typeCheckSOAC (Hist w arrs ops bucket_fun) = do   TC.require (Prim int64) w @@ -851,10 +858,10 @@   pure red_nes'  instance RephraseOp SOAC where-  rephraseInOp r (VJP args vec lam) =-    VJP args vec <$> rephraseLambda r lam-  rephraseInOp r (JVP args vec lam) =-    JVP args vec <$> rephraseLambda r lam+  rephraseInOp r (VJP shape args vec lam) =+    VJP shape args vec <$> rephraseLambda r lam+  rephraseInOp r (JVP shape args vec lam) =+    JVP shape args vec <$> rephraseLambda r lam   rephraseInOp r (WithVJP args lam lam_adj) =     WithVJP args <$> rephraseLambda r lam <*> rephraseLambda r lam_adj   rephraseInOp r (Stream w arrs acc lam) =@@ -882,9 +889,9 @@   Scan <$> rephraseLambda r op <*> pure nes  instance (OpMetrics (Op rep)) => OpMetrics (SOAC rep) where-  opMetrics (VJP _ _ lam) =+  opMetrics (VJP _ _ _ lam) =     inside "VJP" $ lambdaMetrics lam-  opMetrics (JVP _ _ lam) =+  opMetrics (JVP _ _ _ lam) =     inside "JVP" $ lambdaMetrics lam   opMetrics (WithVJP _ lam lam_adj) = do     inside "WithVJP" $ lambdaMetrics lam@@ -901,19 +908,21 @@       lambdaMetrics post_lam  instance (PrettyRep rep) => PP.Pretty (SOAC rep) where-  pretty (VJP args vec lam) =+  pretty (VJP shape args vec lam) =     "vjp"       <> parens         ( PP.align $-            PP.braces (commasep $ map pretty args)+            pretty shape+              <> comma </> PP.braces (commasep $ map pretty args)               <> comma </> PP.braces (commasep $ map pretty vec)               <> comma </> pretty lam         )-  pretty (JVP args vec lam) =+  pretty (JVP shape args vec lam) =     "jvp"       <> parens         ( PP.align $-            PP.braces (commasep $ map pretty args)+            pretty shape+              <> comma </> PP.braces (commasep $ map pretty args)               <> comma </> PP.braces (commasep $ map pretty vec)               <> comma </> pretty lam         )@@ -1017,7 +1026,7 @@   Doc ann ppHist w arrs ops bucket_fun =   "hist"-    <> parens+    <> (parens . align)       ( pretty w           <> comma             </> ppTuple' (map pretty arrs)
src/Futhark/IR/SOACS/Simplify.hs view
@@ -91,16 +91,18 @@ simplifySOAC ::   (Simplify.SimplifiableRep rep) =>   Simplify.SimplifyOp rep (SOAC (Wise rep))-simplifySOAC (VJP arr vec lam) = do-  (lam', hoisted) <- Engine.simplifyLambda mempty lam+simplifySOAC (VJP shape arr vec lam) = do+  shape' <- traverse Engine.simplify shape   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec-  pure (VJP arr' vec' lam', hoisted)-simplifySOAC (JVP arr vec lam) = do   (lam', hoisted) <- Engine.simplifyLambda mempty lam+  pure (VJP shape' arr' vec' lam', hoisted)+simplifySOAC (JVP shape arr vec lam) = do+  shape' <- traverse Engine.simplify shape   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec-  pure (JVP arr' vec' lam', hoisted)+  (lam', hoisted) <- Engine.simplifyLambda mempty lam+  pure (JVP shape' arr' vec' lam', hoisted) simplifySOAC (WithVJP args lam lam_adj) = do   args' <- mapM Engine.simplify args   (lam', hoisted) <- Engine.simplifyLambda mempty lam@@ -412,6 +414,10 @@     Just (used_arrs, map_lam') <- remove map_lam arrs =       Simplify . auxing aux . letBind pat . Op $         soacOp (Screma w used_arrs (ScremaForm map_lam' scan reduce post_lam))+  | Just (Hist w arrs ops map_lam :: SOAC rep) <- asSOAC op,+    Just (used_arrs, map_lam') <- remove map_lam arrs =+      Simplify . auxing aux . letBind pat . Op $+        soacOp (Hist w used_arrs ops map_lam')   where     used_in_body map_lam = freeIn $ lambdaBody map_lam     usedInput map_lam (param, _) = paramName param `nameIn` used_in_body map_lam@@ -789,10 +795,10 @@   | Just (Screma w arrs form) <- asSOAC op,     Constant (IntValue (Int64Value k)) <- w,     "unroll" `inAttrs` stmAuxAttrs aux =-      -- We unroll maps in a more direct way, and pass everything else on to-      -- general sequentialisation.+      -- We unroll maps over non-accumulators in a more direct way, and pass+      -- everything else on to general sequentialisation.       case isMapSOAC form of-        Just map_lam -> Simplify $ do+        Just map_lam | not $ any isAcc $ lambdaReturnType map_lam -> Simplify $ do           arrs_elems <- fmap transpose . forM [0 .. k - 1] $ \i -> do             map_lam' <- renameLambda map_lam             eLambda map_lam' $ map (`eIndex` [eSubExp (constant i)]) arrs
src/Futhark/IR/SegOp.hs view
@@ -969,6 +969,31 @@   where     bound_here = namesFromList $ M.keys $ scopeOfSegSpace space +-- | We are willing to hoist potentially unsafe statements out of segops, but+-- they must be protected by adding a branch on top of them.+protectSegOpHoisted ::+  (Engine.SimplifiableRep rep) =>+  SegSpace ->+  Engine.SimpleM rep (a, b, Stms (Wise rep)) ->+  Engine.SimpleM rep (a, b, Stms (Wise rep))+protectSegOpHoisted space m = do+  (x, y, stms) <- m+  ops <- asks $ Engine.protectHoistedOpS . fst+  stms' <- runBuilder_ $ do+    if not $ all (safeExp . stmExp) stms+      then do+        is_nonempty <- checkIfNonEmpty+        mapM_ (Engine.protectIf ops (not . safeExp) is_nonempty) stms+      else addStms stms+  pure (x, y, stms')+  where+    checkIfNonEmpty = do+      segop_size <-+        letSubExp "segop_size"+          =<< foldBinOp (Mul Int64 OverflowWrap) (intConst Int64 1) (segSpaceDims space)+      letSubExp "segop_nonempty" . BasicOp $+        CmpOp (CmpSlt Int64) (intConst Int64 0) segop_size+ simplifyKernelBody ::   (Engine.SimplifiableRep rep, BodyDec rep ~ ()) =>   SegSpace ->@@ -979,7 +1004,8 @@    -- Ensure we do not try to use anything that is consumed in the result.   (body_res, body_stms, hoisted) <--    Engine.localVtable (segSpaceSymbolTable space)+    protectSegOpHoisted space+      . Engine.localVtable (segSpaceSymbolTable space)       . Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True})       . Engine.enterLoop       $ Engine.blockIf blocker stms
src/Futhark/IR/Syntax.hs view
@@ -610,7 +610,7 @@  -- | Information about the inputs and outputs (return value) of an entry -- point.-type EntryPoint = (Name, [EntryParam], EntryResult)+type EntryPoint = (Name, [EntryParam], EntryResult, Maybe T.Text)  -- | An entire Futhark program. data Prog rep = Prog
src/Futhark/IR/Syntax/Core.hs view
@@ -625,8 +625,7 @@  -- | The representation of an opaque type. data OpaqueType-  = OpaqueType [ValueType]-  | -- | Note that the field ordering here denote the actual+  = -- | Note that the field ordering here denote the actual     -- representation - make sure it is preserved.     OpaqueRecord [(Name, EntryPointType)]   | -- | Constructor ordering also denotes representation, in that the@@ -646,7 +645,7 @@   deriving (Eq, Ord, Show)  -- | Names of opaque types and their representation.-newtype OpaqueTypes = OpaqueTypes [(Name, OpaqueType)]+newtype OpaqueTypes = OpaqueTypes [(Name, (OpaqueType, Maybe T.Text))]   deriving (Eq, Ord, Show)  instance Monoid OpaqueTypes where
src/Futhark/IR/TypeCheck.hs view
@@ -111,7 +111,7 @@   show (InvalidPatError pat t desc) =     "Pat\n"       ++ prettyString pat-      ++ "\ncannot match value of type\n"+      ++ "\ncannot match expression of type\n"       ++ T.unpack (prettyTupleLines t)       ++ end     where@@ -549,7 +549,7 @@ checkOpaques (OpaqueTypes types) = descend [] types   where     descend _ [] = pure ()-    descend known ((name, t) : ts) = do+    descend known ((name, (t, _)) : ts) = do       check known t       descend (name : known) ts     check known (OpaqueRecord fs) =@@ -561,8 +561,6 @@     check known (OpaqueRecordArray _ v fs) = do       checkEntryPointType known (TypeOpaque v)       mapM_ (checkEntryPointType known . snd) fs-    check _ (OpaqueType _) =-      pure ()     checkEntryPointType known (TypeOpaque s) =       unless (s `elem` known) $         Left . Error [] . TypeError $@@ -1033,13 +1031,12 @@         </> indent 2 (pretty $ map fst rettype_annot)   consumeArgs paramtypes argflows checkExp (Loop merge form loopbody) = do-  let (mergepat, mergeexps) = unzip merge-  mergeargs <- mapM checkArg mergeexps+  let mergepat = map fst merge    checkLoopArgs    binding (scopeOfLoopForm form) $ do-    form_consumable <- checkForm mergeargs form+    checkForm form      let rettype = map paramDeclType mergepat         consumable =@@ -1047,7 +1044,6 @@           | param <- mergepat,             unique $ paramDeclType param           ]-            ++ form_consumable      context "Inside the loop body"       $ checkFun'@@ -1074,16 +1070,9 @@           map (`namesSubtract` bound_here)             <$> mapM (subExpAliasesM . resSubExp) (bodyResult loopbody)   where-    checkForm mergeargs (ForLoop loopvar it boundexp) = do-      iparam <- primFParam loopvar $ IntType it-      let mergepat = map fst merge-          funparams = iparam : mergepat-          paramts = map paramDeclType funparams--      boundarg <- checkArg boundexp-      checkFuncall Nothing paramts $ boundarg : mergeargs-      pure mempty-    checkForm mergeargs (WhileLoop cond) = do+    checkForm (ForLoop _ _ boundexp) =+      void $ checkArg boundexp+    checkForm (WhileLoop cond) = do       case find ((== cond) . paramName . fst) merge of         Just (condparam, _) ->           unless (paramType condparam == Prim Bool) $@@ -1096,11 +1085,6 @@         Nothing ->           -- Implies infinite loop, but that's OK.           pure ()-      let mergepat = map fst merge-          funparams = mergepat-          paramts = map paramDeclType funparams-      checkFuncall Nothing paramts mergeargs-      pure mempty      checkLoopArgs = do       let (params, args) = unzip merge@@ -1320,20 +1304,6 @@    unless (rettype' == ts) problem -validApply ::-  (ArrayShape shape) =>-  [TypeBase shape Uniqueness] ->-  [TypeBase shape NoUniqueness] ->-  Bool-validApply expected got =-  length got == length expected-    && and-      ( zipWith-          subtypeOf-          (map rankShaped got)-          (map (fromDecl . rankShaped) expected)-      )- type Arg = (Type, Names)  argType :: Arg -> Type@@ -1362,7 +1332,7 @@   TypeM rep () checkFuncall fname paramts args = do   let argts = map argType args-  unless (validApply paramts argts) $+  unless (map fromDecl paramts == argts) $     bad $       ParameterMismatch fname (map fromDecl paramts) $         map argType args
src/Futhark/Internalise/ApplyTypeAbbrs.hs view
@@ -43,8 +43,8 @@   applySubst (`M.lookup` types)  substEntry :: Types -> EntryPoint -> EntryPoint-substEntry types (EntryPoint params ret) =-  EntryPoint (map onEntryParam params) (onEntryType ret)+substEntry types (EntryPoint params ret doc) =+  EntryPoint (map onEntryParam params) (onEntryType ret) doc   where     onEntryParam (EntryParam v t) =       EntryParam v $ onEntryType t
src/Futhark/Internalise/Defunctorise.hs view
@@ -275,8 +275,8 @@ transformExp = transformNames  transformEntry :: EntryPoint -> TransformM EntryPoint-transformEntry (EntryPoint params ret) =-  EntryPoint <$> mapM onEntryParam params <*> onEntryType ret+transformEntry (EntryPoint params ret doc) =+  EntryPoint <$> mapM onEntryParam params <*> onEntryType ret <*> pure doc   where     onEntryParam (EntryParam v t) =       EntryParam v <$> onEntryType t
src/Futhark/Internalise/Entry.hs view
@@ -11,6 +11,7 @@ import Data.Bifunctor (first) import Data.List (find, intersperse) import Data.Map qualified as M+import Data.Text qualified as T import Futhark.IR qualified as I import Futhark.Internalise.TypesValues (internaliseSumTypeRep, internalisedTypeSize) import Futhark.Util (chunks)@@ -81,17 +82,17 @@ runGenOpaque :: GenOpaque a -> (a, I.OpaqueTypes) runGenOpaque = flip runState mempty -addType :: Name -> I.OpaqueType -> GenOpaque ()-addType name t = modify $ \(I.OpaqueTypes ts) ->-  case find ((== name) . fst) ts of-    Just (_, t')+addType :: Name -> Maybe T.Text -> I.OpaqueType -> GenOpaque ()+addType name doc t = modify $ \(I.OpaqueTypes ts) ->+  case lookup name ts of+    Just (t', _)       | t /= t' ->           error . unlines $             [ "Duplicate definition of entry point type " <> E.prettyString name,               show t,               show t'             ]-    _ -> I.OpaqueTypes ts <> I.OpaqueTypes [(name, t)]+    _ -> I.OpaqueTypes ts <> I.OpaqueTypes [(name, (t, doc))]  isRecord ::   VisibleTypes ->@@ -129,6 +130,13 @@   where     opaqueField e_t i_ts = snd <$> entryPointType types e_t i_ts +teArrayOf :: Int -> E.TypeExp d vn -> E.TypeExp d vn+teArrayOf rank t =+  foldl+    (\x _ -> E.TEArray (E.SizeExpAny mempty) x mempty)+    t+    [0 .. rank - 1]+ opaqueRecordArray ::   VisibleTypes ->   Int ->@@ -141,10 +149,11 @@   f' <- opaqueField t f_ts   ((f, f') :) <$> opaqueRecordArray types rank fs ts'   where-    opaqueField (E.EntryType e_t _) i_ts =-      snd <$> entryPointType types (E.EntryType e_t' Nothing) i_ts+    opaqueField (E.EntryType e_t ascribed) i_ts =+      snd <$> entryPointType types (E.EntryType e_t' ascribed') i_ts       where         e_t' = E.arrayOf (E.Shape (replicate rank $ E.anySize 0)) e_t+        ascribed' = teArrayOf rank <$> ascribed  isSum ::   VisibleTypes ->@@ -218,38 +227,37 @@       pure (u, I.TypeTransparent $ I.ValueType I.Signed r ts0)   | otherwise = do       case E.entryType t of-        E.Scalar (E.Record fs)-          | not $ null fs -> do-              let fs' = recordFields types fs $ E.entryAscribed t-              addType desc . I.OpaqueRecord =<< opaqueRecord types fs' ts+        E.Scalar (E.Record fs) -> do+          let fs' = recordFields types fs $ E.entryAscribed t+          addType desc doc . I.OpaqueRecord =<< opaqueRecord types fs' ts         E.Scalar (E.Sum cs) -> do           let (_, places) = internaliseSumTypeRep cs               cs' = sumConstrs types cs $ E.entryAscribed t               cs'' = zip (map fst cs') (zip (map snd cs') (map snd places))               ts' = if length cs == 1 then ts else drop 1 ts-          addType desc . I.OpaqueSum (map valueType ts)+          addType desc doc . I.OpaqueSum (map valueType ts)             =<< opaqueSum types cs'' ts'-        E.Array _ shape (E.Record fs)-          | not $ null fs -> do-              let fs' = recordFields types fs $ E.entryAscribed t-                  rank = E.shapeRank shape-                  ts' = map (strip rank) ts-                  record_t = E.Scalar (E.Record fs)-                  record_te = rowTypeExp rank =<< E.entryAscribed t-              ept <- snd <$> entryPointType types (E.EntryType record_t record_te) ts'-              addType desc . I.OpaqueRecordArray rank (entryPointTypeName ept)-                =<< opaqueRecordArray types rank fs' ts+        E.Array _ shape (E.Record fs) | not $ null fs -> do+          let rank = E.shapeRank shape+              fs' = recordFields types fs $ rowTypeExp rank =<< E.entryAscribed t+              ts' = map (strip rank) ts+              record_t = E.Scalar (E.Record fs)+              record_te = rowTypeExp rank =<< E.entryAscribed t+          ept <- snd <$> entryPointType types (E.EntryType record_t record_te) ts'+          addType desc doc . I.OpaqueRecordArray rank (entryPointTypeName ept)+            =<< opaqueRecordArray types rank fs' ts         E.Array _ shape et -> do           let ts' = map (strip (E.shapeRank shape)) ts               rank = E.shapeRank shape               elem_te = rowTypeExp rank =<< E.entryAscribed t           ept <- snd <$> entryPointType types (E.EntryType (E.Scalar et) elem_te) ts'-          addType desc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $+          addType desc doc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $             map valueType ts-        _ -> addType desc $ I.OpaqueType $ map valueType ts+        _ -> error $ "entryPointType: " <> E.prettyString (E.entryType t)        pure (u, I.TypeOpaque desc)   where+    doc = Nothing     u = foldl max Nonunique $ map I.uniqueness ts     desc =       maybe (nameFromText $ prettyTextOneLine t') typeExpOpaqueName $@@ -262,14 +270,15 @@ entryPoint ::   VisibleTypes ->   Name ->+  Maybe T.Text ->   [(E.EntryParam, [I.Param I.DeclType])] ->   ( E.EntryType,     [[I.TypeBase I.Rank I.Uniqueness]]   ) ->   (I.EntryPoint, I.OpaqueTypes)-entryPoint types name params (eret, crets) =+entryPoint types name doc params (eret, crets) =   runGenOpaque $-    (name,,)+    (name,,,doc)       <$> mapM onParam params       <*> ( uncurry I.EntryResult               <$> entryPointType types eret (concat crets)
src/Futhark/Internalise/Exps.hs view
@@ -20,7 +20,6 @@ import Futhark.Internalise.AccurateSizes import Futhark.Internalise.Bindings import Futhark.Internalise.Entry-import Futhark.Internalise.Lambdas import Futhark.Internalise.Monad as I import Futhark.Internalise.TypesValues import Futhark.Transform.Rename as I@@ -117,7 +116,7 @@     zeroExts ts = generaliseExtTypes ts ts  generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM ()-generateEntryPoint types (E.EntryPoint e_params e_rettype) vb = do+generateEntryPoint types (E.EntryPoint e_params e_rettype doc) vb = do   let (E.ValBind _ ofname _ _ (Info rettype) tparams params _ _ attrs _) = vb   bindingFParams tparams params $ \shapeparams params' -> do     let all_params = map pure shapeparams ++ concat params'@@ -127,6 +126,7 @@           entryPoint             types             (baseName ofname)+            doc             (zip e_params $ map (foldMap toList) params')             (e_rettype, map (map I.rankShaped) entry_rettype)         args = map (I.Var . I.paramName) $ foldMap (foldMap toList) params'@@ -1304,6 +1304,26 @@     negone = constant (-1 :: Int64)     one = constant (1 :: Int64) +internaliseFoldLambda ::+  E.Exp ->+  [I.Type] ->+  [I.Type] ->+  InternaliseM (I.Lambda SOACS)+internaliseFoldLambda lam acctypes arrtypes = do+  let rowtypes = map I.rowType arrtypes+  (params, body, rettype) <- internaliseLambda lam $ acctypes ++ rowtypes+  let rettype' =+        [ t `I.setArrayShape` I.arrayShape shape+        | (t, shape) <- zip rettype acctypes+        ]+  -- The result of the body must have the exact same shape as the+  -- initial accumulator.+  mkLambda params $+    ensureResultShape+      (ErrorMsg [ErrorString "shape of result does not match shape of initial value"])+      rettype'+      =<< bodyBind body+ internaliseScanOrReduce ::   Name ->   Name ->@@ -1322,7 +1342,7 @@       ne'   nests <- mapM I.subExpType nes'   arrts <- mapM lookupType arrs-  lam' <- internaliseFoldLambda internaliseLambda lam nests arrts+  lam' <- internaliseFoldLambda lam nests arrts   w <- arraysSize 0 <$> mapM lookupType arrs   letValExp' desc . I.Op =<< f w lam' nes' arrs @@ -1353,7 +1373,7 @@       n   ne_ts <- mapM I.subExpType ne_shp   his_ts <- mapM (fmap (I.stripArray (dim - 1)) . lookupType) hist'-  op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts+  op' <- internaliseFoldLambda op ne_ts his_ts    -- reshape return type of bucket function to have same size as neutral element   -- (modulo the index)@@ -1654,7 +1674,10 @@   where     stmsscope = scopeOf stms -internaliseLambda :: InternaliseLambda+internaliseLambda ::+  E.Exp ->+  [I.Type] ->+  InternaliseM ([I.LParam SOACS], I.Body SOACS, [I.Type]) internaliseLambda (E.Parens e _) rowtypes =   internaliseLambda e rowtypes internaliseLambda (E.Lambda params body _ (Info (RetType _ rettype)) _) rowtypes =@@ -1810,16 +1833,6 @@       lam' <- internaliseLambdaCoerce lam $ map rowType arr_ts       let w = arraysSize 0 arr_ts       letTupExp' desc . I.Op . I.Screma w arr' =<< I.mapSOAC lam'-    handleSOACs [k, lam, arr] "partition" = do-      k' <- fromIntegral <$> fromInt32 k-      Just $ \_desc -> do-        arrs <- internaliseExpToVars "partition_input" arr-        lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs-        uncurry (++) <$> partitionWithSOACS (fromIntegral k') lam' arrs-      where-        fromInt32 (Literal (SignedValue (Int32Value k')) _) = Just k'-        fromInt32 (IntLit k' (Info (E.Scalar (E.Prim (E.Signed Int32)))) _) = Just $ fromInteger k'-        fromInt32 _ = Nothing     handleSOACs [lam, ne, arr] "reduce" = Just $ \desc ->       internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr)       where@@ -1857,14 +1870,24 @@     handleAccs _ _ = Nothing      handleAD [f, x, v] fname-      | fname `elem` ["jvp2", "vjp2"] = Just $ \desc -> do+      | fname `elem` ["jvp2", "vjp2", "jmp2", "mjp2"] = Just $ \desc -> do           x' <- internaliseExp "ad_x" x           v' <- internaliseExp "ad_v" v+          x_t <- subExpType $ head x'+          v_t <- subExpType $ head v'           lam <- internaliseLambdaCoerce f =<< mapM subExpType x'           fmap (map I.Var) . letTupExp desc . Op $             case fname of-              "jvp2" -> JVP x' v' lam-              _ -> VJP x' v' lam+              "jvp2" -> JVP mempty x' v' lam+              "vjp2" -> VJP mempty x' v' lam+              "jmp2" ->+                JVP (vecShape x_t v_t) x' v' lam+              "mjp2" ->+                VJP (vecShape (head (lambdaReturnType lam)) v_t) x' v' lam+              _ -> error "handleAD: not supposed to happen."+      where+        vecShape t1 t2 =+          I.Shape $ take (I.arrayRank t2 - I.arrayRank t1) (I.arrayDims t2)     handleAD [f, f_adj, x] "with_vjp" = Just $ \desc -> do       x' <- internaliseExp "ad_x" x       lam <- internaliseLambdaCoerce f =<< mapM subExpType x'@@ -2168,114 +2191,6 @@ askSafety = do   check <- asks envDoBoundsChecks   pure $ if check then I.Safe else I.Unsafe---- Implement partitioning using maps, scans and writes.-partitionWithSOACS :: Int -> I.Lambda SOACS -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp])-partitionWithSOACS k lam arrs = do-  arr_ts <- mapM lookupType arrs-  let w = arraysSize 0 arr_ts-  classes_and_increments <--    letTupExp "increments"-      . I.Op-      . I.Screma w arrs-      =<< mapSOAC lam-  (classes, increments) <- case classes_and_increments of-    classes : increments -> pure (classes, take k increments)-    _ -> error "partitionWithSOACS"--  add_lam_x_params <--    replicateM k $ newParam "x" (I.Prim int64)-  add_lam_y_params <--    replicateM k $ newParam "y" (I.Prim int64)-  add_lam_body <- runBodyBuilder $-    localScope (scopeOfLParams $ add_lam_x_params ++ add_lam_y_params) $-      fmap subExpsRes $-        forM (zip add_lam_x_params add_lam_y_params) $ \(x, y) ->-          letSubExp "z" $-            I.BasicOp $-              I.BinOp-                (I.Add Int64 I.OverflowUndef)-                (I.Var $ I.paramName x)-                (I.Var $ I.paramName y)-  let add_lam =-        I.Lambda-          { I.lambdaBody = add_lam_body,-            I.lambdaParams = add_lam_x_params ++ add_lam_y_params,-            I.lambdaReturnType = replicate k $ I.Prim int64-          }-      nes = replicate (length increments) $ intConst Int64 0--  scan <- I.scanSOAC [I.Scan add_lam nes]-  all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w increments scan--  -- We have the offsets for each of the partitions, but we also need-  -- the total sizes, which are the last elements in the offests.  We-  -- just have to be careful in case the array is empty.-  last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int64 OverflowUndef) w $ constant (1 :: Int64)-  let nonempty_body = runBodyBuilder $-        fmap subExpsRes $-          forM all_offsets $ \offset_array ->-            letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array $ Slice [I.DimFix last_index]-      empty_body = resultBodyM $ replicate k $ constant (0 :: Int64)-  is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int64) w $ constant (0 :: Int64)-  sizes <--    letTupExp "partition_size" =<< eIf (eSubExp is_empty) empty_body nonempty_body--  -- The total size of all partitions must necessarily be equal to the-  -- size of the input array.--  -- Create scratch arrays for the result.-  blanks <- forM arr_ts $ \arr_t ->-    letExp "partition_dest" $-      I.BasicOp $-        Scratch (I.elemType arr_t) (w : drop 1 (I.arrayDims arr_t))--  -- Now write into the result.-  results <--    doScatter "partition_res" 1 blanks (classes : all_offsets ++ arrs) $ \params -> do-      let ([c_param], offset_params, value_params) =-            splitAt3 1 (length all_offsets) params-      offset <--        mkOffsetLambdaBody-          (map I.Var sizes)-          (I.Var $ I.paramName c_param)-          0-          offset_params-      pure $ offset : map (I.Var . I.paramName) value_params-  sizes' <--    letSubExp "partition_sizes" $-      I.BasicOp $-        I.ArrayLit (map I.Var sizes) $-          I.Prim int64-  pure (map I.Var results, [sizes'])-  where-    mkOffsetLambdaBody ::-      [SubExp] ->-      SubExp ->-      Int ->-      [I.LParam SOACS] ->-      InternaliseM SubExp-    mkOffsetLambdaBody _ _ _ [] =-      pure $ constant (-1 :: Int64)-    mkOffsetLambdaBody sizes c i (p : ps) = do-      is_this_one <--        letSubExp "is_this_one" $-          I.BasicOp $-            I.CmpOp (CmpEq int64) c $-              intConst Int64 $-                toInteger i-      next_one <- mkOffsetLambdaBody sizes c (i + 1) ps-      this_one <--        letSubExp "this_offset"-          =<< foldBinOp-            (Add Int64 OverflowUndef)-            (constant (-1 :: Int64))-            (I.Var (I.paramName p) : take i sizes)-      letSubExp "total_res"-        =<< eIf-          (eSubExp is_this_one)-          (resultBodyM [this_one])-          (resultBodyM [next_one])  sizeExpForError :: E.Size -> InternaliseM [ErrorMsgPart SubExp] sizeExpForError e
− src/Futhark/Internalise/Lambdas.hs 
@@ -1,82 +0,0 @@-module Futhark.Internalise.Lambdas-  ( InternaliseLambda,-    internaliseFoldLambda,-    internalisePartitionLambda,-  )-where--import Data.Maybe (listToMaybe)-import Futhark.IR.SOACS as I-import Futhark.Internalise.AccurateSizes-import Futhark.Internalise.Monad-import Language.Futhark as E---- | A function for internalising lambdas.-type InternaliseLambda =-  E.Exp -> [I.Type] -> InternaliseM ([I.LParam SOACS], I.Body SOACS, [I.Type])--internaliseFoldLambda ::-  InternaliseLambda ->-  E.Exp ->-  [I.Type] ->-  [I.Type] ->-  InternaliseM (I.Lambda SOACS)-internaliseFoldLambda internaliseLambda lam acctypes arrtypes = do-  let rowtypes = map I.rowType arrtypes-  (params, body, rettype) <- internaliseLambda lam $ acctypes ++ rowtypes-  let rettype' =-        [ t `I.setArrayShape` I.arrayShape shape-        | (t, shape) <- zip rettype acctypes-        ]-  -- The result of the body must have the exact same shape as the-  -- initial accumulator.-  mkLambda params $-    ensureResultShape-      (ErrorMsg [ErrorString "shape of result does not match shape of initial value"])-      rettype'-      =<< bodyBind body---- Given @k@ lambdas, this will return a lambda that returns an--- (k+2)-element tuple of integers.  The first element is the--- equivalence class ID in the range [0,k].  The remaining are all zero--- except for possibly one element.-internalisePartitionLambda ::-  InternaliseLambda ->-  Int ->-  E.Exp ->-  [I.SubExp] ->-  InternaliseM (I.Lambda SOACS)-internalisePartitionLambda internaliseLambda k lam args = do-  argtypes <- mapM I.subExpType args-  let rowtypes = map I.rowType argtypes-  (params, body, _) <- internaliseLambda lam rowtypes-  body' <--    localScope (scopeOfLParams params) $-      lambdaWithIncrement body-  pure $ I.Lambda params rettype body'-  where-    rettype = replicate (k + 2) $ I.Prim int64-    result i =-      map constant $-        fromIntegral i-          : (replicate i 0 ++ [1 :: Int64] ++ replicate (k - i) 0)--    mkResult _ i | i >= k = pure $ result i-    mkResult eq_class i = do-      is_i <--        letSubExp "is_i" $-          BasicOp $-            CmpOp (CmpEq int64) eq_class $-              intConst Int64 $-                toInteger i-      letTupExp' "part_res"-        =<< eIf-          (eSubExp is_i)-          (pure $ resultBody $ result i)-          (resultBody <$> mkResult eq_class (i + 1))--    lambdaWithIncrement :: I.Body SOACS -> InternaliseM (I.Body SOACS)-    lambdaWithIncrement lam_body = runBodyBuilder $ do-      eq_class <--        maybe (intConst Int64 0) resSubExp . listToMaybe <$> bodyBind lam_body-      subExpsRes <$> mkResult eq_class 0
src/Futhark/Internalise/Monomorphise.hs view
@@ -1225,9 +1225,7 @@           M.insert (valBindName valbind) (removeEntryPoint valbind') $             envPolyBindings env,         envGlobalScope = global <> envGlobalScope env,-        envScope =-          S.insert (valBindName valbind) global-            <> envScope env+        envScope = S.insert (valBindName valbind) global <> envScope env       }  transformValBinds :: [ValBind] -> MonoM ()
src/Futhark/Optimise/Fusion.hs view
@@ -596,6 +596,14 @@     relevant (_, e) = isDep e     fuses_with = map fst $ filter relevant $ G.lpre g node_to_fuse_id +doSoacThroughTransFusion :: DepGraphAug FusionM+doSoacThroughTransFusion dg =+  applyAugs+    [ SF.trySoacThroughTransIntoWithAcc doFusionInLambda fusedSomething wacc_id+    | (wacc_id, StmNode (Let _ _ (WithAcc {}))) <- G.labNodes (dgGraph dg)+    ]+    dg+ doVerticalFusion :: DepGraphAug FusionM doVerticalFusion dg = applyAugs (map tryFuseNodeInGraph $ reverse $ filter relevant $ G.labNodes (dgGraph dg)) dg   where@@ -651,7 +659,8 @@ doAllFusion :: DepGraphAug FusionM doAllFusion =   keepTrying . applyAugs $-    [ doVerticalFusion,+    [ doSoacThroughTransFusion,+      doVerticalFusion,       doHorizontalFusion,       doInnerFusion,       removeUnusedOutputs@@ -667,12 +676,12 @@     cases' <- mapM (traverse $ renameBody <=< (`doFusionWithDelayed` to_fuse)) cases     defbody' <- doFusionWithDelayed defbody to_fuse     pure (incoming, node, MatchNode (Let pat aux (Match cond cases' defbody' dec)) [], outgoing)-  StmNode (Let pat aux (Op (Futhark.VJP args vec lam))) -> doFuseScans $ do+  StmNode (Let pat aux (Op (Futhark.VJP shape args vec lam))) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam-    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP args vec lam'))), outgoing)-  StmNode (Let pat aux (Op (Futhark.JVP args vec lam))) -> doFuseScans $ do+    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP shape args vec lam'))), outgoing)+  StmNode (Let pat aux (Op (Futhark.JVP shape args vec lam))) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam-    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP args vec lam'))), outgoing)+    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP shape args vec lam'))), outgoing)   StmNode (Let pat aux (WithAcc inputs lam)) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam     pure (incoming, node, StmNode (Let pat aux (WithAcc inputs lam')), outgoing)
src/Futhark/Optimise/Fusion/RulesWithAccs.hs view
@@ -4,16 +4,18 @@ --     that involves WithAcc constructs. --   Currently, we support two non-trivial --   transformations:---     I. map-flatten-scatter: a map nest produces---        multi-dimensional index and values arrays---        that are then flattened and used in a---        scatter consumer. Such pattern can be fused---        by re-writing the scatter by means of a WithAcc---        containing a map-nest, thus eliminating the flatten---        operations. The obtained WithAcc can then be fused---        with the producer map nest, e.g., benefiting intra-group---        kernels. The eloquent target for this rule is---        an efficient implementation of radix-sort.+--     I. SOAC-through-Trans-into-WithAcc fusion: a SoacNode is fused+--        into a WithAcc atomically when all dependency paths between+--        them go through TransNodes (reshapes/rearranges) that have no+--        other consumers. The canonical example is a map nest producing+--        multi-dimensional index and value arrays that are flattened+--        and consumed by a scatter. This must be done atomically to+--        avoid an infinite loop: absorbing only the TransNode causes+--        simplifyLambda to hoist the cheap reshape back out, recreating+--        the same TransNode and triggering the rule again. The strategy+--        is to prepend the SoacNode and all TransNode statements into+--        the WithAcc lambda body and run doFusionInLambda to fuse+--        further where possible. -- --    II. WithAcc-WithAcc fusion: two withaccs can be --        fused as long as the common accumulators use@@ -27,16 +29,22 @@ --        they can be transformed by various optimizations passes. module Futhark.Optimise.Fusion.RulesWithAccs   ( tryFuseWithAccs,+    trySoacThroughTransIntoWithAcc,   ) where  import Control.Monad+import Data.Graph.Inductive.Graph qualified as G import Data.List qualified as L import Data.Map.Strict qualified as M+import Data.Maybe (maybeToList)+import Futhark.Analysis.HORep.SOAC qualified as H import Futhark.Construct import Futhark.IR.SOACS hiding (SOAC (..))+import Futhark.Optimise.Fusion.GraphRep import Futhark.Transform.Rename import Futhark.Transform.Substitute+import Futhark.Util (nubOrd)  --------------------------------------------------- --- II. WithAcc-WithAcc Fusion@@ -211,6 +219,97 @@ -- tryFuseWithAccs _ _ _ =   Nothing++---------------------------------------------------+--- I. SOAC-through-Trans-into-WithAcc Fusion+---------------------------------------------------++-- | See the module-level description of transformation I.+-- The @doFusionInLambda@ and @fusedSomething@ arguments are callbacks+-- from Fusion.hs to avoid a circular import.+trySoacThroughTransIntoWithAcc ::+  (HasScope SOACS m, MonadFreshNames m) =>+  (Lambda SOACS -> m (Lambda SOACS, Bool)) ->+  (NodeT -> m (Maybe NodeT)) ->+  G.Node ->+  DepGraphAug m+trySoacThroughTransIntoWithAcc doFusionInLambda fusedSomething wacc_id dg@DepGraph {dgGraph = g}+  | not (G.gelem wacc_id g) = pure dg+  | Just (StmNode (Let pat2 aux2 (WithAcc w_inps lam0))) <- G.lab g wacc_id = do+      -- Edges go FROM consumers TO producers:+      --   G.lpre g n = consumers of n; G.lsuc g n = producers n depends on.+      -- realConsumers n: consumers of n, excluding Alias edges and self-loops.+      let wacc_cons_nms = namesFromList $ concatMap (\(_, nms, _) -> nms) w_inps+          realConsumers n =+            nubOrd $+              map fst $+                filter (\(m, e) -> m /= n && case e of Alias {} -> False; _ -> True) $+                  G.lpre g n+          -- TransNodes that directly feed wacc and are exclusively consumed by wacc.+          trans_preds = do+            (tn_id, _) <- G.lsuc g wacc_id+            TransNode out tr inp <- maybeToList $ G.lab g tn_id+            guard $ realConsumers tn_id == [wacc_id]+            pure (tn_id, out, tr, inp)+          trans_ids = map (\(a, _, _, _) -> a) trans_preds+          trans_out_nms = namesFromList $ map (\(_, out, _, _) -> out) trans_preds+          -- The unique Screma SoacNode that feeds all TransNodes and has no+          -- other consumers besides those TransNodes.+          soac_preds = do+            (tn_id, _, _, _) <- trans_preds+            (sn_id, _) <- G.lsuc g tn_id+            guard $ sn_id `notElem` trans_ids+            guard $ all (`elem` trans_ids) (realConsumers sn_id)+            pure sn_id+          prod_ids = nubOrd soac_preds+      case (trans_preds, prod_ids) of+        (_ : _, [prod_id])+          | Just (SoacNode ots1 pat1 soac@(H.Screma {}) aux1) <- G.lab g prod_id,+            ots1 == mempty,+            all ((`notNameIn` wacc_cons_nms) . H.inputArray) (H.inputs soac),+            not $ namesIntersect trans_out_nms wacc_cons_nms ->+              attempt trans_ids trans_preds prod_id pat1 aux1 soac pat2 aux2 w_inps lam0+        _ -> pure dg+  | otherwise = pure dg+  where+    attempt trans_ids trans_preds prod_id pat1 aux1 soac pat2 aux2 w_inps lam0 = do+      let trans_info = map (\(_, out, tr, inp) -> (out, tr, inp)) trans_preds+      lam' <- renameLambda <=< runLambdaBuilder (lambdaParams lam0) $ do+        soac' <- H.toExp soac+        addStm $ Let pat1 aux1 soac'+        forM_ trans_info $ \(out, tr, inp) -> do+          (tr_aux, tr_exp) <- H.transformToExp tr inp+          auxing tr_aux $ letBindNames [out] tr_exp+        bodyBind $ lambdaBody lam0+      -- Run inner fusion. We always proceed with onSuccess because embedding+      -- the SoacNode + TransNodes into the WithAcc is itself a valid fusion+      -- step and avoids potential infinite loops from simplifyLambda hoisting+      -- a reshape back out.+      lam'' <- fst <$> doFusionInLambda lam'+      onSuccess trans_ids prod_id pat2 aux2 w_inps lam''++    onSuccess trans_ids prod_id pat2 aux2 w_inps lam'' = do+      void $ fusedSomething (StmNode $ Let pat2 aux2 $ WithAcc w_inps lam'')+      -- Rebuild the graph: remove absorbed nodes and rewire wacc's edges.+      -- G.context returns (in_adj, n, label, out_adj) where both adjacency+      -- lists are [(EdgeT, Node)]. G.lsuc/lpre are (Node, EdgeT).+      let to_remove = prod_id : trans_ids+          new_wacc = StmNode $ Let pat2 aux2 $ WithAcc w_inps lam''+          g' = foldr G.delNode g to_remove+          (wacc_preds, _, _, wacc_succs) = G.context g wacc_id+          -- Inherit producers of the absorbed nodes that are still in g'.+          removed_succs =+            nubOrd $+              concatMap (filter ((`G.gelem` g') . fst) . G.lsuc g) to_remove+          new_succs =+            nubOrd $+              removed_succs+                ++ map (\(e, n) -> (n, e)) (filter ((`G.gelem` g') . snd) wacc_succs)+          new_preds = filter ((`G.gelem` g') . snd) wacc_preds+          g'' = G.insNode (wacc_id, new_wacc) $ G.delNode wacc_id g'+          g''' = foldr (\(e, n) gr -> G.insEdge (n, wacc_id, e) gr) g'' new_preds+          g'''' = foldr (\(n, e) gr -> G.insEdge (wacc_id, n, e) gr) g''' new_succs+      pure dg {dgGraph = g''''}  ------------------------------- --- simple helper functions ---
src/Futhark/Optimise/Fusion/Screma.hs view
@@ -149,6 +149,9 @@       is_fusible =         fuseIsVarish inp_c out_p           && not (forbidden_c `namesIntersect` forbidden_p)+          && all+            (`notElem` mapMaybe SOAC.isVarishInput inp_c)+            (take num_red_p out_p)   unless is_fusible (fail "Scremas are not fusible.")   where     pre_pars_c = oneName . paramName <$> lambdaParams pre_c@@ -165,6 +168,7 @@     post_scan_pars_p = take num_scan_p $ paramName <$> lambdaParams post_p     num_scan_c = scanResults $ scremaScans form_c     num_red_c = redResults $ scremaReduces form_c+    num_red_p = redResults $ scremaReduces form_p     num_scan_p = scanResults $ scremaScans form_p  -- | Given two scremas that are fusible, fuse them into a super
src/Futhark/Optimise/Fusion/TryFusion.hs view
@@ -18,6 +18,7 @@ import Control.Monad import Control.Monad.Reader import Control.Monad.State+import Data.Either (partitionEithers) import Data.List (find) import Data.Map.Strict qualified as M import Data.Maybe@@ -457,7 +458,7 @@   TryFusion (SOAC, SOAC.ArrayTransforms)  optimizations :: [Optimization]-optimizations = [iswim]+optimizations = [iswim, unflattenAccOnlyMap]  iswim ::   Maybe [VName] ->@@ -506,6 +507,87 @@ iswim _ _ _ =   fail "ISWIM does not apply." +-- | When a pure-map Screma returns exclusively accumulator results and some+-- non-accumulator inputs carry a 2D-to-1D flattening Reshape transform, we+-- can "unflatten" the map:+--+--   Screma(n*m, {flat_a1:[n*m]t1, ..., acc_p:acc(...)}, lam)+--   where lam : (t1, ..., acc) → acc+--+-- becomes+--+--   Screma(n, {a1:[n][m]t1, ..., acc_p:acc(...)}, outer_lam)+--   where outer_lam = \(row_a1:[m]t1, ..., acc_p:acc(...)) →+--     Screma(m, {row_a1, ..., acc_p}, lam)+--+-- This exposes the 2D inputs directly, enabling the standard fusion rules to+-- fuse the resulting Screma(n,...) with an upstream Screma(n,...) producer.+unflattenAccOnlyMap ::+  Maybe [VName] ->+  SOAC ->+  SOAC.ArrayTransforms ->+  TryFusion (SOAC, SOAC.ArrayTransforms)+unflattenAccOnlyMap (Just outVars) (SOAC.Screma _nm inps form) ots = do+  lam <- liftMaybe $ isMapSOAC form+  -- All results must be accumulator types.+  guard $ all isAcc $ lambdaReturnType lam+  -- Only apply when the producer outputs non-scalar rows (rank > 1), meaning+  -- pullReshape cannot handle this case (it requires scalar-leaf map nests).+  -- When the producer outputs scalars, the simpler prepend approach works fine.+  outVarTypes <- mapM lookupType outVars+  guard $ any ((> 1) . arrayRank) outVarTypes+  -- Partition inputs paired with their lambda params: those with a 2D→1D+  -- flattening Reshape vs. those that pass through unchanged (acc params).+  -- A flattening reshape: base type is 2D, first transform collapses it to 1D.+  let classifyInp (inp@(SOAC.Input ts _v base_t), p)+        | SOAC.Reshape _aux ns SOAC.:< ts' <- SOAC.viewf ts,+          arrayRank base_t == 2,+          shapeRank (newShape ns) == 1 =+            Left (SOAC.Input ts' _v base_t, p)+        | otherwise =+            Right (inp, p)+      (flat_pairs, pass_pairs) =+        partitionEithers $ zipWith (curry classifyInp) inps (lambdaParams lam)+  -- Need at least one flattened input.+  guard $ not (null flat_pairs)+  -- The non-flattened inputs must be accumulators, because we are changing the+  -- width of the SOAC.+  guard $ all (isAcc . SOAC.inputType . fst) pass_pairs+  -- All flattened inputs must agree on the outer dim n and inner dim m.+  let dims2d base_t = (arraySize 0 base_t, arraySize 1 base_t)+      getBaseTy (SOAC.Input _ _ base_t, _) = base_t+      (n, m) = dims2d (getBaseTy (head flat_pairs))+  guard $ all ((== (n, m)) . dims2d . getBaseTy) flat_pairs+  -- The lambda params for the flat inputs get their type changed from [n*m]t+  -- to [m]t (a single row).  Pass-through params are unchanged.+  let mkRowParam (_, p) = p {paramDec = rowType (paramDec p)}+      flat_row_params = map mkRowParam flat_pairs+      pass_params = map snd pass_pairs+      inner_lam_params = flat_row_params ++ pass_params+  inner_lam <- renameLambda $ lam {lambdaParams = inner_lam_params}+  inner_form <- mapSOAC inner_lam+  -- Inner Screma over m: plain-variable inputs for the row params, then+  -- plain-variable inputs for the pass-through (acc) params.+  let inner_inps =+        map (SOAC.identInput . paramToIdent) flat_row_params+          ++ map (SOAC.identInput . paramToIdent) pass_params+      inner_soac = SOAC.Screma m inner_inps inner_form+  -- Outer lambda: same param names but outer params have type [m]t (rows).+  let outer_lam_params = flat_row_params ++ pass_params+  outer_lam <- runLambdaBuilder outer_lam_params $ do+    inner_exp <- SOAC.toExp inner_soac+    res <- letTupExp "inner_acc" inner_exp+    pure $ map (subExpRes . Var) res+  outer_form <- mapSOAC outer_lam+  -- Outer Screma over n: 2D inputs (flatten reshape stripped) then pass-through.+  let outer_inps = map fst flat_pairs ++ map fst pass_pairs+  pure (SOAC.Screma n outer_inps outer_form, ots)+unflattenAccOnlyMap _ _ _ =+  fail "unflattenAccOnlyMap does not apply."++paramToIdent :: Param Type -> Ident+paramToIdent p = Ident (paramName p) (paramType p)+ removeParamOuterDim :: LParam SOACS -> LParam SOACS removeParamOuterDim param =   let t = rowType $ paramType param@@ -538,11 +620,21 @@     Just (Just mot, inps') -> first (mot SOAC.<|) $ commonTransforms' $ reverse inps'     _ -> (SOAC.noTransforms, map snd inps)   where+    -- Two reshapes with the same shape are compatible even if their certs differ;+    -- merge the certs to produce a single common reshape transform.+    compatibleTransforms (SOAC.Reshape aux1 shape1) (SOAC.Reshape aux2 shape2)+      | shape1 == shape2 =+          Just $ SOAC.Reshape (aux1 <> aux2) shape1+    compatibleTransforms ot1 ot2+      | ot1 == ot2 = Just ot1+    compatibleTransforms _ _ = Nothing+     inspect (mot, prev) (True, inp) =       case (mot, inputToOutput inp) of         (Nothing, Just (ot, inp')) -> Just (Just ot, (True, inp') : prev)         (Just ot1, Just (ot2, inp'))-          | ot1 == ot2 -> Just (Just ot2, (True, inp') : prev)+          | Just combined <- compatibleTransforms ot1 ot2 ->+              Just (Just combined, (True, inp') : prev)         _ -> Nothing     inspect (mot, prev) inp = Just (mot, inp : prev)
src/Futhark/Optimise/Simplify/Engine.hs view
@@ -42,6 +42,8 @@     asksEngineEnv,     askVtable,     localVtable,+    Protect,+    protectIf,      -- * Building blocks     SimplifiableRep,@@ -284,6 +286,8 @@     zeroIfContext (Var v) | v `elem` ctx_names = intConst Int64 0     zeroIfContext se = se +-- | Protect a hoisted statement by enclosing it in a branch (or doing something+-- smarter for certain statements). protectIf ::   (MonadBuilder m) =>   Protect m ->
src/Futhark/Optimise/Simplify/Rules/BasicOp.hs view
@@ -304,6 +304,15 @@                 Rearrange v2 (map (subtract num_dims) rest_perm)             letBind pat $ BasicOp $ Replicate dims v +-- Rearranging a replicate of primitives is the same as just reshaping it.+ruleBasicOp vtable pat aux (Rearrange v1 perm)+  | Just (BasicOp (Replicate _ se), v1_cs) <- ST.lookupExp v1 vtable,+    Just old_shape <- arrayShape <$> ST.lookupType v1 vtable,+    Just (Prim _) <- ST.lookupSubExpType se vtable =+      Simplify . certifying v1_cs . auxing aux $ do+        let new_shape = Shape $ rearrangeShape perm $ shapeDims old_shape+        letBind pat $ BasicOp $ Reshape v1 $ reshapeAll old_shape new_shape+ -- Simplify away 0<=i when 'i' is from a loop of form 'for i < n'. ruleBasicOp vtable pat aux (CmpOp CmpSle {} x y)   | Constant (IntValue (Int64Value 0)) <- x,
src/Futhark/Optimise/TileLoops.hs view
@@ -122,6 +122,7 @@       -- 2D tiling of redomap.       | (gtids, kdims) <- unzip $ unSegSpace initial_space,         Just (w, arrs, form) <- tileable stm_to_tile,+        all (\gtid -> subExpInvariantTo gtid variance w) gtids,         Just inputs <-           mapM (invariantToOneOfTwoInnerDims branch_variant variance gtids) arrs,         not $ null $ tiledInputs inputs,@@ -145,6 +146,7 @@       -- 1D tiling of redomap.       | (gtid, kdim) : top_space_rev <- reverse $ unSegSpace initial_space,         Just (w, arrs, form) <- tileable stm_to_tile,+        subExpInvariantTo gtid variance w,         inputs <- map (is1DTileable gtid variance) arrs,         not $ null $ tiledInputs inputs,         gtid `notNameIn` branch_variant,@@ -745,9 +747,21 @@   let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims   pure $ TileReturns mempty tile_dims arr' -is1DTileable :: VName -> M.Map VName Names -> VName -> InputArray+lookupVariance :: VName -> VarianceTable -> Names+lookupVariance v variance = oneName v <> M.findWithDefault mempty v variance++-- | Is the second name invariant to the first?+varInvariantTo :: VName -> VarianceTable -> VName -> Bool+varInvariantTo gtid variance v =+  not $ nameIn gtid $ lookupVariance v variance++subExpInvariantTo :: VName -> VarianceTable -> SubExp -> Bool+subExpInvariantTo gtid variance (Var v) = varInvariantTo gtid variance v+subExpInvariantTo _ _ Constant {} = True++is1DTileable :: VName -> VarianceTable -> VName -> InputArray is1DTileable gtid variance arr-  | not $ nameIn gtid $ M.findWithDefault mempty arr variance =+  | varInvariantTo gtid variance arr =       InputTile [0] arr   | otherwise =       InputDontTile arr@@ -957,7 +971,7 @@  invariantToOneOfTwoInnerDims ::   Names ->-  M.Map VName Names ->+  VarianceTable ->   [VName] ->   VName ->   Maybe InputArray
src/Futhark/Pass/AD.hs view
@@ -36,20 +36,22 @@     certifying cs $ letBindNames [v] $ BasicOp $ SubExp se  onStm :: Bool -> Mode -> Scope SOACS -> Stm SOACS -> PassM (Stms SOACS)-onStm _ mode scope (Let pat aux (Op (VJP args vec lam))) = do+onStm _ mode scope (Let pat aux (Op (VJP shape args vec lam))) = do   lam' <- onLambda True mode scope lam   if mode == All || lam == lam'     then do-      lam'' <- (`runReaderT` scope) . simplifyLambda =<< revVJP scope lam'+      lam'' <-+        (`runReaderT` scope) . simplifyLambda+          =<< revVJP scope shape (stmAuxAttrs aux) lam'       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope-    else pure $ oneStm $ Let pat aux $ Op $ VJP args vec lam'-onStm _ mode scope (Let pat aux (Op (JVP args vec lam))) = do+    else pure $ oneStm $ Let pat aux $ Op $ VJP shape args vec lam'+onStm _ mode scope (Let pat aux (Op (JVP shape args vec lam))) = do   lam' <- onLambda True mode scope lam   if mode == All || lam == lam'     then do-      lam'' <- fwdJVP scope lam'+      lam'' <- fwdJVP scope shape (stmAuxAttrs aux) lam'       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope-    else pure $ oneStm $ Let pat aux $ Op $ JVP args vec lam'+    else pure $ oneStm $ Let pat aux $ Op $ JVP shape args vec lam' -- -- This corresponds to a WithVJP that is not inside of a differential operator. -- FIXME: this assumption will go bad when we don't inline so much.
src/Futhark/Test/Property.hs view
@@ -478,15 +478,15 @@           then loop i           else do             let failmsg =-                  "PBT FAIL: "+                  "Property "                     <> propName-                    <> " size="+                    <> " (size="                     <> showText size-                    <> " seed="+                    <> ", seed="                     <> showText (configSeed config)-                    <> " after "+                    <> ") failed after "                     <> showText i-                    <> " tests\n"+                    <> " tests:\n"              shrinkRes <- case psShrink s of               Nothing ->
src/Futhark/Tools.hs view
@@ -13,6 +13,8 @@     partitionChunkedFoldParameters,     withAcc,     doScatter,+    addBinOp,+    addLambda,      -- * Primitive expressions     module Futhark.Analysis.PrimExp.Convert,@@ -20,11 +22,58 @@ where  import Control.Monad+import Data.List qualified as L+import Data.Maybe import Futhark.Analysis.PrimExp.Convert import Futhark.Construct import Futhark.IR import Futhark.IR.SOACS.SOAC+import Futhark.Util (mapAccumLM) +splitScanOrRedomap ::+  (MonadFreshNames m) =>+  [PatElem Type] ->+  SubExp ->+  Lambda rep ->+  [[SubExp]] ->+  m (Pat Type, Pat Type, [VName], Lambda rep)+splitScanOrRedomap pes w map_lam nes = do+  let (nonmap_pes, map_pes) =+        splitAt (length $ concat nes) pes+      (nonmap_ts, map_ts) =+        splitAt (length (concat nes)) $ lambdaReturnType map_lam+      (nonmap_res, map_res) =+        splitAt (length (concat nes)) $ bodyResult $ lambdaBody map_lam++  -- Put some care into not having duplicate results from the map function.+  (red_arrs, acc_info) <-+    unzip . snd+      <$> mapAccumLM+        accMapPatElem+        (zip map_res map_pes)+        (zip3 nonmap_pes nonmap_ts nonmap_res)++  let (nonmap_tmppes, nonmap_ts', nonmap_res') =+        L.unzip3 $ catMaybes acc_info+      map_lam' =+        map_lam+          { lambdaBody = (lambdaBody map_lam) {bodyResult = nonmap_res' <> map_res},+            lambdaReturnType = nonmap_ts' <> map_ts+          }+      map_pat = nonmap_tmppes <> map_pes++  pure (Pat map_pat, Pat nonmap_pes, red_arrs, map_lam')+  where+    accMapPatElem res_to_pe (pe, nonmap_t, res) =+      case res `L.lookup` res_to_pe of+        Just pe' -> pure (res_to_pe, (patElemName pe', Nothing))+        Nothing -> do+          pe' <-+            PatElem+              <$> newVName (baseName (patElemName pe) <> "_map_acc")+              <*> pure (nonmap_t `arrayOfRow` w)+          pure ((res, pe') : res_to_pe, (patElemName pe', Just (pe', nonmap_t, res)))+ -- | Turns a binding of a @redomap@ into two seperate bindings, a -- @map@ binding and a @reduce@ binding (returned in that order). --@@ -32,6 +81,7 @@ -- pattern with new 'Ident's for the result of the @map@. redomapToMapAndReduce ::   ( MonadFreshNames m,+    LetDec rep ~ Type,     Buildable rep,     ExpDec rep ~ (),     Op rep ~ SOAC rep@@ -44,9 +94,10 @@   ) ->   m (Stm rep, Stm rep) redomapToMapAndReduce (Pat pes) (w, reds, map_lam, arrs) = do-  (map_pat, red_pat, red_arrs) <-+  (map_pat, red_pat, red_arrs, map_lam') <-     splitScanOrRedomap pes w map_lam $ map redNeutral reds-  map_stm <- mkLet map_pat . Op . Screma w arrs <$> mapSOAC map_lam++  map_stm <- Let map_pat (defAux ()) . Op . Screma w arrs <$> mapSOAC map_lam'   red_stm <-     Let red_pat (defAux ()) . Op       <$> (Screma w red_arrs <$> reduceSOAC reds)@@ -55,6 +106,7 @@ scanomapToMapAndScan ::   ( MonadFreshNames m,     Buildable rep,+    LetDec rep ~ Type,     ExpDec rep ~ (),     Op rep ~ SOAC rep   ) =>@@ -66,9 +118,9 @@   ) ->   m (Stm rep, Stm rep) scanomapToMapAndScan (Pat pes) (w, scans, map_lam, arrs) = do-  (map_pat, scan_pat, scan_arrs) <-+  (map_pat, scan_pat, scan_arrs, map_lam') <-     splitScanOrRedomap pes w map_lam $ map scanNeutral scans-  map_stm <- mkLet map_pat . Op . Screma w arrs <$> mapSOAC map_lam+  map_stm <- Let map_pat (defAux ()) . Op . Screma w arrs <$> mapSOAC map_lam'   scan_stm <-     Let scan_pat (defAux ()) . Op       <$> (Screma w scan_arrs <$> scanSOAC scans)@@ -100,27 +152,6 @@   where     tempRes res = newIdent "temp_res" $ res `arrayOfRow` w -splitScanOrRedomap ::-  (Typed dec, MonadFreshNames m) =>-  [PatElem dec] ->-  SubExp ->-  Lambda rep ->-  [[SubExp]] ->-  m ([Ident], Pat dec, [VName])-splitScanOrRedomap pes w map_lam nes = do-  let (acc_pes, arr_pes) =-        splitAt (length $ concat nes) pes-      (acc_ts, _arr_ts) =-        splitAt (length (concat nes)) $ lambdaReturnType map_lam-  map_accpat <- zipWithM accMapPatElem acc_pes acc_ts-  map_arrpat <- mapM arrMapPatElem arr_pes-  let map_pat = map_accpat ++ map_arrpat-  pure (map_pat, Pat acc_pes, map identName map_accpat)-  where-    accMapPatElem pe acc_t =-      newIdent (baseName (patElemName pe) <> "_map_acc") $ acc_t `arrayOfRow` w-    arrMapPatElem = pure . patElemIdent- -- | Turn a Screma into a maposcanomap (possibly with mapout parts) and a -- Redomap.  This is used to handle Scremas that are so complicated -- that we cannot directly generate efficient parallel code for them.@@ -304,3 +335,39 @@       =<< mapSOAC map_lam    letTupExp desc $ WithAcc [(acc_shape, [v], Nothing) | v <- dest] withacc_lam++-- | The most addition-like binary operator for some primitive type.+addBinOp :: PrimType -> BinOp+addBinOp (IntType it) = Add it OverflowWrap+addBinOp (FloatType ft) = FAdd ft+addBinOp Bool = LogAnd+addBinOp Unit = LogAnd++-- | Construct a lambda for adding two values of the given type, Using SOACs to handle arrays.+addLambda ::+  ( OpC (Rep m) ~ SOAC,+    MonadBuilder m,+    Buildable (Rep m)+  ) =>+  TypeBase Shape NoUniqueness ->+  m (Lambda (Rep m))+addLambda (Prim pt) = binOpLambda (addBinOp pt) pt+addLambda t@Array {} = do+  xs_p <- newParam "xs" t+  ys_p <- newParam "ys" t+  lam <- addLambda $ rowType t+  body <- insertStmsM $ do+    res <-+      letSubExp "lam_map"+        . Op+        . Screma (arraySize 0 t) [paramName xs_p, paramName ys_p]+        =<< mapSOAC lam+    pure $ resultBody [res]+  pure+    Lambda+      { lambdaParams = [xs_p, ys_p],+        lambdaReturnType = [t],+        lambdaBody = body+      }+addLambda t =+  error $ "addLambda: " ++ show t
src/Futhark/Util.hs view
@@ -55,6 +55,8 @@     topologicalSort,     debugTraceM,     ensureCacheDirectory,+    interleave,+    unterleave,   ) where @@ -375,6 +377,15 @@   | isInfinite v, v < 0 = -1 / 0   | isNaN v = 0 / 0   | otherwise = fromRational $ toRational v++-- | Interleave two lists.+interleave :: [a] -> [a] -> [a]+interleave xs ys = concat $ L.transpose [xs, ys]++-- | The inverse of interleave.+unterleave :: [a] -> ([a], [a])+unterleave (x : y : xys) = bimap (x :) (y :) $ unterleave xys+unterleave _ = ([], [])  -- Z-encoding from https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/SymbolNames --
src/Language/Futhark/Interpreter.hs view
@@ -22,6 +22,7 @@     Value,     fromTuple,     isEmptyArray,+    asByteString,     prettyEmptyArray,     prettyValue,     valueText,@@ -37,6 +38,7 @@ import Data.Array import Data.Bifunctor import Data.Bitraversable+import Data.ByteString qualified as BS import Data.Functor (($>), (<&>)) import Data.List   ( find,@@ -298,6 +300,14 @@  type Value = Language.Futhark.Interpreter.Values.Value EvalM +-- | If the value represents an array of type @[]i8@, then return those bytes.+asByteString :: Value -> Maybe BS.ByteString+asByteString (ValueArray _ vals) = BS.pack <$> mapM asU8 (elems vals)+  where+    asU8 (ValuePrim (UnsignedValue (Int8Value x))) = Just $ fromIntegral x+    asU8 _ = Nothing+asByteString _ = Nothing+ asInteger :: Value -> Integer asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v asInteger (ValuePrim (UnsignedValue v)) =@@ -2154,6 +2164,18 @@     def "manifest" = Just $ fun1 pure     def "jvp2" = Just $ fun3 doJVP2     def "vjp2" = Just $ fun3 doVJP2+    def "jmp2" = Just $ fun3 $ \f x seeds -> do+      v <- apply noLoc mempty f x+      dvs <-+        toArray' (valueShape v) . map (project "1")+          <$> mapM (doJVP2 f x) (snd (fromArray seeds))+      pure $ toTuple [v, dvs]+    def "mjp2" = Just $ fun3 $ \f x seeds -> do+      v <- apply noLoc mempty f x+      dvs <-+        toArray' (valueShape x) . map (project "1")+          <$> mapM (doVJP2 f x) (snd (fromArray seeds))+      pure $ toTuple [v, dvs]     def "with_vjp" = Just $ fun3 $ \f _ arg ->       -- XXX? We simply ignore the custom derivative. This is correct, but makes       -- it more of a hassle to test them.
src/Language/Futhark/Prop.hs view
@@ -973,6 +973,34 @@                   $ Scalar                   $ tupleRecord [Scalar $ t_b Nonunique, Scalar $ t_a Nonunique]               ),+              ( "jmp2",+                IntrinsicPolyFun+                  [tp_a, tp_b, sp_n]+                  [ Scalar (t_a mempty) `arr` Scalar (t_b Nonunique),+                    Scalar (t_a Observe),+                    array_a Observe $ shape [n]+                  ]+                  $ RetType []+                  $ Scalar+                  $ tupleRecord+                    [ Scalar $ t_b Nonunique,+                      array_b Unique $ shape [n]+                    ]+              ),+              ( "mjp2",+                IntrinsicPolyFun+                  [tp_a, tp_b, sp_n]+                  [ Scalar (t_a mempty) `arr` Scalar (t_b Nonunique),+                    Scalar (t_a Observe),+                    array_b Observe $ shape [n]+                  ]+                  $ RetType []+                  $ Scalar+                  $ tupleRecord+                    [ Scalar $ t_b Nonunique,+                      array_a Unique $ shape [n]+                    ]+              ),               ( "with_vjp",                 IntrinsicPolyFun                   [tp_a, tp_b]
src/Language/Futhark/Syntax.hs view
@@ -1092,7 +1092,8 @@ -- points, so the types can be either ascribed or inferred. data EntryPoint = EntryPoint   { entryParams :: [EntryParam],-    entryReturn :: EntryType+    entryReturn :: EntryType,+    entryDoc :: Maybe T.Text   }   deriving (Show)
src/Language/Futhark/TypeChecker.hs view
@@ -599,10 +599,14 @@           TypeBind name' l tps' te' (Info elab_t) doc loc         ) -entryPoint :: [Pat ParamType] -> Maybe (TypeExp Exp VName) -> ResRetType -> EntryPoint-entryPoint params orig_ret_te (RetType _ret orig_ret) =-  EntryPoint (map patternEntry params ++ more_params) rettype'+entryPoint :: Maybe DocComment -> [Pat ParamType] -> Maybe (TypeExp Exp VName) -> ResRetType -> EntryPoint+entryPoint doc params orig_ret_te (RetType _ret orig_ret) =+  EntryPoint (map patternEntry params ++ more_params) rettype' doc'   where+    doc' = case doc of+      Just (DocComment t _) -> Just t+      _ -> Nothing+     (more_params, rettype') = onRetType orig_ret_te $ toStruct orig_ret      patternEntry (PatParens p _) =@@ -690,7 +694,7 @@   (tparams', params', maybe_tdecl', rettype, body') <-     checkFunDef (fname, maybe_tdecl, tparams, params, body, loc) -  let entry' = Info (entryPoint params' maybe_tdecl' rettype) <$ entry+  let entry' = Info (entryPoint doc params' maybe_tdecl' rettype) <$ entry   case entry' of     Just _ -> checkEntryPoint loc tparams' params' rettype     _ -> pure ()
src/Language/Futhark/TypeChecker/Consumption.hs view
@@ -580,7 +580,7 @@ applyArg :: TypeAliases -> TypeAliases -> TypeAliases applyArg (Scalar (Arrow closure_als _ d _ (RetType _ rettype))) arg_als =   returnType closure_als rettype d arg_als-applyArg t _ = error $ "applyArg: " <> show t+applyArg _ arg_als = arg_als  applyLoopArg :: Aliases -> ParamType -> TypeAliases -> ResType -> TypeAliases applyLoopArg appres (Scalar (Record pfs)) (Scalar (Record afs)) (Scalar (Record rfs)) =@@ -770,7 +770,7 @@ -- checkExp (AppExp (Apply f args loc) appres) = do   (f', f_als) <- checkExp f-  (args', args_als) <- NE.unzip <$> checkArgs (toRes Nonunique f_als) args+  (args', args_als) <- NE.unzip <$> checkArgs (diets $ toRes Nonunique f_als) args   res_als <- checkFuncall loc (fname f) f_als args_als   pure (AppExp (Apply f' args' loc) appres, res_als)   where@@ -781,13 +781,16 @@       (e', e_als) <- checkArg prev (second (const d) (typeOf e)) e       pure ((Info p, e'), e_als) -    checkArgs (Scalar (Arrow _ _ d _ (RetType _ rt))) (x NE.:| args') = do+    diets (Scalar (Arrow _ _ d _ (RetType _ rt))) =+      d : diets rt+    diets _ = repeat Observe++    checkArgs ds (x NE.:| args') = do+      let (d, ds') = fromMaybe (Observe, []) $ L.uncons ds       -- Note Futhark uses right-to-left evaluation of applications.-      args'' <- maybe (pure []) (fmap NE.toList . checkArgs rt) $ NE.nonEmpty args'+      args'' <- maybe (pure []) (fmap NE.toList . checkArgs ds') $ NE.nonEmpty args'       (x', x_als) <- checkArg' (map (first snd) args'') d x       pure $ (x', x_als) NE.:| args''-    checkArgs t _ =-      error $ "checkArgs: " <> prettyString t  -- checkExp (AppExp (Loop sparams pat loopinit form body loc) appres) = do
src/Language/Futhark/TypeChecker/Match.hs view
@@ -73,10 +73,20 @@ isConstr (MatchConstr (Constr c) _ _) = Just c isConstr _ = Nothing +isWild :: Match t -> Bool+isWild MatchWild {} = True+isWild _ = False+ isBool :: Match t -> Maybe Bool isBool (MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) _ _) = Just b isBool _ = Nothing +isStructuralConstr :: Match t -> Bool+isStructuralConstr (MatchConstr (Constr _) _ _) = True+isStructuralConstr (MatchConstr ConstrTuple _ _) = True+isStructuralConstr (MatchConstr (ConstrRecord _) _ _) = True+isStructuralConstr _ = False+ complete :: [Match StructType] -> Bool complete xs   | Just x <- maybeHead xs,@@ -147,28 +157,56 @@         MatchWild _ ->           MatchWild () : u -    incompleteCase pt cs = do-      u <- findUnmatched (defaultMat pmat) (n - 1)-      if null cs-        then pure $ MatchWild () : u-        else case pt of-          Scalar (Sum all_cs) -> do-            -- Figure out which constructors are missing.-            let sigma = mapMaybe isConstr cs-                notCovered (k, _) = k `notElem` sigma-            (cname, ts) <- filter notCovered $ M.toList all_cs-            pure $ MatchConstr (Constr cname) (map (const (MatchWild ())) ts) () : u-          Scalar (Prim Bool) -> do-            -- Figure out which constants are missing.-            let sigma = mapMaybe isBool cs-            b <- filter (`notElem` sigma) [True, False]-            pure $ MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) [] () : u-          _ -> do-            -- FIXME: this is wrong in the unlikely case where someone-            -- is pattern-matching every single possible number for-            -- some numeric type.  It should be handled more like Bool-            -- above.-            pure $ MatchWild () : u+    -- When wildcards make the column complete, recurse on unique constructor+    -- heads (wildcard rows are included via specialise).+    recurseOnConstrHeads cs = do+      c@(MatchConstr c' args _) <- nubOrd $ filter isStructuralConstr cs+      let ats = map matchType args+          a_k = length ats+          pmat' = specialise ats c pmat+      u <- findUnmatched pmat' (a_k + n - 1)+      let (r, rest) = splitAt a_k u+      pure $ MatchConstr c' r () : rest++    incompleteCase pt cs+      | null cs = do+          u <- findUnmatched (defaultMat pmat) (n - 1)+          pure $ MatchWild () : u+      | any isWild cs,+        Scalar (Sum all_cs) <- pt =+          let sigma = mapMaybe isConstr cs+              notCovered (k, _) = k `notElem` sigma+              missingCs = filter notCovered $ M.toList all_cs+              -- Sigma constructors: check sub-patterns (wildcard rows included via specialise).+              sigmaWitnesses = recurseOnConstrHeads cs+              -- Missing constructors: wildcard covers them, but sub-patterns+              -- may still be unmatched (checked via default matrix).+              missingWitnesses = do+                u <- findUnmatched (defaultMat pmat) (n - 1)+                (cname, ts) <- missingCs+                pure $ MatchConstr (Constr cname) (map (const (MatchWild ())) ts) () : u+           in sigmaWitnesses ++ missingWitnesses+      | any isWild cs,+        any isStructuralConstr cs =+          recurseOnConstrHeads cs+      | Scalar (Sum all_cs) <- pt = do+          let sigma = mapMaybe isConstr cs+              notCovered (k, _) = k `notElem` sigma+          (cname, ts) <- filter notCovered $ M.toList all_cs+          u <- findUnmatched (defaultMat pmat) (n - 1)+          pure $ MatchConstr (Constr cname) (map (const (MatchWild ())) ts) () : u+      | Scalar (Prim Bool) <- pt = do+          u <- findUnmatched (defaultMat pmat) (n - 1)+          let sigma = mapMaybe isBool cs+          b <- filter (`notElem` sigma) [True, False]+          pure $ MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) [] () : u+      | otherwise = do+          u <- findUnmatched (defaultMat pmat) (n - 1)+          -- FIXME: this is wrong in the unlikely case where someone+          -- is pattern-matching every single possible number for+          -- some numeric type.  It should be handled more like Bool+          -- above.+          pure $ MatchWild () : u findUnmatched [] n = [replicate n $ MatchWild ()] findUnmatched _ _ = []