diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/man/futhark-hip.rst b/docs/man/futhark-hip.rst --- a/docs/man/futhark-hip.rst +++ b/docs/man/futhark-hip.rst @@ -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 ======== diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -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 diff --git a/prelude/ad.fut b/prelude/ad.fut --- a/prelude/ad.fut +++ b/prelude/ad.fut @@ -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 diff --git a/rts/c/backends/hip.h b/rts/c/backends/hip.h --- a/rts/c/backends/hip.h +++ b/rts/c/backends/hip.h @@ -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]); } diff --git a/rts/c/scalar_f16.h b/rts/c/scalar_f16.h --- a/rts/c/scalar_f16.h +++ b/rts/c/scalar_f16.h @@ -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 +#elif defined(__HIP_DEVICE_COMPILE__) +#include #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); } diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -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. diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs --- a/src/Futhark/AD/Rev.hs +++ b/src/Futhark/AD/Rev.hs @@ -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. diff --git a/src/Futhark/AD/Rev/Acc.hs b/src/Futhark/AD/Rev/Acc.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Acc.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Hist.hs b/src/Futhark/AD/Rev/Hist.hs --- a/src/Futhark/AD/Rev/Hist.hs +++ b/src/Futhark/AD/Rev/Hist.hs @@ -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) diff --git a/src/Futhark/AD/Rev/Loop.hs b/src/Futhark/AD/Rev/Loop.hs --- a/src/Futhark/AD/Rev/Loop.hs +++ b/src/Futhark/AD/Rev/Loop.hs @@ -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' = diff --git a/src/Futhark/AD/Rev/Map.hs b/src/Futhark/AD/Rev/Map.hs --- a/src/Futhark/AD/Rev/Map.hs +++ b/src/Futhark/AD/Rev/Map.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Monad.hs b/src/Futhark/AD/Rev/Monad.hs --- a/src/Futhark/AD/Rev/Monad.hs +++ b/src/Futhark/AD/Rev/Monad.hs @@ -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] -- diff --git a/src/Futhark/AD/Rev/Reduce.hs b/src/Futhark/AD/Rev/Reduce.hs --- a/src/Futhark/AD/Rev/Reduce.hs +++ b/src/Futhark/AD/Rev/Reduce.hs @@ -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 diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs --- a/src/Futhark/AD/Rev/SOAC.hs +++ b/src/Futhark/AD/Rev/SOAC.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs --- a/src/Futhark/AD/Rev/Scan.hs +++ b/src/Futhark/AD/Rev/Scan.hs @@ -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 diff --git a/src/Futhark/AD/Shared.hs b/src/Futhark/AD/Shared.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Shared.hs @@ -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 diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -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 diff --git a/src/Futhark/CLI/REPL.hs b/src/Futhark/CLI/REPL.hs --- a/src/Futhark/CLI/REPL.hs +++ b/src/Futhark/CLI/REPL.hs @@ -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", diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs --- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericC/Code.hs b/src/Futhark/CodeGen/Backends/GenericC/Code.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Code.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Code.hs @@ -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'|] diff --git a/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs b/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs --- a/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs @@ -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)) diff --git a/src/Futhark/CodeGen/Backends/GenericC/Server.hs b/src/Futhark/CodeGen/Backends/GenericC/Server.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Server.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Server.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericC/Types.hs b/src/Futhark/CodeGen/Backends/GenericC/Types.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Types.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Types.hs @@ -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) = diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs --- a/src/Futhark/CodeGen/Backends/GenericPython.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericWASM.hs b/src/Futhark/CodeGen/Backends/GenericWASM.hs --- a/src/Futhark/CodeGen/Backends/GenericWASM.hs +++ b/src/Futhark/CodeGen/Backends/GenericWASM.hs @@ -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) diff --git a/src/Futhark/CodeGen/Backends/MulticoreWASM.hs b/src/Futhark/CodeGen/Backends/MulticoreWASM.hs --- a/src/Futhark/CodeGen/Backends/MulticoreWASM.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreWASM.hs @@ -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, diff --git a/src/Futhark/CodeGen/Backends/SequentialWASM.hs b/src/Futhark/CodeGen/Backends/SequentialWASM.hs --- a/src/Futhark/CodeGen/Backends/SequentialWASM.hs +++ b/src/Futhark/CodeGen/Backends/SequentialWASM.hs @@ -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, diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -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' diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -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 -> diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs @@ -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 diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -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 diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs --- a/src/Futhark/IR/MC/Op.hs +++ b/src/Futhark/IR/MC/Op.hs @@ -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 diff --git a/src/Futhark/IR/Mem/LMAD.hs b/src/Futhark/IR/Mem/LMAD.hs --- a/src/Futhark/IR/Mem/LMAD.hs +++ b/src/Futhark/IR/Mem/LMAD.hs @@ -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) => diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -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 diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -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) = diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs --- a/src/Futhark/IR/Prop.hs +++ b/src/Futhark/IR/Prop.hs @@ -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 _ _ _) = diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -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) diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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 diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -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 diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -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 diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs --- a/src/Futhark/IR/Syntax/Core.hs +++ b/src/Futhark/IR/Syntax/Core.hs @@ -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 diff --git a/src/Futhark/IR/TypeCheck.hs b/src/Futhark/IR/TypeCheck.hs --- a/src/Futhark/IR/TypeCheck.hs +++ b/src/Futhark/IR/TypeCheck.hs @@ -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 diff --git a/src/Futhark/Internalise/ApplyTypeAbbrs.hs b/src/Futhark/Internalise/ApplyTypeAbbrs.hs --- a/src/Futhark/Internalise/ApplyTypeAbbrs.hs +++ b/src/Futhark/Internalise/ApplyTypeAbbrs.hs @@ -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 diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs --- a/src/Futhark/Internalise/Defunctorise.hs +++ b/src/Futhark/Internalise/Defunctorise.hs @@ -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 diff --git a/src/Futhark/Internalise/Entry.hs b/src/Futhark/Internalise/Entry.hs --- a/src/Futhark/Internalise/Entry.hs +++ b/src/Futhark/Internalise/Entry.hs @@ -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) diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -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 diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs deleted file mode 100644 --- a/src/Futhark/Internalise/Lambdas.hs +++ /dev/null @@ -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 diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -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 () diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -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) diff --git a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs --- a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs +++ b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs @@ -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 --- diff --git a/src/Futhark/Optimise/Fusion/Screma.hs b/src/Futhark/Optimise/Fusion/Screma.hs --- a/src/Futhark/Optimise/Fusion/Screma.hs +++ b/src/Futhark/Optimise/Fusion/Screma.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -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) diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -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 -> diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs --- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -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, diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -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 diff --git a/src/Futhark/Pass/AD.hs b/src/Futhark/Pass/AD.hs --- a/src/Futhark/Pass/AD.hs +++ b/src/Futhark/Pass/AD.hs @@ -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. diff --git a/src/Futhark/Test/Property.hs b/src/Futhark/Test/Property.hs --- a/src/Futhark/Test/Property.hs +++ b/src/Futhark/Test/Property.hs @@ -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 -> diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -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 diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -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 -- diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -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. diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -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] diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs --- a/src/Language/Futhark/Syntax.hs +++ b/src/Language/Futhark/Syntax.hs @@ -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) diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -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 () diff --git a/src/Language/Futhark/TypeChecker/Consumption.hs b/src/Language/Futhark/TypeChecker/Consumption.hs --- a/src/Language/Futhark/TypeChecker/Consumption.hs +++ b/src/Language/Futhark/TypeChecker/Consumption.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Match.hs b/src/Language/Futhark/TypeChecker/Match.hs --- a/src/Language/Futhark/TypeChecker/Match.hs +++ b/src/Language/Futhark/TypeChecker/Match.hs @@ -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 _ _ = []