futhark 0.25.22 → 0.25.23

raw patch · 29 files changed

+1558/−257 lines, 29 filesdep ~futhark-dataPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: futhark-data

API changes (from Hackage documentation)

Files

• CHANGELOG.md +20/−4
• docs/language-reference.rst +3/−2
• futhark.cabal +4/−2
• rts/c/atomics.h +47/−18
• rts/c/tuning.h +1/−1
• rts/c/values.h +7/−26
• rts/python/values.py +7/−4
• src/Futhark/Analysis/HORep/MapNest.hs +1/−1
• src/Futhark/Builder.hs +9/−7
• src/Futhark/CodeGen/ImpGen/GPU.hs +3/−5
• src/Futhark/IR/SOACS/Simplify.hs +1/−1
• src/Futhark/IR/Syntax.hs +3/−1
• src/Futhark/Internalise/Exps.hs +6/−6
• src/Futhark/Internalise/Lambdas.hs +1/−1
• src/Futhark/Internalise/Monomorphise.hs +26/−23
• src/Futhark/Optimise/BlkRegTiling.hs +4/−15
• src/Futhark/Optimise/Fusion.hs +145/−35
• src/Futhark/Optimise/Fusion/GraphRep.hs +26/−0
• src/Futhark/Optimise/Fusion/RulesWithAccs.hs +607/−0
• src/Futhark/Optimise/Fusion/TryFusion.hs +2/−1
• src/Futhark/Optimise/Simplify/Engine.hs +7/−1
• src/Futhark/Optimise/TileLoops.hs +4/−11
• src/Futhark/Optimise/TileLoops/Shared.hs +4/−3
• src/Futhark/Pass/ExtractKernels/Interchange.hs +5/−6
• src/Futhark/Script.hs +3/−3
• src/Futhark/Transform/FirstOrderTransform.hs +5/−5
• src/Language/Futhark/Interpreter.hs +243/−74
• src/Language/Futhark/Interpreter/AD.hs +320/−0
• src/Language/Futhark/Interpreter/Values.hs +44/−1

@@ -5,6 +5,26 @@ 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.25.23]++### Added++* Trailing commas are now allowed for arrays, records, and tuples in+  the textual value format and in FutharkScript.++* Faster floating-point atomics with OpenCL backend on AMD and NVIDIA+  GPUs. This affects histogram workloads.++* AD is now supported by the interpreter (thanks to Marcus Jensen).++### Fixed++* Some instances of invalid copy removal. (Again.)++* An issue related to entry points with nontrivial sizes in their+  arguments, where the entry points were also used as normal functions+  elsewhere. (#2184)+ ## [0.25.22]  ### Added@@ -12,10 +32,6 @@ * `futhark script` now supports an `-f` option.  * `futhark script` now supports the builtin procedure `$store`.--### Removed--### Changed  ### Fixed 

@@ -880,7 +880,7 @@ ``#c x y z`` ............ -Apply the sum type constructor ``#x`` to the payload ``x``, ``y``, and+Apply the sum type constructor ``#c`` to the payload ``x``, ``y``, and ``z``.  A constructor application is always assumed to be saturated, i.e. its entire payload provided.  This means that constructors may not be partially applied.@@ -1462,7 +1462,8 @@   def consumes_first_arg (a: *[]i32) (b: []i32) = ...  For bulk in-place updates with multiple values, use the ``scatter``-function in the basis library.+function from the `prelude+<https://futhark-lang.org/docs/prelude/doc/prelude/soacs.html>`_.  Alias Analysis ~~~~~~~~~~~~~~

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.22+version:        0.25.23 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -313,6 +313,7 @@       Futhark.Optimise.Fusion       Futhark.Optimise.Fusion.Composing       Futhark.Optimise.Fusion.GraphRep+      Futhark.Optimise.Fusion.RulesWithAccs       Futhark.Optimise.Fusion.TryFusion       Futhark.Optimise.GenRedOpt       Futhark.Optimise.HistAccs@@ -396,6 +397,7 @@       Language.Futhark       Language.Futhark.Core       Language.Futhark.Interpreter+      Language.Futhark.Interpreter.AD       Language.Futhark.Interpreter.Values       Language.Futhark.FreeVars       Language.Futhark.Parser@@ -457,7 +459,7 @@     , file-embed >=0.0.14.0     , filepath >=1.4.1.1     , free >=5.1.10-    , futhark-data >= 1.1.0.0+    , futhark-data >= 1.1.1.0     , futhark-server >= 1.2.2.1     , futhark-manifest >= 1.5.0.0     , githash >=0.1.6.1

@@ -78,16 +78,29 @@ SCALAR_FUN_ATTR float atomic_fadd_f32_global(volatile __global float *p, float x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((float*)p, x);+  // On OpenCL, use technique from+  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html+#elif defined(cl_nv_pragma_unroll)+  // use hardware-supported atomic addition on Nvidia GPUs with inline+  // PTX assembly+  float ret;+  asm volatile("atom.global.add.f32 %0,[%1],%2;":"=f"(ret):"l"(p),"f"(x):"memory");+  return ret;+#elif defined(__opencl_c_ext_fp32_global_atomic_add)+  // use hardware-supported atomic addition on some Intel GPUs+  return atomic_fetch_add_explicit((volatile __global atomic_float*)p,+                                   x,+                                   memory_order_relaxed);+#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f32)+  // use hardware-supported atomic addition on some AMD GPUs+  return __builtin_amdgcn_global_atomic_fadd_f32(p, x); #else-  union { int32_t i; float f; } old;-  union { int32_t i; float f; } assumed;-  old.f = *p;-  do {-    assumed.f = old.f;-    old.f = old.f + x;-    old.i = atomic_cmpxchg_i32_global((volatile __global int32_t*)p, assumed.i, old.i);-  } while (assumed.i != old.i);-  return old.f;+  // fallback emulation:+  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5+  float old = x;+  float ret;+  while ((old=atomic_xchg(p, ret=atomic_xchg(p, 0.0f)+old))!=0.0f);+  return ret; #endif } @@ -306,16 +319,32 @@ SCALAR_FUN_ATTR double atomic_fadd_f64_global(volatile __global double *p, double x) { #if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 || defined(FUTHARK_HIP)   return atomicAdd((double*)p, x);+  // On OpenCL, use technique from+  // https://pipinspace.github.io/blog/atomic-float-addition-in-opencl.html+#elif defined(cl_nv_pragma_unroll)+  // use hardware-supported atomic addition on Nvidia GPUs with inline+  // PTX assembly+  double ret;+  asm volatile("atom.global.add.f64 %0,[%1],%2;":"=d"(ret):"l"(p),"d"(x):"memory");+  return ret;+#elif __has_builtin(__builtin_amdgcn_global_atomic_fadd_f64)+  // use hardware-supported atomic addition on some AMD GPUs+  return __builtin_amdgcn_global_atomic_fadd_f64(p, x); #else-  union { int64_t i; double f; } old;-  union { int64_t i; double f; } assumed;-  old.f = *p;-  do {-    assumed.f = old.f;-    old.f = old.f + x;-    old.i = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed.i, old.i);-  } while (assumed.i != old.i);-  return old.f;+  // fallback emulation:+  // https://forums.developer.nvidia.com/t/atomicadd-float-float-atomicmul-float-float/14639/5+  union {int64_t i; double f;} old;+  union {int64_t i; double f;} ret;+  old.f = x;+  while (1) {+    ret.i = atom_xchg((volatile __global int64_t*)p, (int64_t)0);+    ret.f += old.f;+    old.i = atom_xchg((volatile __global int64_t*)p, ret.i);+    if (old.i == 0) {+      break;+    }+  }+  return ret.f; #endif } 

@@ -2,7 +2,7 @@   int is_blank_line_or_comment(const char *s) {-  size_t i = strspn(s, " \t");+  size_t i = strspn(s, " \t\n");   return s[i] == '\0' || // Line is blank.          strncmp(s + i, "--", 2) == 0; // Line is comment. }

@@ -104,16 +104,16 @@ static int read_str_array_elems(FILE *f,                                 char *buf, int bufsize,                                 struct array_reader *reader, int64_t dims) {-  int ret;-  int first = 1;+  int ret = 1;+  int expect_elem = 1;   char *knows_dimsize = (char*) calloc((size_t)dims, sizeof(char));   int cur_dim = (int)dims-1;   int64_t *elems_read_in_dim = (int64_t*) calloc((size_t)dims, sizeof(int64_t));    while (1) {     next_token(f, buf, bufsize);-     if (strcmp(buf, "]") == 0) {+      expect_elem = 0;       if (knows_dimsize[cur_dim]) {         if (reader->shape[cur_dim] != elems_read_in_dim[cur_dim]) {           ret = 1;@@ -130,14 +130,14 @@         cur_dim--;         elems_read_in_dim[cur_dim]++;       }-    } else if (strcmp(buf, ",") == 0) {-      next_token(f, buf, bufsize);+    } else if (!expect_elem && strcmp(buf, ",") == 0) {+      expect_elem = 1;+    } else if (expect_elem) {       if (strcmp(buf, "[") == 0) {         if (cur_dim == dims - 1) {           ret = 1;           break;         }-        first = 1;         cur_dim++;         elems_read_in_dim[cur_dim] = 0;       } else if (cur_dim == dims - 1) {@@ -145,30 +145,11 @@         if (ret != 0) {           break;         }+        expect_elem = 0;         elems_read_in_dim[cur_dim]++;       } else {         ret = 1;         break;-      }-    } else if (strlen(buf) == 0) {-      // EOF-      ret = 1;-      break;-    } else if (first) {-      if (strcmp(buf, "[") == 0) {-        if (cur_dim == dims - 1) {-          ret = 1;-          break;-        }-        cur_dim++;-        elems_read_in_dim[cur_dim] = 0;-      } else {-        ret = read_str_elem(buf, reader);-        if (ret != 0) {-          break;-        }-        elems_read_in_dim[cur_dim]++;-        first = 0;       }     } else {       ret = 1;

@@ -107,14 +107,17 @@         return False  -def sepBy(p, sep, *args):+def sepEndBy(p, sep, *args):     elems = []     x = optional(p, *args)     if x != None:         elems += [x]         while optional(sep, *args) != None:-            x = p(*args)-            elems += [x]+            x = optional(p, *args)+            if x == None:+                break+            else:+                elems += [x]     return elems  @@ -372,7 +375,7 @@     except ValueError:         return read_str_empty_array(f, type_name, rank)     else:-        xs = sepBy(elem_reader, read_str_comma, f)+        xs = sepEndBy(elem_reader, read_str_comma, f)         skip_spaces(f)         parse_specific_char(f, b"]")         return xs

@@ -154,7 +154,7 @@       letBindNames nres         =<< SOAC.toExp         =<< toSOAC (MapNest nw lam ns $ map (SOAC.identInput . paramIdent) nparams)-      pure $ resultBody $ map Var nres+      pure $ varsRes nres   let outerlam =         Lambda           { lambdaParams = nparams,

@@ -88,7 +88,7 @@ instance MonadTrans (BuilderT rep) where   lift = BuilderT . lift --- | The most commonly used binder monad.+-- | The most commonly used builder monad. type Builder rep = BuilderT rep (State VNameSource)  instance (MonadFreshNames m) => MonadFreshNames (BuilderT rep m) where@@ -138,7 +138,7 @@     BuilderT $ put (old_stms, old_scope)     pure (x, new_stms) --- | Run a binder action given an initial scope, returning a value and+-- | Run a builder action given an initial scope, returning a value and -- the statements added ('addStm') during the action. runBuilderT ::   (MonadFreshNames m) =>@@ -175,7 +175,7 @@   m (Stms rep) runBuilderT'_ = fmap snd . runBuilderT' --- | Run a binder action, returning a value and the statements added+-- | Run a builder action, returning a value and the statements added -- ('addStm') during the action.  Assumes that the current monad -- provides initial scope and name source. runBuilder ::@@ -194,17 +194,19 @@   m (Stms rep) runBuilder_ = fmap snd . runBuilder --- | Run a binder that produces a t'Body', and prefix that t'Body' by--- the statements produced during execution of the action.+-- | Run a builder that produces a 'Result' and construct a body that+-- contains that result alongside the statements produced during the+-- builder. runBodyBuilder ::   ( Buildable rep,     MonadFreshNames m,     HasScope somerep m,     SameScope somerep rep   ) =>-  Builder rep (Body rep) ->+  Builder rep Result ->   m (Body rep)-runBodyBuilder = fmap (uncurry $ flip insertStms) . runBuilder+runBodyBuilder =+  fmap (uncurry $ flip insertStms) . runBuilder . fmap (mkBody mempty)  -- | Given lambda parameters, Run a builder action that produces the -- statements and returns the 'Result' of the lambda body.

@@ -45,6 +45,7 @@   where     opencl64 =       [ (Add Int64 OverflowUndef, Imp.AtomicAdd Int64),+        (FAdd Float64, Imp.AtomicFAdd Float64),         (SMax Int64, Imp.AtomicSMax Int64),         (SMin Int64, Imp.AtomicSMin Int64),         (UMax Int64, Imp.AtomicUMax Int64),@@ -55,6 +56,7 @@       ]     opencl32 =       [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32),+        (FAdd Float32, Imp.AtomicFAdd Float32),         (SMax Int32, Imp.AtomicSMax Int32),         (SMin Int32, Imp.AtomicSMin Int32),         (UMax Int32, Imp.AtomicUMax Int32),@@ -64,11 +66,7 @@         (Xor Int32, Imp.AtomicXor Int32)       ]     opencl = opencl32 ++ opencl64-    cuda =-      opencl-        ++ [ (FAdd Float32, Imp.AtomicFAdd Float32),-             (FAdd Float64, Imp.AtomicFAdd Float64)-           ]+    cuda = opencl  compileProg ::   (MonadFreshNames m) =>

@@ -154,7 +154,7 @@   body <- runBodyBuilder $     localScope (scopeOfLParams $ lambdaParams lam) $ do       zipWithM_ maybeFix (lambdaParams lam) fixes'-      pure $ lambdaBody lam+      bodyBind $ lambdaBody lam   pure     lam       { lambdaBody = body,

@@ -448,7 +448,9 @@   | -- | Create accumulators backed by the given arrays (which are     -- consumed) and pass them to the lambda, which must return the     -- updated accumulators and possibly some extra values.  The-    -- accumulators are turned back into arrays.  The t'Shape' is the+    -- accumulators are turned back into arrays.  In the lambda, the result+    -- accumulators come first, and are ordered in a manner consistent with+    -- that of the input (accumulator) arguments. The t'Shape' is the     -- write index space.  The corresponding arrays must all have this     -- shape outermost.  This construct is not part of t'BasicOp'     -- because we need the @rep@ parameter.

@@ -710,9 +710,9 @@   forM (zip ses ts) $ \(e', t') -> do     dims <- arrayDims <$> subExpType e'     let parts =-          ["Value of (core language) shape ("]-            ++ intersperse ", " (map (ErrorVal int64) dims)-            ++ [") cannot match shape of type `"]+          ["Value of (desugared) shape ["]+            ++ intersperse "][" (map (ErrorVal int64) dims)+            ++ ["] cannot match shape of type `"]             ++ dt'             ++ ["`."]     ensureExtShape (errorMsg parts) loc (I.fromDecl t') desc e'@@ -1591,7 +1591,7 @@                     and_lam <- binOpLambda I.LogAnd I.Bool                     reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]]                     all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems [cmps] reduce-                    pure $ resultBody [all_equal]+                    pure $ subExpsRes [all_equal]                letSubExp "arrays_equal"                 =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])@@ -2094,7 +2094,7 @@     replicateM k $ newParam "y" (I.Prim int64)   add_lam_body <- runBodyBuilder $     localScope (scopeOfLParams $ add_lam_x_params ++ add_lam_y_params) $-      fmap resultBody $+      fmap subExpsRes $         forM (zip add_lam_x_params add_lam_y_params) $ \(x, y) ->           letSubExp "z" $             I.BasicOp $@@ -2118,7 +2118,7 @@   -- 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 resultBody $+        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)

@@ -80,4 +80,4 @@     lambdaWithIncrement lam_body = runBodyBuilder $ do       eq_class <-         maybe (intConst Int64 0) resSubExp . listToMaybe <$> bodyBind lam_body-      resultBody <$> mkResult eq_class 0+      subExpsRes <$> mkResult eq_class 0

@@ -34,6 +34,7 @@ import Data.List (partition) import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M+import Data.Maybe (isJust, isNothing) import Data.Sequence qualified as Seq import Data.Set qualified as S import Futhark.MonadFreshNames@@ -51,7 +52,8 @@ -- parameters. newtype PolyBinding   = PolyBinding-      ( VName,+      ( Maybe EntryPoint,+        VName,         [TypeParam],         [Pat ParamType],         ResRetType,@@ -60,10 +62,10 @@         SrcLoc       ) --- | To deduplicate size expressions, we want a looser notation of+-- | To deduplicate size expressions, we want a looser notion of -- equality than the strict syntactical equality provided by the Eq--- instance on Exp.  This newtype wrapper provides such a looser--- notion of equality.+-- instance on Exp. This newtype wrapper provides such a looser notion+-- of equality. newtype ReplacedExp = ReplacedExp {unReplaced :: Exp}   deriving (Show) @@ -397,7 +399,7 @@         (Nothing, Nothing) -> pure $ var fname t'         -- A polymorphic function.         (Nothing, Just funbind) -> do-          (fname', infer, funbind') <- monomorphiseBinding False funbind mono_t+          (fname', infer, funbind') <- monomorphiseBinding funbind mono_t           tell $ Seq.singleton (qualLeaf fname, funbind')           addLifted (qualLeaf fname) mono_t (fname', infer)           applySizeArgs fname' (toRes Nonunique t') <$> infer t'@@ -982,11 +984,10 @@ -- list. Monomorphises the body of the function as well. Returns the fresh name -- of the generated monomorphic function and its 'ValBind' representation. monomorphiseBinding ::-  Bool ->   PolyBinding ->   MonoType ->   MonoM (VName, InferSizeArgs, ValBind)-monomorphiseBinding entry (PolyBinding (name, tparams, params, rettype, body, attrs, loc)) inst_t = isolateNormalisation $ do+monomorphiseBinding (PolyBinding (entry, name, tparams, params, rettype, body, attrs, loc)) inst_t = isolateNormalisation $ do   let bind_t = funType params rettype   (substs, t_shape_params) <-     typeSubstsM loc bind_t $ noNamedParams inst_t@@ -1028,14 +1029,18 @@    seen_before <- elem name . map (fst . fst) <$> getLifts   name' <--    if null tparams && not entry && not seen_before+    if null tparams && isNothing entry && not seen_before       then pure name       else newName name    pure     ( name',-      inferSizeArgs shape_params_explicit bind_t'' bind_r,-      if entry+      -- If the function is an entry point, then it cannot possibly+      -- need any explicit size arguments (checked by type checker).+      if isJust entry+        then const $ pure []+        else inferSizeArgs shape_params_explicit bind_t'' bind_r,+      if isJust entry         then           toValBinding             name'@@ -1082,7 +1087,7 @@      toValBinding name' tparams' params'' rettype' body'' =       ValBind-        { valBindEntryPoint = Nothing,+        { valBindEntryPoint = Info <$> entry,           valBindName = name',           valBindRetType = Info rettype',           valBindRetDecl = Nothing,@@ -1174,23 +1179,21 @@   PatConstr n (Info tp) ps loc -> PatConstr n (Info $ f tp) ps loc  toPolyBinding :: ValBind -> PolyBinding-toPolyBinding (ValBind _ name _ (Info rettype) tparams params body _ attrs loc) =-  PolyBinding (name, tparams, params, rettype, body, attrs, loc)+toPolyBinding (ValBind entry name _ (Info rettype) tparams params body _ attrs loc) =+  PolyBinding (unInfo <$> entry, name, tparams, params, rettype, body, attrs, loc)  transformValBind :: ValBind -> MonoM Env transformValBind valbind = do   let valbind' = toPolyBinding valbind -  case valBindEntryPoint valbind of-    Nothing -> pure ()-    Just entry -> do-      let t =-            funType (valBindParams valbind) $-              unInfo $-                valBindRetType valbind-      (name, infer, valbind'') <- monomorphiseBinding True valbind' $ monoType t-      tell $ Seq.singleton (name, valbind'' {valBindEntryPoint = Just entry})-      addLifted (valBindName valbind) (monoType t) (name, infer)+  when (isJust $ valBindEntryPoint valbind) $ do+    let t =+          funType (valBindParams valbind) $+            unInfo $+              valBindRetType valbind+    (name, infer, valbind'') <- monomorphiseBinding valbind' $ monoType t+    tell $ Seq.singleton (name, valbind'')+    addLifted (valBindName valbind) (monoType t) (name, infer)    pure     mempty

@@ -128,12 +128,10 @@         let e = le64 i + le64 k * pe64 tr_par         pure (i, k, e)       ---      --       mkCompLoopRxRy fits_ij css_init (a_idx_fn, b_idx_fn) (ltid_y, ltid_x) = do         css <- forLoop ry [css_init] $ \i [css_merge] -> do           css <- forLoop rx [css_merge] $ \j [css_merge'] ->-            resultBodyM-              =<< letTupExp' "foo"+            (resultBodyM <=< letTupExp' "foo")               =<< eIf                 ( toExp $                     if fits_ij@@ -143,16 +141,8 @@                       -- is garbage anyways and should not be written.                       -- so fits_ij should be always true!!! -                        le64 iii-                          + le64 i-                          + pe64 ry-                            * le64 ltid_y-                              .<. pe64 height_A-                              .&&. le64 jjj-                          + le64 j-                          + pe64 rx-                            * le64 ltid_x-                              .<. pe64 width_B+                        (le64 iii + le64 i + pe64 ry * le64 ltid_y .<. pe64 height_A)+                          .&&. (le64 jjj + le64 j + pe64 rx * le64 ltid_x .<. pe64 width_B)                 )                 ( do                     a <- a_idx_fn ltid_y i@@ -184,8 +174,7 @@             css_init <- index "css_init" css_merge [ltid_y, ltid_x]              css <- forLoop tk [css_init] $ \k [acc_merge] ->-              resultBodyM-                =<< letTupExp' "foo"+              (resultBodyM <=< letTupExp' "foo")                 =<< eIf                   ( toExp $                       if epilogue

@@ -22,6 +22,7 @@ import Futhark.IR.SOACS qualified as Futhark import Futhark.IR.SOACS.Simplify (simplifyLambda) import Futhark.Optimise.Fusion.GraphRep+import Futhark.Optimise.Fusion.RulesWithAccs qualified as SF import Futhark.Optimise.Fusion.TryFusion qualified as TF import Futhark.Pass import Futhark.Transform.Rename@@ -78,10 +79,6 @@   modify (\s -> s {fuseScans = fs})   pure r -unreachableEitherDir :: DepGraph -> G.Node -> G.Node -> Bool-unreachableEitherDir g a b =-  not (reachable g a b || reachable g b a)- isNotVarInput :: [H.Input] -> [H.Input] isNotVarInput = filter (isNothing . H.isVarInput) @@ -114,14 +111,6 @@   modify $ \s -> s {fusionCount = 1 + fusionCount s}   pure $ Just x -vFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool-vFusionFeasability dg@DepGraph {dgGraph = g} n1 n2 =-  not (any isInf (edgesBetween dg n1 n2))-    && not (any (reachable dg n2) (filter (/= n2) (G.pre g n1)))--hFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool-hFusionFeasability = unreachableEitherDir- vTryFuseNodesInGraph :: G.Node -> G.Node -> DepGraphAug FusionM -- find the neighbors -> verify that fusion causes no cycles -> fuse vTryFuseNodesInGraph node_1 node_2 dg@DepGraph {dgGraph = g}@@ -302,6 +291,115 @@         H.addTransform           (H.Index cs (fullSlice (H.inputType inp) [ds]))           inp+-- Case of fusing a screma with an WithAcc such as to (hopefully) perform+--   more fusion within the WithAcc. This would allow the withAcc to move in+--   the code (since up to now they mostly remain where they were introduced.)+-- We conservatively allow the fusion to fire---i.e., to move the soac inside+--   the withAcc---when the following are not part of withAcc's accumulators:+--    1. the in-dependencies of the soac and+--    2. the result of the soac+--  Note that the soac result is allowed to be part of the `infusible`+--    for as long as it is returned by the withAcc. If `infusible` is empty+--    then the extranous result will be simplified away.+vFuseNodeT+  _edges+  _infusible+  (SoacNode ots1 pat1 soac@(H.Screma _w _form _s_inps) aux1, is1, _os1)+  (StmNode (Let pat2 aux2 (WithAcc w_inps lam0)), _os2)+    | ots1 == mempty,+      wacc_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w_inps,+      soac_prod_nms <- map patElemName $ patElems pat1,+      soac_indep_nms <- map getName is1,+      all (`notNameIn` wacc_cons_nms) (soac_indep_nms ++ soac_prod_nms) =+        do+          lam <- fst <$> doFusionInLambda lam0+          bdy' <-+            runBodyBuilder $ inScopeOf lam $ do+              soac' <- H.toExp soac+              addStm $ Let pat1 aux1 soac'+              lam_res <- bodyBind $ lambdaBody lam+              let pat1_res = map (SubExpRes (Certs []) . Var) soac_prod_nms+              pure $ lam_res ++ pat1_res+          let lam_ret_tp = lambdaReturnType lam ++ map patElemType (patElems pat1)+              pat = Pat $ patElems pat2 ++ patElems pat1+          lam' <- renameLambda $ lam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}+          -- see if bringing the map inside the scatter has actually benefitted fusion+          (lam'', success) <- doFusionInLambda lam'+          if not success+            then pure Nothing+            else do+              -- `aux1` already appear in the moved SOAC stm; is there+              -- any need to add it to the enclosing withAcc stm as well?+              fusedSomething $ StmNode $ Let pat aux2 $ WithAcc w_inps lam''+--+-- The reverse of the case above, i.e., fusing a screma at the back of an+--   WithAcc such as to (hopefully) enable more fusion there.+-- This should be safe as long as the SOAC does not uses any of the+--   accumulator arrays produced by the withAcc.+-- We could not provide a test for this case, due to the very restrictive+--   way in which accumulators can be used at source level.+--+--+vFuseNodeT+  edges+  _infusible+  (StmNode (Let pat1 aux1 (WithAcc w_inps wlam0)), _is1, _os1)+  (SoacNode ots2 pat2 soac@(H.Screma _w _form _s_inps) aux2, _os2)+    | ots2 == mempty,+      n <- length (lambdaParams wlam0) `div` 2,+      pat1_acc_nms <- namesFromList $ take n $ map patElemName $ patElems pat1,+      -- not $ namesIntersect (freeIn soac) pat1_acc_nms+      all ((`notNameIn` pat1_acc_nms) . getName) edges = do+        let empty_aux = StmAux mempty mempty mempty+        wlam <- fst <$> doFusionInLambda wlam0+        bdy' <-+          runBodyBuilder $ inScopeOf wlam $ do+            -- adding stms of withacc's lambda+            wlam_res <- bodyBind $ lambdaBody wlam+            -- add copies of the non-accumulator results of withacc+            let other_pr1 = drop n $ zip (patElems pat1) wlam_res+            forM_ other_pr1 $ \(pat_elm, bdy_res) -> do+              let (nm, se, tp) = (patElemName pat_elm, resSubExp bdy_res, patElemType pat_elm)+                  aux = empty_aux {stmAuxCerts = resCerts bdy_res}+              addStm $ Let (Pat [PatElem nm tp]) aux $ BasicOp $ SubExp se+            -- add the soac stmt+            soac' <- H.toExp soac+            addStm $ Let pat2 aux2 soac'+            -- build the body result+            let pat2_res = map (SubExpRes (Certs []) . Var . patElemName) $ patElems pat2+            pure $ wlam_res ++ pat2_res+        let lam_ret_tp = lambdaReturnType wlam ++ map patElemType (patElems pat2)+            pat = Pat $ patElems pat1 ++ patElems pat2+        wlam' <- renameLambda $ wlam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}+        -- see if bringing the map inside the scatter has actually benefitted fusion+        (wlam'', success) <- doFusionInLambda wlam'+        if not success+          then pure Nothing+          else -- `aux2` already appear in the enclosed SOAC stm; is there+          -- any need to add it to the enclosing withAcc stm as well?+            fusedSomething $ StmNode $ Let pat aux1 $ WithAcc w_inps wlam''+-- the case of fusing two withaccs+vFuseNodeT+  _edges+  infusible+  (StmNode (Let pat1 aux1 (WithAcc w1_inps lam1)), is1, _os1)+  (StmNode (Let pat2 aux2 (WithAcc w2_inps lam2)), _os2)+    | wacc2_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w2_inps,+      wacc1_indep_nms <- map getName is1,+      all (`notNameIn` wacc2_cons_nms) wacc1_indep_nms = do+        -- \^ the other safety checks are done inside `tryFuseWithAccs`+        lam1' <- fst <$> doFusionInLambda lam1+        lam2' <- fst <$> doFusionInLambda lam2+        let stm1 = Let pat1 aux1 (WithAcc w1_inps lam1')+            stm2 = Let pat2 aux2 (WithAcc w2_inps lam2')+        mstm <- SF.tryFuseWithAccs infusible stm1 stm2+        case mstm of+          Just (Let pat aux (WithAcc w_inps wlam)) -> do+            (wlam', success) <- doFusionInLambda wlam+            let new_stm = Let pat aux (WithAcc w_inps wlam')+            if success then fusedSomething (StmNode new_stm) else pure Nothing+          _ -> error "Illegal result of tryFuseWithAccs called from vFuseNodeT."+-- vFuseNodeT _ _ _ _ = pure Nothing  resFromLambda :: Lambda rep -> Result@@ -367,18 +465,26 @@ removeUnusedOutputs = mapAcross removeUnusedOutputsFromContext  tryFuseNodeInGraph :: DepNode -> DepGraphAug FusionM+tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g}+  | not (G.gelem (nodeFromLNode node_to_fuse) g) = pure dg+-- \^ Node might have been fused away since. tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g} = do-  if G.gelem node_to_fuse_id g -- Node might have been fused away since.-    then applyAugs (map (vTryFuseNodesInGraph node_to_fuse_id) fuses_with) dg-    else pure dg+  spec_rule_res <- SF.ruleMFScat node_to_fuse dg+  -- \^ specialized fusion rules such as the one+  --   enabling map-flatten-scatter fusion+  case spec_rule_res of+    Just dg' -> pure dg'+    Nothing -> applyAugs (map (vTryFuseNodesInGraph node_to_fuse_id) fuses_with) dg   where-    fuses_with = map fst $ filter (isDep . snd) $ G.lpre g (nodeFromLNode node_to_fuse)     node_to_fuse_id = nodeFromLNode node_to_fuse+    relevant (n, InfDep _) = isWithAccNodeId n dg+    relevant (_, e) = isDep e+    fuses_with = map fst $ filter relevant $ G.lpre g node_to_fuse_id  doVerticalFusion :: DepGraphAug FusionM doVerticalFusion dg = applyAugs (map tryFuseNodeInGraph $ reverse $ filter relevant $ G.labNodes (dgGraph dg)) dg   where-    relevant (_, StmNode {}) = False+    relevant (_, n@(StmNode {})) = isWithAccNodeT n     relevant (_, ResNode {}) = False     relevant _ = True @@ -436,48 +542,52 @@     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 lam args vec))) -> doFuseScans $ do-    lam' <- doFusionLambda lam+    lam' <- fst <$> doFusionInLambda lam     pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP lam' args vec))), outgoing)   StmNode (Let pat aux (Op (Futhark.JVP lam args vec))) -> doFuseScans $ do-    lam' <- doFusionLambda lam+    lam' <- fst <$> doFusionInLambda lam     pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP lam' args vec))), outgoing)   StmNode (Let pat aux (WithAcc inputs lam)) -> doFuseScans $ do-    lam' <- doFusionLambda lam+    lam' <- fst <$> doFusionInLambda lam     pure (incoming, node, StmNode (Let pat aux (WithAcc inputs lam')), outgoing)   SoacNode ots pat soac aux -> do     let lam = H.lambda soac-    lam' <- localScope (scopeOf lam) $ case soac of+    lam' <- inScopeOf lam $ case soac of       H.Stream {} ->-        dontFuseScans $ doFusionLambda lam+        dontFuseScans $ fst <$> doFusionInLambda lam       _ ->-        doFuseScans $ doFusionLambda lam+        doFuseScans $ fst <$> doFusionInLambda lam     let nodeT' = SoacNode ots pat (H.setLambda lam' soac) aux     pure (incoming, node, nodeT', outgoing)   _ -> pure c   where     doFusionWithDelayed :: Body SOACS -> [(NodeT, [EdgeT])] -> FusionM (Body SOACS)-    doFusionWithDelayed (Body () stms res) extraNodes = localScope (scopeOf stms) $ do+    doFusionWithDelayed (Body () stms res) extraNodes = inScopeOf stms $ do       stm_node <- mapM (finalizeNode . fst) extraNodes       stms' <- fuseGraph (mkBody (mconcat stm_node <> stms) res)       pure $ Body () stms' res++doFusionInLambda :: Lambda SOACS -> FusionM (Lambda SOACS, Bool)+doFusionInLambda lam = do+  -- To clean up previous instances of fusion.+  lam' <- simplifyLambda lam+  prev_count <- gets fusionCount+  newbody <- inScopeOf lam' $ doFusionBody $ lambdaBody lam'+  aft_count <- gets fusionCount+  -- To clean up any inner fusion.+  lam'' <-+    (if prev_count /= aft_count then simplifyLambda else pure)+      lam' {lambdaBody = newbody}+  pure (lam'', prev_count /= aft_count)+  where     doFusionBody :: Body SOACS -> FusionM (Body SOACS)     doFusionBody body = do       stms' <- fuseGraph body       pure $ body {bodyStms = stms'}-    doFusionLambda :: Lambda SOACS -> FusionM (Lambda SOACS)-    doFusionLambda lam = do-      -- To clean up previous instances of fusion.-      lam' <- simplifyLambda lam-      prev_count <- gets fusionCount-      newbody <- localScope (scopeOf lam') $ doFusionBody $ lambdaBody lam'-      aft_count <- gets fusionCount-      -- To clean up any inner fusion.-      (if prev_count /= aft_count then simplifyLambda else pure)-        lam' {lambdaBody = newbody}  -- main fusion function. fuseGraph :: Body SOACS -> FusionM (Stms SOACS)-fuseGraph body = localScope (scopeOf (bodyStms body)) $ do+fuseGraph body = inScopeOf (bodyStms body) $ do   graph_not_fused <- mkDepGraph body   graph_fused <- doAllFusion graph_not_fused   linearizeGraph graph_fused

@@ -41,6 +41,10 @@     mkDepGraph,     mkDepGraphForFun,     pprg,+    isWithAccNodeT,+    isWithAccNodeId,+    vFusionFeasability,+    hFusionFeasability,   ) where @@ -431,3 +435,25 @@ isCons :: EdgeT -> Bool isCons (Cons _) = True isCons _ = False++-- | Is this a withAcc?+isWithAccNodeT :: NodeT -> Bool+isWithAccNodeT (StmNode (Let _ _ (WithAcc _ _))) = True+isWithAccNodeT _ = False++isWithAccNodeId :: G.Node -> DepGraph -> Bool+isWithAccNodeId node_id (DepGraph {dgGraph = g}) =+  let (_, _, nT, _) = G.context g node_id+   in isWithAccNodeT nT++unreachableEitherDir :: DepGraph -> G.Node -> G.Node -> Bool+unreachableEitherDir g a b =+  not (reachable g a b || reachable g b a)++vFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool+vFusionFeasability dg@DepGraph {dgGraph = g} n1 n2 =+  (isWithAccNodeId n2 dg || not (any isInf (edgesBetween dg n1 n2)))+    && not (any (reachable dg n2) (filter (/= n2) (G.pre g n1)))++hFusionFeasability :: DepGraph -> G.Node -> G.Node -> Bool+hFusionFeasability = unreachableEitherDir

@@ -0,0 +1,607 @@+{-# LANGUAGE Strict #-}++-- | This module consists of rules for fusion+--     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.+--+--    II. WithAcc-WithAcc fusion: two withaccs can be+--        fused as long as the common accumulators use+--        the same operator, and as long as the non-accumulator+--        input of an WithAcc is not used as an accumulator in+--        the other. This fusion opens the door for fusing+--        the SOACs appearing inside the WithAccs. This is+--        also intended to demonstrate that it is not so+--        important where exactly the WithAccs were originally+--        introduced in the code, it is more important that+--        they can be transformed by various optimizations passes.+module Futhark.Optimise.Fusion.RulesWithAccs+  ( ruleMFScat,+    tryFuseWithAccs,+  )+where++import Control.Monad+import Data.Graph.Inductive.Graph qualified as G+import Data.Map.Strict qualified as M+import Data.Maybe+import Futhark.Analysis.HORep.SOAC qualified as H+import Futhark.Construct+import Futhark.IR.SOACS hiding (SOAC (..))+import Futhark.IR.SOACS qualified as F+import Futhark.Optimise.Fusion.GraphRep+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Transform.Substitute++se0 :: SubExp+se0 = intConst Int64 0++se1 :: SubExp+se1 = intConst Int64 1++-------------------------------------+--- I. Map-Flatten-Scatter Fusion ---+-------------------------------------++-- helper data structures+type IotaInp = ((VName, LParam SOACS), (SubExp, SubExp, SubExp, IntType))+-- ^           ((array-name, lambda param), (len, start, stride, Int64))++type RshpInp = ((VName, LParam SOACS), (Shape, Shape, Type))+-- ^           ((array-name, lambda param), (flat-shape, unflat-shape, elem-type))++-- | Implements a specialized rule for fusing a pattern+--   formed by a map o flatten o scatter, i.e.,+--      let (inds,   vals) = map-nest f inps+--          (finds, fvals) = (flatten inds, flatten vals)+--          let res = scatter res0 finds fvals+--   where inds & vals have higher rank than finds & fvals.+ruleMFScat ::+  (HasScope SOACS m, MonadFreshNames m) =>+  DepNode ->+  DepGraph ->+  m (Maybe DepGraph)+ruleMFScat node_to_fuse dg@DepGraph {dgGraph = g}+  | soac_nodeT <- snd node_to_fuse,+    scat_node_id <- nodeFromLNode node_to_fuse,+    SoacNode node_out_trsfs scat_pat scat_soac scat_aux <- soac_nodeT,+    H.nullTransforms node_out_trsfs,+    -- \^ for simplicity we do not allow transforms on scatter's result.+    H.Scatter _len scat_inp scat_out scat_lam <- scat_soac,+    -- \^ get the scatter+    scat_trsfs <- map H.inputTransforms (H.inputs scat_soac),+    -- \^ get the transforms on the input+    any (/= mempty) scat_trsfs,+    scat_ctx <- G.context g scat_node_id,+    (out_deps, _, _, inp_deps) <- scat_ctx,+    cons_deps <- filter (isCons . fst) inp_deps,+    drct_deps <- filter (isDep . fst) inp_deps,+    cons_ctxs <- map (G.context g . snd) cons_deps,+    drct_ctxs <- map (G.context g . snd) drct_deps,+    _cons_nTs <- map getNodeTfromCtx cons_ctxs, -- not used!!+    drct_tups0 <- mapMaybe (pairUp (zip drct_ctxs (map fst drct_deps))) scat_inp,+    length drct_tups0 == length scat_inp,+    -- \^ checks that all direct dependencies are also array+    --   inputs to scatter+    (t1s, t2s) <- unzip drct_tups0,+    drct_tups <- zip t1s $ zip t2s (lambdaParams scat_lam),+    (ctxs_iots, drct_iots) <- unzip $ filter (isIota . snd . fst . snd) drct_tups,+    (ctxs_rshp, drct_rshp) <- unzip $ filter (not . isIota . snd . fst . snd) drct_tups,+    length drct_iots + length drct_rshp == length scat_inp,+    -- \^ direct dependencies are either flatten reshapes or iotas.+    rep_iotas <- mapMaybe getRepIota drct_iots,+    length rep_iotas == length drct_iots,+    rep_rshps_certs <- mapMaybe getRepRshpArr drct_rshp,+    (rep_rshps, certs_rshps) <- unzip rep_rshps_certs,+    -- \^ gather the representations for the iotas and reshapes, that use+    --   the helper types `IotaInp` and `RshpInp`+    not (null rep_rshps),+    -- \^ at least one flatten-reshaped array+    length rep_rshps == length drct_rshp,+    (_, (s1, s2, _)) : _ <- rep_rshps,+    all (\(_, (s1', s2', _)) -> s1 == s1' && s2 == s2') rep_rshps,+    -- \^ Check that all unflatten shape dimensions are the same,+    --   so that we can construct a map nest;++    -- check profitability, which is conservatively defined as all+    -- the reshaped and consumer arrays are used solely by the+    -- scatter AND all reshape dependencies originate in the same+    -- map.+    checkSafeAndProfitable dg scat_node_id ctxs_rshp cons_ctxs = do+      -- generate the withAcc statement+      let cons_patels_outs = zip (patElems scat_pat) scat_out+      wacc_stm <- mkWithAccStm rep_iotas rep_rshps cons_patels_outs scat_aux scat_lam+      let all_cert_rshp = mconcat certs_rshps+          aux = stmAux wacc_stm+          aux' = aux {stmAuxCerts = all_cert_rshp <> stmAuxCerts aux}+          wacc_stm' = wacc_stm {stmAux = aux'}+          -- get the input deps of iotas+          fiot acc (_, _, _, inp_deps_iot) =+            acc <> inp_deps_iot+          deps_of_iotas = foldl fiot mempty ctxs_iots+          --+          iota_nms = namesFromList $ map (fst . fst) rep_iotas+          inp_deps_wo_iotas = filter ((`notNameIn` iota_nms) . getName . fst) inp_deps+          -- generate a new node for the with-acc-stmt and its associated context:+          --   add the inp-deps of iotas but remove the iota themselves from deps.+          new_withacc_nT = StmNode wacc_stm'+          inp_deps' = inp_deps_wo_iotas <> deps_of_iotas+          new_withacc_ctx = (out_deps, scat_node_id, new_withacc_nT, inp_deps')+          -- construct the new WithAcc node/graph; do we need to use `fusedSomething` ??+          new_node = G.node' new_withacc_ctx+          dg' = dg {dgGraph = new_withacc_ctx G.& G.delNodes [new_node] g}+      -- result+      pure $ Just dg'+  where+    --+    getNodeTfromCtx (_, _, nT, _) = nT+    findCtxOf ctxes nm+      | [ctxe] <- filter (\x -> nm == getName (snd x)) ctxes =+          Just ctxe+    findCtxOf _ _ = Nothing+    pairUp :: [(DepContext, EdgeT)] -> H.Input -> Maybe (DepContext, (H.Input, NodeT))+    pairUp ctxes inp@(H.Input _arrtrsfs nm _tp)+      | Just (ctx@(_, _, nT, _), _) <- findCtxOf ctxes nm =+          Just (ctx, (inp, nT))+    pairUp _ _ = Nothing+    --+    isIota :: NodeT -> Bool+    isIota (StmNode (Let _ _ (BasicOp (Iota {})))) = True+    isIota _ = False+    --+    getRepIota :: ((H.Input, NodeT), LParam SOACS) -> Maybe IotaInp+    getRepIota ((H.Input iottrsf arr_nm _arr_tp, nt), farg)+      | mempty == iottrsf,+        StmNode (Let _ _ (BasicOp (Iota n x s Int64))) <- nt =+          Just ((arr_nm, farg), (n, x, s, Int64))+    getRepIota _ = Nothing+    --+    getRepRshpArr :: ((H.Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)+    getRepRshpArr ((H.Input outtrsf arr_nm arr_tp, _nt), farg)+      | rshp_trsfm H.:< other_trsfms <- H.viewf outtrsf,+        (H.Reshape c ReshapeArbitrary shp_flat) <- rshp_trsfm,+        other_trsfms == mempty,+        eltp <- paramDec farg,+        Just shp_flat' <- checkShp eltp shp_flat,+        Array _ptp shp_unflat _ <- arr_tp,+        Just shp_unflat' <- checkShp eltp shp_unflat,+        shapeRank shp_flat' == 1,+        shapeRank shp_flat' < shapeRank shp_unflat' =+          Just (((arr_nm, farg), (shp_flat', shp_unflat', eltp)), c)+    getRepRshpArr _ = Nothing+    --+    checkShp (Prim _) shp_arr = Just shp_arr+    checkShp (Array _ptp shp_elm _) shp_arr =+      let dims_elm = shapeDims shp_elm+          dims_arr = shapeDims shp_arr+          (m, n) = (length dims_elm, length dims_arr)+          shp' = Shape $ take (n - m) dims_arr+          dims_com = drop (n - m) dims_arr+       in if all (uncurry (==)) (zip dims_com dims_elm)+            then Just shp'+            else Nothing+    checkShp _ _ = Nothing+-- default fails:+ruleMFScat _ _ = pure Nothing++checkSafeAndProfitable :: DepGraph -> G.Node -> [DepContext] -> [DepContext] -> Bool+checkSafeAndProfitable dg scat_node_id ctxs_rshp@(_ : _) ctxs_cons =+  let all_deps = concatMap (\(x, _, _, _) -> x) $ ctxs_rshp ++ ctxs_cons+      prof1 = all (\(_, dep_id) -> dep_id == scat_node_id) all_deps+      -- \^ scatter is the sole target to all consume & unflatten-reshape deps+      (_, map_node_id, map_nT, _) = head ctxs_rshp+      prof2 = all (\(_, nid, _, _) -> nid == map_node_id) ctxs_rshp+      prof3 = isMap map_nT+      -- \^ all reshapes come from the same node, which is a map+      safe = vFusionFeasability dg map_node_id scat_node_id+   in safe && prof1 && prof2 && prof3+  where+    isMap nT+      | SoacNode out_trsfs _pat soac _ <- nT,+        H.Screma _ _ form <- soac,+        ScremaForm [] [] _ <- form =+          H.nullTransforms out_trsfs+    isMap _ = False+checkSafeAndProfitable _ _ _ _ = False++-- | produces the withAcc statement that constitutes the translation of+--   the scater o flatten o map composition in which the map inputs are+--   reshaped in the same way+mkWithAccStm ::+  (HasScope SOACS m, MonadFreshNames m) =>+  [IotaInp] ->+  [RshpInp] ->+  [(PatElem (LetDec SOACS), (Shape, Int, VName))] ->+  StmAux (ExpDec SOACS) ->+  Lambda SOACS ->+  m (Stm SOACS)+mkWithAccStm iota_inps rshp_inps cons_patels_outs scatter_aux scatter_lam+  -- iotas are assumed to operate on Int64 values+  -- ToDo: maybe simplify rshp_inps+  --       check that the unflat shape is the same across reshapes+  --       check that the rank of the unflatten shape is higher than the flatten+  | rshp_inp : _ <- rshp_inps,+    (_, (_, s_unflat, _)) <- rshp_inp,+    (_ : _) <- shapeDims s_unflat = do+      --+      (cert_params, acc_params) <- fmap unzip $+        forM cons_patels_outs $ \(patel, (shp, _, nm)) -> do+          cert_param <- newParam "acc_cert_p" $ Prim Unit+          let arr_tp = patElemType patel+              acc_tp = stripArray (shapeRank shp) arr_tp+          acc_param <-+            newParam (baseString nm) $+              Acc (paramName cert_param) shp [acc_tp] NoUniqueness+          pure (cert_param, acc_param)+      let cons_params_outs = zip acc_params $ map snd cons_patels_outs+      acc_bdy <- mkWithAccBdy s_unflat iota_inps rshp_inps cons_params_outs scatter_lam+      let withacc_lam =+            Lambda+              { lambdaParams = cert_params ++ acc_params,+                lambdaReturnType = map paramDec acc_params,+                lambdaBody = acc_bdy+              }+          withacc_inps = map (\(_, (shp, _, nm)) -> (shp, [nm], Nothing)) cons_patels_outs+          withacc_pat = Pat $ map fst cons_patels_outs+          stm =+            Let withacc_pat scatter_aux $+              WithAcc withacc_inps withacc_lam+      pure stm+mkWithAccStm _ _ _ _ _ =+  error "Unreachable case reached!"++-- | Wrapper function for constructing the body of the withAcc+--   translation of the scatter+mkWithAccBdy ::+  (HasScope SOACS m, MonadFreshNames m) =>+  Shape ->+  [IotaInp] ->+  [RshpInp] ->+  [(LParam SOACS, (Shape, Int, VName))] ->+  Lambda SOACS ->+  m (Body SOACS)+mkWithAccBdy shp iota_inps rshp_inps cons_params_outs scat_lam = do+  let cons_ps = map fst cons_params_outs+      scat_res_info = map snd cons_params_outs+      static_arg = (iota_inps, rshp_inps, scat_res_info, scat_lam)+      mkParam ((nm, _), (_, s, t)) = Param mempty nm (arrayOfShape t s)+      rshp_ps = map mkParam rshp_inps+  mkWithAccBdy' static_arg (shapeDims shp) [] [] rshp_ps cons_ps++-- | builds a body that essentially consists of a map-nest with accumulators,+--   i.e., one level for each level of the unflatten shape of scatter's reshaped+--   input arrays+mkWithAccBdy' ::+  (HasScope SOACS m, MonadFreshNames m) =>+  ([IotaInp], [RshpInp], [(Shape, Int, VName)], Lambda SOACS) ->+  [SubExp] ->+  [SubExp] ->+  [VName] ->+  [LParam SOACS] ->+  [LParam SOACS] ->+  m (Body SOACS)++-- | the base case below addapts the scatter's lambda+mkWithAccBdy' static_arg [] dims_rev iot_par_nms rshp_ps cons_ps = do+  let (iota_inps, rshp_inps, scat_res_info, scat_lam) = static_arg+      tp_int = Prim $ IntType Int64+  scope <- askScope+  runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do+    -- handle iota args+    let strides_rev = scanl (*) (pe64 se1) $ map pe64 dims_rev+        strides = tail $ reverse strides_rev+        prods = zipWith (*) (map le64 iot_par_nms) strides+        i_pe = sum prods+    i_norm <- letExp "iota_norm_arg" =<< toExp i_pe+    forM_ iota_inps $ \arg -> do+      let ((_, i_par), (_, b, s, _)) = arg+      i_new <- letExp "tmp" =<< toExp (pe64 b + le64 i_norm * pe64 s)+      letBind (Pat [PatElem (paramName i_par) tp_int]) $ BasicOp $ SubExp $ Var i_new+    -- handle rshp args+    let rshp_lam_args = map (snd . fst) rshp_inps+    forM_ (zip rshp_lam_args rshp_ps) $ \(old_par, new_par) -> do+      let pat = Pat [PatElem (paramName old_par) (paramDec old_par)]+      letBind pat $ BasicOp $ SubExp $ Var $ paramName new_par+    -- add the body of the scatter's lambda+    mapM_ addStm $ bodyStms $ lambdaBody scat_lam+    -- add the withAcc update statements+    let iv_ses = groupScatterResults' scat_res_info $ bodyResult $ lambdaBody scat_lam+    res_nms <-+      forM (zip cons_ps iv_ses) $ \(cons_p, (i_ses, v_se)) -> do+        -- i_ses is a list+        let f nm_in i_se =+              letExp (baseString nm_in) $ BasicOp $ UpdateAcc Safe nm_in [resSubExp i_se] [resSubExp v_se]+        foldM f (paramName cons_p) i_ses+    let lam_certs = foldMap resCerts $ bodyResult $ lambdaBody scat_lam+    pure $ map (SubExpRes lam_certs . Var) res_nms+-- \| the recursive case builds a call to a map soac.+mkWithAccBdy' static_arg (dim : dims) dims_rev iot_par_nms rshp_ps cons_ps = do+  scope <- askScope+  runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do+    iota_arr <- letExp "iota_arr" $ BasicOp $ Iota dim se0 se1 Int64+    iota_p <- newParam "iota_arg" $ Prim $ IntType Int64+    rshp_ps' <- forM (zip [0 .. length rshp_ps - 1] (map paramDec rshp_ps)) $+      \(i, arr_tp) ->+        newParam ("rshp_arg_" ++ show i) $ stripArray 1 arr_tp+    cons_ps' <- forM (zip [0 .. length cons_ps - 1] (map paramDec cons_ps)) $+      \(i, arr_tp) ->+        newParam ("acc_arg_" ++ show i) arr_tp+    map_lam_bdy <-+      mkWithAccBdy' static_arg dims (dim : dims_rev) (iot_par_nms ++ [paramName iota_p]) rshp_ps' cons_ps'+    let map_lam = Lambda (rshp_ps' ++ [iota_p] ++ cons_ps') (map paramDec cons_ps') map_lam_bdy+        map_inps = map paramName rshp_ps ++ [iota_arr] ++ map paramName cons_ps+        map_soac = F.Screma dim map_inps $ ScremaForm [] [] map_lam+    res_nms <- letTupExp "acc_res" $ Op map_soac+    pure $ map (subExpRes . Var) res_nms++---------------------------------------------------+--- II. WithAcc-WithAcc Fusion+---------------------------------------------------++-- | Local helper type that tuples together:+--   1.   the pattern element corresponding to one withacc input+--   2.   the withacc input+--   3-5  withacc's lambda corresponding acc-certificate param,+--           argument param and result name+type AccTup =+  ( [PatElem (LetDec SOACS)],+    WithAccInput SOACS,+    LParam SOACS,+    LParam SOACS,+    (VName, Certs)+  )++accTup1 :: AccTup -> [PatElem (LetDec SOACS)]+accTup1 (a, _, _, _, _) = a++accTup2 :: AccTup -> WithAccInput SOACS+accTup2 (_, a, _, _, _) = a++accTup3 :: AccTup -> LParam SOACS+accTup3 (_, _, a, _, _) = a++accTup4 :: AccTup -> LParam SOACS+accTup4 (_, _, _, a, _) = a++accTup5 :: AccTup -> (VName, Certs)+accTup5 (_, _, _, _, a) = a++-- | Simple case for fusing two withAccs (can be extended):+--    let (b1, ..., bm, x1, ..., xq) = withAcc a1 ... am lam1+--    let (d1, ..., dn, y1, ..., yp) = withAcc c1 ... cn lam2+-- Notation: `b1 ... bm` are the accumulator results of the+--     first withAcc and `d1, ..., dn` of the second withAcc.+--     `x1 ... xq` and `y1, ..., yp` are non-accumulator results.+-- Conservative conditions:+--   1. for any bi (i=1..m) either `bi IN {c1, ..., cm}` OR+--        `bi NOT-IN FV(lam2)`, i.e., perfect producer-consumer+--        relation on accums. Of course the binary-op should+--        be the same.+--   2. The `bs` that are also accumulated upon in lam2+--        do NOT belong to the `infusible` set (they are destroyed)+--   3. x1 ... xq do not overlap with c1 ... cn+-- Fusion will create one withacc that accumulates on the+--   union of `a1 ... am` and `c1 ... cn` and returns, in addition+--   to the accumulator arrays the union of regular variables+--   `x1 ... xq` and `y1, ..., yp`+tryFuseWithAccs ::+  (HasScope SOACS m, MonadFreshNames m) =>+  [VName] ->+  Stm SOACS ->+  Stm SOACS ->+  m (Maybe (Stm SOACS))+tryFuseWithAccs+  infusible+  (Let pat1 aux1 (WithAcc w_inps1 lam1))+  (Let pat2 aux2 (WithAcc w_inps2 lam2))+    | (pat1_els, pat2_els) <- (patElems pat1, patElems pat2),+      (acc_tup1, other_pr1) <- groupAccs pat1_els w_inps1 lam1,+      (acc_tup2, other_pr2) <- groupAccs pat2_els w_inps2 lam2,+      (tup_common, acc_tup1', acc_tup2') <-+        groupCommonAccs acc_tup1 acc_tup2,+      -- safety 0: make sure that the accs from acc_tup1' and+      --           acc_tup2' do not overlap+      pnms_1' <- map patElemName $ concatMap (\(nms, _, _, _, _) -> nms) acc_tup1',+      winp_2' <- concatMap (\(_, (_, nms, _), _, _, _) -> nms) acc_tup2',+      not $ namesIntersect (namesFromList pnms_1') (namesFromList winp_2'),+      -- safety 1: we have already determined the commons;+      --           now we also need to check NOT-IN FV(lam2)+      not $ namesIntersect (namesFromList pnms_1') (freeIn lam2),+      -- safety 2:+      -- bs <- map patElemName $ concatMap accTup1 acc_tup1,+      bs <- map patElemName $ concatMap (accTup1 . fst) tup_common,+      all (`notElem` infusible) bs,+      -- safety 3:+      cs <- namesFromList $ concatMap ((\(_, xs, _) -> xs) . accTup2) acc_tup2,+      all ((`notNameIn` cs) . patElemName . fst) other_pr1 = do+        let getCertPairs (t1, t2) = (paramName (accTup3 t2), paramName (accTup3 t1))+            tab_certs = M.fromList $ map getCertPairs tup_common+            lam2_bdy' = substituteNames tab_certs (lambdaBody lam2)+            rcrt_params = map (accTup3 . fst) tup_common ++ map accTup3 acc_tup1' ++ map accTup3 acc_tup2'+            racc_params = map (accTup4 . fst) tup_common ++ map accTup4 acc_tup1' ++ map accTup4 acc_tup2'+            (comm_res_nms, comm_res_certs2) = unzip $ map (accTup5 . snd) tup_common+            (_, comm_res_certs1) = unzip $ map (accTup5 . fst) tup_common+            com_res_certs = zipWith (\x y -> Certs (unCerts x ++ unCerts y)) comm_res_certs1 comm_res_certs2+            bdyres_certs = com_res_certs ++ map (snd . accTup5) (acc_tup1' ++ acc_tup2')+            bdyres_accse = map Var comm_res_nms ++ map (Var . fst . accTup5) (acc_tup1' ++ acc_tup2')+            bdy_res_accs = zipWith SubExpRes bdyres_certs bdyres_accse+            bdy_res_others = map snd $ other_pr1 ++ other_pr2+        scope <- askScope+        lam_bdy <-+          runBodyBuilder $ do+            localScope (scope <> scopeOfLParams (rcrt_params ++ racc_params)) $ do+              -- add the stms of lam1+              mapM_ addStm $ stmsToList $ bodyStms $ lambdaBody lam1+              -- add the copy stms for the common accumulator+              forM_ tup_common $ \(tup1, tup2) -> do+                let (lpar1, lpar2) = (accTup4 tup1, accTup4 tup2)+                    ((nm1, _), nm2, tp_acc) = (accTup5 tup1, paramName lpar2, paramDec lpar1)+                letBind (Pat [PatElem nm2 tp_acc]) $ BasicOp $ SubExp $ Var nm1+              -- add copy stms to bring in scope x1 ... xq+              forM_ other_pr1 $ \(pat_elm, bdy_res) -> do+                let (nm, se, tp) = (patElemName pat_elm, resSubExp bdy_res, patElemType pat_elm)+                certifying (resCerts bdy_res) $+                  letBind (Pat [PatElem nm tp]) $+                    BasicOp (SubExp se)+              -- add the statements of lam2 (in which the acc-certificates have been substituted)+              mapM_ addStm $ stmsToList $ bodyStms lam2_bdy'+              -- build the result of body+              pure $ bdy_res_accs ++ bdy_res_others+        let tp_res_other = map (patElemType . fst) (other_pr1 ++ other_pr2)+            res_lam =+              Lambda+                { lambdaParams = rcrt_params ++ racc_params,+                  lambdaBody = lam_bdy,+                  lambdaReturnType = map paramDec racc_params ++ tp_res_other+                }+        res_lam' <- renameLambda res_lam+        let res_pat =+              concatMap (accTup1 . snd) tup_common+                ++ concatMap accTup1 (acc_tup1' ++ acc_tup2')+                ++ map fst (other_pr1 ++ other_pr2)+            res_w_inps = map (accTup2 . fst) tup_common ++ map accTup2 (acc_tup1' ++ acc_tup2')+        res_w_inps' <- mapM renameLamInWAccInp res_w_inps+        let stm_res = Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'+        pure $ Just stm_res+    where+      -- local helpers:++      groupAccs ::+        [PatElem (LetDec SOACS)] ->+        [WithAccInput SOACS] ->+        Lambda SOACS ->+        ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])+      groupAccs pat_els wacc_inps wlam =+        let lam_params = lambdaParams wlam+            n = length lam_params+            (lam_par_crts, lam_par_accs) = splitAt (n `div` 2) lam_params+            lab_res_ses = bodyResult $ lambdaBody wlam+         in groupAccsHlp pat_els wacc_inps lam_par_crts lam_par_accs lab_res_ses+      groupAccsHlp ::+        [PatElem (LetDec SOACS)] ->+        [WithAccInput SOACS] ->+        [LParam SOACS] ->+        [LParam SOACS] ->+        [SubExpRes] ->+        ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])+      groupAccsHlp pat_els [] [] [] lam_res_ses+        | length pat_els == length lam_res_ses =+            ([], zip pat_els lam_res_ses)+      groupAccsHlp+        pat_els+        (winp@(_, inp, _) : wacc_inps)+        (par_crt : lam_par_crts)+        (par_acc : lam_par_accs)+        (res_se : lam_res_ses)+          | n <- length inp,+            (n <= length pat_els) && (n <= (1 + length lam_res_ses)),+            Var res_nm <- resSubExp res_se =+              let (pat_els_cur, pat_els') = splitAt n pat_els+                  (rec1, rec2) = groupAccsHlp pat_els' wacc_inps lam_par_crts lam_par_accs lam_res_ses+               in ((pat_els_cur, winp, par_crt, par_acc, (res_nm, resCerts res_se)) : rec1, rec2)+      groupAccsHlp _ _ _ _ _ =+        error "Unreachable case reached in groupAccsHlp!"+      --+      groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])+      groupCommonAccs [] tup_accs2 =+        ([], [], tup_accs2)+      groupCommonAccs (tup_acc1 : tup_accs1) tup_accs2+        | commons2 <- filter (matchingAccTup tup_acc1) tup_accs2,+          length commons2 <= 1 =+            let (rec1, rec2, rec3) =+                  groupCommonAccs tup_accs1 $+                    if null commons2+                      then tup_accs2+                      else filter (not . matchingAccTup tup_acc1) tup_accs2+             in if null commons2+                  then (rec1, tup_acc1 : rec2, rec3)+                  else ((tup_acc1, head commons2) : rec1, tup_accs1, rec3)+      groupCommonAccs _ _ =+        error "Unreachable case reached in groupCommonAccs!"+      renameLamInWAccInp (shp, inps, Just (lam, se)) = do+        lam' <- renameLambda lam+        pure (shp, inps, Just (lam', se))+      renameLamInWAccInp winp = pure winp+--+tryFuseWithAccs _ _ _ =+  pure Nothing++-------------------------------+--- simple helper functions ---+-------------------------------++equivLambda ::+  M.Map VName VName ->+  Lambda SOACS ->+  Lambda SOACS ->+  Bool+equivLambda stab lam1 lam2+  | (ps1, ps2) <- (lambdaParams lam1, lambdaParams lam2),+    (nms1, nms2) <- (map paramName ps1, map paramName ps2),+    map paramDec ps1 == map paramDec ps2,+    map paramAttrs ps1 == map paramAttrs ps2,+    lambdaReturnType lam1 == lambdaReturnType lam2,+    (bdy1, bdy2) <- (lambdaBody lam1, lambdaBody lam2),+    bodyDec bdy1 == bodyDec bdy2 =+      let insert tab (x, k) = M.insert k x tab+          stab' = foldl insert stab $ zip nms1 nms2+          fStm (vtab, False) _ = (vtab, False)+          fStm (vtab, True) (s1, s2) = equivStm vtab s1 s2+          (stab'', success) =+            foldl fStm (stab', True) $+              zip (stmsToList (bodyStms bdy1)) $+                stmsToList (bodyStms bdy2)+          sres2 = substInSEs stab'' $ map resSubExp $ bodyResult bdy2+       in success && map resSubExp (bodyResult bdy1) == sres2+equivLambda _ _ _ =+  False++equivStm ::+  M.Map VName VName ->+  Stm SOACS ->+  Stm SOACS ->+  (M.Map VName VName, Bool)+equivStm+  stab+  (Let pat1 aux1 (BasicOp (BinOp bop1 se11 se12)))+  (Let pat2 aux2 (BasicOp (BinOp bop2 se21 se22)))+    | [se11, se12] == substInSEs stab [se21, se22],+      (pels1, pels2) <- (patElems pat1, patElems pat2),+      map patElemDec pels1 == map patElemDec pels2,+      bop1 == bop2 && aux1 == aux2 =+        let stab_new =+              M.fromList $+                zip (map patElemName pels2) (map patElemName pels1)+         in (M.union stab_new stab, True)+-- To Be Continued ...+equivStm vtab _ _ = (vtab, False)++matchingAccTup :: AccTup -> AccTup -> Bool+matchingAccTup+  (pat_els1, (shp1, _winp_arrs1, mlam1), _, _, _)+  (_, (shp2, winp_arrs2, mlam2), _, _, _) =+    shapeDims shp1 == shapeDims shp2+      && map patElemName pat_els1 == winp_arrs2+      && case (mlam1, mlam2) of+        (Nothing, Nothing) -> True+        (Just (lam1, see1), Just (lam2, see2)) ->+          (see1 == see2) && equivLambda M.empty lam1 lam2+        _ -> False++substInSEs :: M.Map VName VName -> [SubExp] -> [SubExp]+substInSEs vtab = map substInSE+  where+    substInSE (Var x)+      | Just y <- M.lookup x vtab = Var y+    substInSE z = z

@@ -789,7 +789,8 @@              inner_body <-               runBodyBuilder $-                eBody [SOAC.toExp $ inner $ map (SOAC.identInput . paramIdent) ps]+                varsRes+                  <$> (letTupExp "x" <=< SOAC.toExp $ inner $ map (SOAC.identInput . paramIdent) ps)             let inner_fun =                   Lambda                     { lambdaParams = ps,

@@ -892,8 +892,14 @@         pure (Just (op_lam', nes'), op_lam_stms)     (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs)   let noteAcc = ST.noteAccTokens (zip (map paramName (lambdaParams lam)) inputs')-  (lam', lam_stms) <- simplifyLambdaWith noteAcc (isFalse True) usage lam+  (lam', lam_stms) <-+    consumeInput inputs' $+      simplifyLambdaWith noteAcc (isFalse True) usage lam   pure (WithAcc inputs' lam', mconcat inputs_stms <> lam_stms)+  where+    inputArrs (_, arrs, _) = arrs+    consumeInput =+      localVtable . flip (foldl' (flip ST.consume)) . concatMap inputArrs simplifyExp _ _ e = do   e' <- simplifyExpBase e   pure (e', mempty)

@@ -429,7 +429,7 @@          loopbody' <-           localScope (scopeOfFParams mergeparams') . runBodyBuilder $-            resultBody . map Var <$> tiledBody private' privstms'+            varsRes <$> tiledBody private' privstms'         accs' <-           letTupExp "tiled_inside_loop" $             Loop merge' (ForLoop i it bound) loopbody'@@ -711,7 +711,7 @@                 map (paramName . fst) merge               tile_args =                 ProcessTileArgs privstms red_comm red_lam map_lam tile accs (Var tile_id)-          resultBody . map Var <$> tilingProcessTile tiling tile_args+          varsRes <$> tilingProcessTile tiling tile_args        accs <- letTupExp "accs" $ Loop merge loopform loopbody @@ -906,8 +906,7 @@       -- updates its accumulator.       let tile_args =             ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles-      resultBody . map Var-        <$> processTile1D gid gtid kdim residual_input grid tile_args+      varsRes <$> processTile1D gid gtid kdim residual_input grid tile_args  tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp tiling1d dims_on_top gtid kdim w = do@@ -1172,13 +1171,7 @@        -- Now each thread performs a traversal of the tile and       -- updates its accumulator.-      resultBody . map Var-        <$> processTile2D-          gids-          gtids-          kdims-          tile_size-          tile_args+      varsRes <$> processTile2D gids gtids kdims tile_size tile_args  tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp) tiling2d dims_on_top (gtid_x, gtid_y) (kdim_x, kdim_y) w = do

@@ -62,9 +62,10 @@   loop_inits <- mapM (\merge_t -> newParam "merge" $ toDecl merge_t Unique) merge_ts    loop_body <--    runBodyBuilder . localScope (scopeOfLoopForm loop_form <> scopeOfFParams loop_inits) $-      body i $-        map paramName loop_inits+    insertStmsM $+      localScope (scopeOfLoopForm loop_form <> scopeOfFParams loop_inits) $+        body i $+          map paramName loop_inits    letTupExp "loop" $     Loop (zip loop_inits $ map Var merge) loop_form loop_body

@@ -213,14 +213,13 @@         branch_pat' =           Pat $ map (fmap (`arrayOfRow` w)) $ patElems branch_pat -        mkBranch branch = (renameBody =<<) $ do+        mkBranch branch = (renameBody =<<) $ runBodyBuilder $ do           let lam = Lambda params lam_ret branch-              res = varsRes $ patNames branch_pat'-              map_stm = Let branch_pat' aux $ Op $ Screma w arrs $ mapSOAC lam-          pure $ mkBody (oneStm map_stm) res+          addStm $ Let branch_pat' aux $ Op $ Screma w arrs $ mapSOAC lam+          pure $ varsRes $ patNames branch_pat' -    cases' <- mapM (traverse $ runBodyBuilder . mkBranch) cases-    defbody' <- runBodyBuilder $ mkBranch defbody+    cases' <- mapM (traverse mkBranch) cases+    defbody' <- mkBranch defbody     pure . Branch [0 .. patSize pat - 1] pat' cond cases' defbody' $       MatchDec ret' if_sort 

@@ -182,9 +182,9 @@      pAtom =       choice-        [ try $ inParens sep (mkTuple <$> (parseExp sep `sepBy` pComma)),+        [ try $ inParens sep (mkTuple <$> (parseExp sep `sepEndBy` pComma)),           inParens sep $ parseExp sep,-          inBraces sep (Record <$> (pField `sepBy` pComma)),+          inBraces sep (Record <$> (pField `sepEndBy` pComma)),           StringLit . T.pack <$> lexeme sep ("\"" *> manyTill charLiteral "\""),           Const <$> V.parseValue sep,           Call <$> parseFunc <*> pure []@@ -192,7 +192,7 @@      pPat =       choice-        [ inParens sep $ pVarName `sepBy` pComma,+        [ inParens sep $ pVarName `sepEndBy` pComma,           pure <$> pVarName         ] 

@@ -213,7 +213,7 @@           letwith (map paramName mapout_params) (Var i) $             map resSubExp map_res -      pure . mkBody mempty . concat $+      pure . concat $         [ scan_res',           varsRes scan_outarrs,           red_res',@@ -278,7 +278,7 @@         certifying cs . letSubExp "mapout_res" . BasicOp $           Update Unsafe (paramName p) (fullSlice (paramType p) slice) se -      mkBodyM mempty $ subExpsRes $ res' ++ mapout_res'+      pure $ subExpsRes $ res' ++ mapout_res'    letBind pat $ Loop merge loop_form loop_body transformSOAC pat (Scatter len ivs as lam) = do@@ -307,7 +307,7 @@                   Update Safe arr' (fullSlice arr_t $ map (DimFix . resSubExp) indexCur) valueCur          foldM saveInArray arr indexes'-      pure $ resultBody (map Var ress)+      pure $ varsRes ress   letBind pat $ Loop merge (ForLoop iter Int64 len) loopBody transformSOAC pat (Hist len imgs ops bucket_fun) = do   iter <- newVName "iter"@@ -361,7 +361,7 @@            pure $ varsRes hist' -    pure $ resultBody $ map Var $ concat hists_out''+    pure $ varsRes $ concat hists_out''    -- Wrap up the above into a for-loop.   letBind pat $ Loop merge (ForLoop iter Int64 len) loopBody@@ -380,7 +380,7 @@   m (AST.Lambda rep) transformLambda (Lambda params rettype body) = do   body' <--    runBodyBuilder $+    fmap fst . runBuilder $       localScope (scopeOfLParams params) $         transformBody body   pure $ Lambda params rettype body'

@@ -56,6 +56,7 @@ import Futhark.Util.Pretty hiding (apply) import Language.Futhark hiding (Shape, matchDims) import Language.Futhark qualified as F+import Language.Futhark.Interpreter.AD qualified as AD import Language.Futhark.Interpreter.Values hiding (Value) import Language.Futhark.Interpreter.Values qualified import Language.Futhark.Primitive (floatValue, intValue)@@ -263,6 +264,9 @@ asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v asInteger (ValuePrim (UnsignedValue v)) =   toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64)+asInteger (ValueAD d v)+  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v =+      P.valueIntegral v' asInteger v = error $ "Unexpectedly not an integer: " <> show v  asInt :: Value -> Int@@ -270,13 +274,17 @@  asSigned :: Value -> IntValue asSigned (ValuePrim (SignedValue v)) = v-asSigned v = error $ "Unexpected not a signed integer: " <> show v+asSigned (ValueAD d v)+  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v'+asSigned v = error $ "Unexpectedly not a signed integer: " <> show v  asInt64 :: Value -> Int64 asInt64 = fromIntegral . asInteger  asBool :: Value -> Bool asBool (ValuePrim (BoolValue x)) = x+asBool (ValueAD d v)+  | P.BoolValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v' asBool v = error $ "Unexpectedly not a boolean: " <> show v  lookupInEnv ::@@ -937,6 +945,12 @@         Just v' -> pure v'         Nothing -> match v cs' +zeroOfType :: PrimType -> Value+zeroOfType (Signed it) = ValuePrim $ SignedValue $ P.intValue it (0 :: Int)+zeroOfType (Unsigned it) = ValuePrim $ UnsignedValue $ P.intValue it (0 :: Int)+zeroOfType (FloatType ft) = ValuePrim $ FloatValue $ P.floatValue ft (0 :: Int)+zeroOfType Bool = ValuePrim $ BoolValue False+ eval :: Env -> Exp -> EvalM Value eval _ (Literal v _) = pure $ ValuePrim v eval env (Hole (Info t) loc) =@@ -1008,28 +1022,15 @@     Scalar (Prim (FloatType ft)) ->       pure $ ValuePrim $ FloatValue $ floatValue ft v     _ -> error $ "eval: nonsensical type for float literal: " <> prettyString t-eval env (Negate e _) = do-  ev <- eval env e-  ValuePrim <$> case ev of-    ValuePrim (SignedValue (Int8Value v)) -> pure $ SignedValue $ Int8Value (-v)-    ValuePrim (SignedValue (Int16Value v)) -> pure $ SignedValue $ Int16Value (-v)-    ValuePrim (SignedValue (Int32Value v)) -> pure $ SignedValue $ Int32Value (-v)-    ValuePrim (SignedValue (Int64Value v)) -> pure $ SignedValue $ Int64Value (-v)-    ValuePrim (UnsignedValue (Int8Value v)) -> pure $ UnsignedValue $ Int8Value (-v)-    ValuePrim (UnsignedValue (Int16Value v)) -> pure $ UnsignedValue $ Int16Value (-v)-    ValuePrim (UnsignedValue (Int32Value v)) -> pure $ UnsignedValue $ Int32Value (-v)-    ValuePrim (UnsignedValue (Int64Value v)) -> pure $ UnsignedValue $ Int64Value (-v)-    ValuePrim (FloatValue (Float16Value v)) -> pure $ FloatValue $ Float16Value (-v)-    ValuePrim (FloatValue (Float32Value v)) -> pure $ FloatValue $ Float32Value (-v)-    ValuePrim (FloatValue (Float64Value v)) -> pure $ FloatValue $ Float64Value (-v)-    _ -> error $ "Cannot negate " <> show ev-eval env (Not e _) = do-  ev <- eval env e-  ValuePrim <$> case ev of-    ValuePrim (BoolValue b) -> pure $ BoolValue $ not b-    ValuePrim (SignedValue iv) -> pure $ SignedValue $ P.doComplement iv-    ValuePrim (UnsignedValue iv) -> pure $ UnsignedValue $ P.doComplement iv-    _ -> error $ "Cannot logically negate " <> show ev+eval env (Negate e loc) =+  -- -x = 0-x+  case typeOf e of+    Scalar (Prim pt) -> do+      ev <- eval env e+      apply2 loc env intrinsicsMinus (zeroOfType pt) ev+    t -> error $ "Cannot negate expression of type  " <> prettyString t+eval env (Not e loc) =+  apply loc env intrinsicsNot =<< eval env e eval env (Update src is v loc) =   maybe oob pure     =<< writeArray <$> mapM (evalDimIndex env) is <*> eval env src <*> eval env v@@ -1276,44 +1277,44 @@     types = M.mapMaybeWithKey (const . tdef . baseString) intrinsics      sintOp f =-      [ (getS, putS, P.doBinOp (f Int8)),-        (getS, putS, P.doBinOp (f Int16)),-        (getS, putS, P.doBinOp (f Int32)),-        (getS, putS, P.doBinOp (f Int64))+      [ (getS, putS, P.doBinOp (f Int8), adBinOp $ AD.OpBin (f Int8)),+        (getS, putS, P.doBinOp (f Int16), adBinOp $ AD.OpBin (f Int16)),+        (getS, putS, P.doBinOp (f Int32), adBinOp $ AD.OpBin (f Int32)),+        (getS, putS, P.doBinOp (f Int64), adBinOp $ AD.OpBin (f Int64))       ]     uintOp f =-      [ (getU, putU, P.doBinOp (f Int8)),-        (getU, putU, P.doBinOp (f Int16)),-        (getU, putU, P.doBinOp (f Int32)),-        (getU, putU, P.doBinOp (f Int64))+      [ (getU, putU, P.doBinOp (f Int8), adBinOp $ AD.OpBin (f Int8)),+        (getU, putU, P.doBinOp (f Int16), adBinOp $ AD.OpBin (f Int16)),+        (getU, putU, P.doBinOp (f Int32), adBinOp $ AD.OpBin (f Int32)),+        (getU, putU, P.doBinOp (f Int64), adBinOp $ AD.OpBin (f Int64))       ]     intOp f = sintOp f ++ uintOp f     floatOp f =-      [ (getF, putF, P.doBinOp (f Float16)),-        (getF, putF, P.doBinOp (f Float32)),-        (getF, putF, P.doBinOp (f Float64))+      [ (getF, putF, P.doBinOp (f Float16), adBinOp $ AD.OpBin (f Float16)),+        (getF, putF, P.doBinOp (f Float32), adBinOp $ AD.OpBin (f Float32)),+        (getF, putF, P.doBinOp (f Float64), adBinOp $ AD.OpBin (f Float64))       ]     arithOp f g = Just $ bopDef $ intOp f ++ floatOp g -    flipCmps = map (\(f, g, h) -> (f, g, flip h))+    flipCmps = map (\(f, g, h, o) -> (f, g, flip h, flip o))     sintCmp f =-      [ (getS, Just . BoolValue, P.doCmpOp (f Int8)),-        (getS, Just . BoolValue, P.doCmpOp (f Int16)),-        (getS, Just . BoolValue, P.doCmpOp (f Int32)),-        (getS, Just . BoolValue, P.doCmpOp (f Int64))+      [ (getS, Just . BoolValue, P.doCmpOp (f Int8), adBinOp $ AD.OpCmp (f Int8)),+        (getS, Just . BoolValue, P.doCmpOp (f Int16), adBinOp $ AD.OpCmp (f Int16)),+        (getS, Just . BoolValue, P.doCmpOp (f Int32), adBinOp $ AD.OpCmp (f Int32)),+        (getS, Just . BoolValue, P.doCmpOp (f Int64), adBinOp $ AD.OpCmp (f Int64))       ]     uintCmp f =-      [ (getU, Just . BoolValue, P.doCmpOp (f Int8)),-        (getU, Just . BoolValue, P.doCmpOp (f Int16)),-        (getU, Just . BoolValue, P.doCmpOp (f Int32)),-        (getU, Just . BoolValue, P.doCmpOp (f Int64))+      [ (getU, Just . BoolValue, P.doCmpOp (f Int8), adBinOp $ AD.OpCmp (f Int8)),+        (getU, Just . BoolValue, P.doCmpOp (f Int16), adBinOp $ AD.OpCmp (f Int16)),+        (getU, Just . BoolValue, P.doCmpOp (f Int32), adBinOp $ AD.OpCmp (f Int32)),+        (getU, Just . BoolValue, P.doCmpOp (f Int64), adBinOp $ AD.OpCmp (f Int64))       ]     floatCmp f =-      [ (getF, Just . BoolValue, P.doCmpOp (f Float16)),-        (getF, Just . BoolValue, P.doCmpOp (f Float32)),-        (getF, Just . BoolValue, P.doCmpOp (f Float64))+      [ (getF, Just . BoolValue, P.doCmpOp (f Float16), adBinOp $ AD.OpCmp (f Float16)),+        (getF, Just . BoolValue, P.doCmpOp (f Float32), adBinOp $ AD.OpCmp (f Float32)),+        (getF, Just . BoolValue, P.doCmpOp (f Float64), adBinOp $ AD.OpCmp (f Float64))       ]-    boolCmp f = [(getB, Just . BoolValue, P.doCmpOp f)]+    boolCmp f = [(getB, Just . BoolValue, P.doCmpOp f, adBinOp $ AD.OpCmp f)]      getV (SignedValue x) = Just $ P.IntValue x     getV (UnsignedValue x) = Just $ P.IntValue x@@ -1344,6 +1345,17 @@     putB (P.BoolValue x) = Just $ BoolValue x     putB _ = Nothing +    getAD (ValuePrim v) = AD.Constant <$> getV v+    getAD (ValueAD d v) = Just $ AD.Variable d v+    getAD _ = Nothing+    putAD (AD.Variable d s) = ValueAD d s+    putAD (AD.Constant v) = ValuePrim $ putV v++    adToPrim v = putV $ AD.primitive v++    adBinOp op x y = AD.doOp op [x, y]+    adUnOp op x = AD.doOp op [x]+     fun1 f =       TermValue Nothing $ ValueFun $ \x -> f x @@ -1408,6 +1420,12 @@           | Just z <- msum $ map (`bopDef'` (x', y')) fs -> do               breakOnNaN [x', y'] z               pure $ ValuePrim z+        _+          | Just x' <- getAD x,+            Just y' <- getAD y,+            Just z <- msum $ map (`bopDefAD'` (x', y')) fs -> do+              breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z+              pure $ putAD z         _ ->           bad noLoc mempty . docText $             "Cannot apply operator to arguments"@@ -1416,10 +1434,11 @@               <+> dquotes (prettyValue y)               <> "."       where-        bopDef' (valf, retf, op) (x, y) = do+        bopDef' (valf, retf, op, _) (x, y) = do           x' <- valf x           y' <- valf y           retf =<< op x' y'+        bopDefAD' (_, _, _, dop) (x, y) = dop x y      unopDef fs = fun1 $ \x ->       case x of@@ -1427,17 +1446,23 @@           | Just r <- msum $ map (`unopDef'` x') fs -> do               breakOnNaN [x'] r               pure $ ValuePrim r+        _+          | Just x' <- getAD x,+            Just r <- msum $ map (`unopDefAD'` x') fs -> do+              breakOnNaN [adToPrim x'] $ adToPrim r+              pure $ putAD r         _ ->           bad noLoc mempty . docText $             "Cannot apply function to argument"               <+> dquotes (prettyValue x)               <> "."       where-        unopDef' (valf, retf, op) x = do+        unopDef' (valf, retf, op, _) x = do           x' <- valf x           retf =<< op x'+        unopDefAD' (_, _, _, dop) = dop -    tbopDef f = fun1 $ \v ->+    tbopDef op f = fun1 $ \v ->       case fromTuple v of         Just [ValuePrim x, ValuePrim y]           | Just x' <- getV x,@@ -1445,6 +1470,12 @@             Just z <- putV <$> f x' y' -> do               breakOnNaN [x, y] z               pure $ ValuePrim z+        Just [x, y]+          | Just x' <- getAD x,+            Just y' <- getAD y,+            Just z <- AD.doOp op [x', y'] -> do+              breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z+              pure $ putAD z         _ ->           bad noLoc mempty . docText $             "Cannot apply operator to argument"@@ -1454,15 +1485,15 @@     def "!" =       Just $         unopDef-          [ (getS, putS, P.doUnOp $ P.Complement Int8),-            (getS, putS, P.doUnOp $ P.Complement Int16),-            (getS, putS, P.doUnOp $ P.Complement Int32),-            (getS, putS, P.doUnOp $ P.Complement Int64),-            (getU, putU, P.doUnOp $ P.Complement Int8),-            (getU, putU, P.doUnOp $ P.Complement Int16),-            (getU, putU, P.doUnOp $ P.Complement Int32),-            (getU, putU, P.doUnOp $ P.Complement Int64),-            (getB, putB, P.doUnOp P.Not)+          [ (getS, putS, P.doUnOp $ P.Complement Int8, adUnOp $ AD.OpUn $ P.Complement Int8),+            (getS, putS, P.doUnOp $ P.Complement Int16, adUnOp $ AD.OpUn $ P.Complement Int16),+            (getS, putS, P.doUnOp $ P.Complement Int32, adUnOp $ AD.OpUn $ P.Complement Int32),+            (getS, putS, P.doUnOp $ P.Complement Int64, adUnOp $ AD.OpUn $ P.Complement Int64),+            (getU, putU, P.doUnOp $ P.Complement Int8, adUnOp $ AD.OpUn $ P.Complement Int8),+            (getU, putU, P.doUnOp $ P.Complement Int16, adUnOp $ AD.OpUn $ P.Complement Int16),+            (getU, putU, P.doUnOp $ P.Complement Int32, adUnOp $ AD.OpUn $ P.Complement Int32),+            (getU, putU, P.doUnOp $ P.Complement Int64, adUnOp $ AD.OpUn $ P.Complement Int64),+            (getB, putB, P.doUnOp P.Not, adUnOp $ AD.OpUn P.Not)           ]     def "+" = arithOp (`P.Add` P.OverflowWrap) P.FAdd     def "-" = arithOp (`P.Sub` P.OverflowWrap) P.FSub@@ -1542,16 +1573,16 @@               ++ boolCmp P.CmpLle     def s       | Just bop <- find ((s ==) . prettyString) P.allBinOps =-          Just $ tbopDef $ P.doBinOp bop+          Just $ tbopDef (AD.OpBin bop) $ P.doBinOp bop       | Just unop <- find ((s ==) . prettyString) P.allCmpOps =-          Just $ tbopDef $ \x y -> P.BoolValue <$> P.doCmpOp unop x y+          Just $ tbopDef (AD.OpCmp unop) $ \x y -> P.BoolValue <$> P.doCmpOp unop x y       | Just cop <- find ((s ==) . prettyString) P.allConvOps =-          Just $ unopDef [(getV, Just . putV, P.doConvOp cop)]+          Just $ unopDef [(getV, Just . putV, P.doConvOp cop, adUnOp $ AD.OpConv cop)]       | Just unop <- find ((s ==) . prettyString) P.allUnOps =-          Just $ unopDef [(getV, Just . putV, P.doUnOp unop)]+          Just $ unopDef [(getV, Just . putV, P.doUnOp unop, adUnOp $ AD.OpUn unop)]       | Just (pts, _, f) <- M.lookup s P.primFuns =           case length pts of-            1 -> Just $ unopDef [(getV, Just . putV, f . pure)]+            1 -> Just $ unopDef [(getV, Just . putV, f . pure, adUnOp $ AD.OpFn s)]             _ -> Just $               fun1 $ \x -> do                 let getV' (ValuePrim v) = Just v@@ -1742,14 +1773,22 @@           ( ValueAcc shape op acc_arr,             ValuePrim (SignedValue (Int64Value i'))             ) ->-              if i' >= 0 && i' < arrayLength acc_arr-                then do-                  let x = acc_arr ! fromIntegral i'-                  res <- op x v-                  pure $ ValueAcc shape op $ acc_arr // [(fromIntegral i', res)]-                else pure acc+              write acc v shape op acc_arr i'+          ( ValueAcc shape op acc_arr,+            adv@(ValueAD {})+            )+              | Just (SignedValue (Int64Value i')) <- putV . AD.primitive <$> getAD adv ->+                  write acc v shape op acc_arr i'           _ ->             error $ "acc_write invalid arguments: " <> prettyString (show acc, show i, show v)+      where+        write acc v shape op acc_arr i' =+          if i' >= 0 && i' < arrayLength acc_arr+            then do+              let x = acc_arr ! fromIntegral i'+              res <- op x v+              pure $ ValueAcc shape op $ acc_arr // [(fromIntegral i', res)]+            else pure acc     --     def "flat_index_2d" = Just . fun6 $ \arr offset n1 s1 n2 s2 -> do       let offset' = asInt64 offset@@ -1926,11 +1965,129 @@           else pure $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) xs'     def "manifest" = Just $ fun1 pure     def "vjp2" = Just $-      fun3 $-        \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff."+      -- TODO: This could be much better. Currently, it is very inefficient+      -- Perhaps creating VJPValues could be abstracted into a function+      -- exposed by the AD module?+      fun3 $ \f v s -> do+        -- Get the depth+        depth <- length <$> stacktrace++        -- Augment the values+        let v' =+              fromMaybe (error $ "vjp: invalid values " ++ show v) $+                modifyValueM (\i lv -> ValueAD depth . AD.VJP . AD.VJPValue . AD.TapeID i <$> getAD lv) v+        -- Turn the seeds into a list of ADValues+        let s' =+              fromMaybe (error $ "vjp: invalid seeds " ++ show s) $+                mapM getAD $+                  fst $+                    valueAccum (\a b -> (b : a, b)) [] s++        -- Run the function, and turn its outputs into a list of Values+        o <- apply noLoc mempty f v'+        let o' = fst $ valueAccum (\a b -> (b : a, b)) [] o++        -- For each output..+        let m =+              fromMaybe (error "vjp: differentiation failed") $+                zipWithM+                  ( \on sn -> case on of+                      -- If it is a VJP variable of the correct depth, run deriveTape on it- and its corresponding seed+                      (ValueAD d (AD.VJP (AD.VJPValue t))) | d == depth -> (putAD $ AD.tapePrimal t,) <$> AD.deriveTape t sn+                      -- Otherwise, its partial derivatives are all 0+                      _ -> Just (on, M.empty)+                  )+                  o'+                  s'++        -- Add together every derivative+        let drvs = M.map (Just . putAD) $ M.unionsWith add $ map snd m++        -- Extract the output values, and the partial derivatives+        let ov = modifyValue (\i _ -> fst $ m !! i) o+        let od =+              fromMaybe (error "vjp: differentiation failed") $+                modifyValueM (\i vo -> M.findWithDefault (ValuePrim . putV . P.blankPrimValue . P.primValueType . AD.primitive <$> getAD vo) i drvs) v++        -- Return a tuple of the output values, and partial derivatives+        pure $ toTuple [ov, od]+      where+        modifyValue f v = snd $ valueAccum (\a b -> (a + 1, f a b)) 0 v+        modifyValueM f v =+          snd+            <$> valueAccumLM+              ( \a b -> do+                  b' <- f a b+                  pure (a + 1, b')+              )+              0+              v++        -- TODO: Perhaps this could be fully abstracted by AD?+        -- Making addFor private would be nice..+        add x y =+          fromMaybe (error "jvp: illtyped add") $+            AD.doOp (AD.OpBin $ AD.addFor $ P.primValueType $ AD.primitive x) [x, y]     def "jvp2" = Just $-      fun3 $-        \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff."+      -- TODO: This could be much better. Currently, it is very inefficient+      -- Perhaps creating JVPValues could be abstracted into a function+      -- exposed by the AD module?+      fun3 $ \f v s -> do+        -- Get the depth+        depth <- length <$> stacktrace++        -- Turn the seeds into a list of ADValues+        let s' =+              expectJust ("jvp: invalid seeds " ++ show s) $+                mapM getAD $+                  fst $+                    valueAccum (\a b -> (b : a, b)) [] s+        -- Augment the values+        let v' =+              expectJust ("jvp: invalid values " ++ show v) $+                modifyValueM+                  ( \i lv -> do+                      lv' <- getAD lv+                      pure $ ValueAD depth . AD.JVP . AD.JVPValue lv' $ s' !! (length s' - 1 - i)+                  )+                  v++        -- Run the function, and turn its outputs into a list of Values+        o <- apply noLoc mempty f v'+        let o' = fst $ valueAccum (\a b -> (b : a, b)) [] o++        -- For each output..+        let m =+              expectJust "jvp: differentiation failed" $+                mapM+                  ( \on -> case on of+                      -- If it is a JVP variable of the correct depth, return its primal and derivative+                      (ValueAD d (AD.JVP (AD.JVPValue pv dv))) | d == depth -> Just (putAD pv, putAD dv)+                      -- Otherwise, its partial derivatives are all 0+                      _ -> (on,) . ValuePrim . putV . P.blankPrimValue . P.primValueType . AD.primitive <$> getAD on+                  )+                  o'++        -- Extract the output values, and the partial derivatives+        let ov = modifyValue (\i _ -> fst $ m !! (length m - 1 - i)) o+            od = modifyValue (\i _ -> snd $ m !! (length m - 1 - i)) o++        -- Return a tuple of the output values, and partial derivatives+        pure $ toTuple [ov, od]+      where+        modifyValue f v = snd $ valueAccum (\a b -> (a + 1, f a b)) 0 v+        modifyValueM f v =+          snd+            <$> valueAccumLM+              ( \a b -> do+                  b' <- f a b+                  pure (a + 1, b')+              )+              0+              v++        expectJust _ (Just v) = v+        expectJust s Nothing = error s     def "acc" = Nothing     def s | nameFromString s `M.member` namesToPrimTypes = Nothing     def s = error $ "Missing intrinsic: " ++ s@@ -1943,6 +2100,18 @@       let n = ValuePrim $ SignedValue $ Int64Value $ arrayLength xs        in apply2 noLoc mempty f n arg     stream _ arg = error $ "Cannot stream: " <> show arg++intrinsicVal :: Name -> Value+intrinsicVal name =+  case M.lookup (intrinsicVar name) $ envTerm $ ctxEnv initialCtx of+    Just (TermValue _ v) -> v+    _ -> error $ "intrinsicVal: " <> prettyString name++intrinsicsMinus :: Value+intrinsicsMinus = intrinsicVal "-"++intrinsicsNot :: Value+intrinsicsNot = intrinsicVal "!"  interpretExp :: Ctx -> Exp -> F ExtOp Value interpretExp ctx e = runEvalM (ctxImports ctx) $ eval (ctxEnv ctx) e

@@ -0,0 +1,320 @@+module Language.Futhark.Interpreter.AD+  ( Op (..),+    ADVariable (..),+    ADValue (..),+    Tape (..),+    VJPValue (..),+    JVPValue (..),+    doOp,+    addFor,+    primal,+    tapePrimal,+    primitive,+    deriveTape,+  )+where++import Control.Monad (foldM, zipWithM)+import Data.Either (isRight)+import Data.List (find)+import Data.Map qualified as M+import Data.Maybe (fromMaybe)+import Futhark.AD.Derivatives (pdBinOp, pdBuiltin, pdUnOp)+import Futhark.Analysis.PrimExp (PrimExp (..))+import Language.Futhark.Core (VName (..), nameFromString)+import Language.Futhark.Primitive++-- Mathematical operations subject to AD.+data Op+  = OpBin BinOp+  | OpCmp CmpOp+  | OpUn UnOp+  | OpFn String+  | OpConv ConvOp+  deriving (Show)++-- Checks if an operation matches the types of its operands+opTypeMatch :: Op -> [PrimType] -> Bool+opTypeMatch (OpBin op) p = all (\x -> binOpType op == x) p+opTypeMatch (OpCmp op) p = all (\x -> cmpOpType op == x) p+opTypeMatch (OpUn op) p = all (\x -> unOpType op == x) p+opTypeMatch (OpConv op) p = all (\x -> fst (convOpType op) == x) p+opTypeMatch (OpFn fn) p = case M.lookup fn primFuns of+  Just (t, _, _) -> and $ zipWith (==) t p+  Nothing -> error "opTypeMatch" -- It is assumed that the function exists++-- Gets the return type of an operation+opReturnType :: Op -> PrimType+opReturnType (OpBin op) = binOpType op+opReturnType (OpCmp op) = cmpOpType op+opReturnType (OpUn op) = unOpType op+opReturnType (OpConv op) = snd $ convOpType op+opReturnType (OpFn fn) = case M.lookup fn primFuns of+  Just (_, t, _) -> t+  Nothing -> error "opReturnType" -- It is assumed that the function exists++-- Returns the operation which performs addition (or an+-- equivalent operation) on the given type+addFor :: PrimType -> BinOp+addFor (IntType t) = Add t OverflowWrap+addFor (FloatType t) = FAdd t+addFor Bool = LogOr+addFor t = error $ "addFor: " ++ show t++-- Returns the function which performs multiplication+-- (or an equivalent operation) on the given type+mulFor :: PrimType -> BinOp+mulFor (IntType t) = Mul t OverflowWrap+mulFor (FloatType t) = FMul t+mulFor Bool = LogAnd+mulFor t = error $ "mulFor: " ++ show t++-- Types and utility functions--+-- When taking the partial derivative of a function, we+-- must differentiate between the values which are kept+-- constant, and those which are not+data ADValue+  = Variable Int ADVariable+  | Constant PrimValue+  deriving (Show)++-- When performing automatic differentiation, each derived+-- variable must be augmented with additional data. This+-- value holds the primitive value of the variable, as well+-- as its data+data ADVariable+  = VJP VJPValue+  | JVP JVPValue+  deriving (Show)++depth :: ADValue -> Int+depth (Variable d _) = d+depth (Constant _) = 0++primal :: ADValue -> ADValue+primal (Variable _ (VJP (VJPValue t))) = tapePrimal t+primal (Variable _ (JVP (JVPValue v _))) = primal v+primal (Constant v) = Constant v++primitive :: ADValue -> PrimValue+primitive v@(Variable _ _) = primitive $ primal v+primitive (Constant v) = v++-- Evaluates a PrimExp using doOp+evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Maybe ADValue+evalPrimExp m (LeafExp n _) = M.lookup n m+evalPrimExp _ (ValueExp pv) = Just $ Constant pv+evalPrimExp m (BinOpExp op x y) = do+  x' <- evalPrimExp m x+  y' <- evalPrimExp m y+  doOp (OpBin op) [x', y']+evalPrimExp m (CmpOpExp op x y) = do+  x' <- evalPrimExp m x+  y' <- evalPrimExp m y+  doOp (OpCmp op) [x', y']+evalPrimExp m (UnOpExp op x) = do+  x' <- evalPrimExp m x+  doOp (OpUn op) [x']+evalPrimExp m (ConvOpExp op x) = do+  x' <- evalPrimExp m x+  doOp (OpConv op) [x']+evalPrimExp m (FunExp fn p _) = do+  p' <- mapM (evalPrimExp m) p+  doOp (OpFn fn) p'++-- Returns a list of PrimExps calculating the partial+-- derivative of each operands of a given operation+lookupPDs :: Op -> [PrimExp VName] -> Maybe [PrimExp VName]+lookupPDs (OpBin op) [x, y] = Just $ do+  let (a, b) = pdBinOp op x y+  [a, b]+lookupPDs (OpUn op) [x] = Just [pdUnOp op x]+lookupPDs (OpFn fn) p = pdBuiltin (nameFromString fn) p+lookupPDs _ _ = Nothing++-- Shared AD logic--+-- This function performs a mathematical operation on a+-- list of operands, performing automatic differentiation+-- if one or more operands is a Variable (of depth > 0)+doOp :: Op -> [ADValue] -> Maybe ADValue+doOp op o+  | not $ opTypeMatch op (map primValueType pv) =+      -- This function may be called with arguments of invalid types,+      -- because it is used as part of an overloaded operator.+      Nothing+  | otherwise = do+      let dep = case op of+            OpCmp _ -> 0 -- AD is not well-defined for comparason operations+            -- There are no derivatives for those written in+            -- PrimExp (check lookupPDs)+            _ -> maximum (map depth o)+      if dep == 0 then constCase else nonconstCase dep+  where+    pv = map primitive o++    divideDepths :: Int -> ADValue -> Either ADValue ADVariable+    divideDepths _ v@(Constant {}) = Left v+    divideDepths d v@(Variable d' v') = if d' < d then Left v else Right v'++    -- TODO: There may be a more graceful way of+    -- doing this+    extractVJP :: Either ADValue ADVariable -> Either ADValue VJPValue+    extractVJP (Right (VJP v)) = Right v+    extractVJP (Left v) = Left v+    extractVJP _ =+      -- This will never be called when the maximum depth layer is JVP+      error "extractVJP"++    -- TODO: There may be a more graceful way of+    -- doing this+    extractJVP :: Either ADValue ADVariable -> Either ADValue JVPValue+    extractJVP (Right (JVP v)) = Right v+    extractJVP (Left v) = Left v+    extractJVP _ =+      -- This will never be called when the maximum depth layer is VJP+      error "extractJVP"++    -- In this case, every operand is a constant, and the+    -- mathematical operation can be applied as it would be+    -- otherwise+    constCase =+      Constant <$> case (op, pv) of+        (OpBin op', [x, y]) -> doBinOp op' x y+        (OpCmp op', [x, y]) -> BoolValue <$> doCmpOp op' x y+        (OpUn op', [x]) -> doUnOp op' x+        (OpConv op', [x]) -> doConvOp op' x+        (OpFn fn, _) -> do+          (_, _, f) <- M.lookup fn primFuns+          f pv+        _ -> error "doOp: opTypeMatch"++    nonconstCase dep = do+      -- In this case, some values are variables. We therefore+      -- have to perform the necessary steps for AD++      -- First, we calculate the value for the previous depth+      let oprev = map primal o+      vprev <- doOp op oprev++      -- Then we separate the values of the maximum depth from+      -- those of a lower depth+      let o' = map (divideDepths dep) o+      -- Then we find out what type of AD is being performed+      case find isRight o' of+        -- Finally, we perform the necessary steps for the given+        -- type of AD+        Just (Right (VJP {})) ->+          Just . Variable dep . VJP . VJPValue $ vjpHandleOp op (map extractVJP o') vprev+        Just (Right (JVP {})) ->+          Variable dep . JVP . JVPValue vprev <$> jvpHandleFn op (map extractJVP o')+        _ ->+          -- Since the maximum depth is non-zero, there must be at+          -- least one variable of depth > 0+          error "find isRight"++calculatePDs :: Op -> [ADValue] -> Maybe [ADValue]+calculatePDs op p = do+  -- Create a unique VName for each operand+  let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length p]+  -- Put the operands in the environment+  let m = M.fromList $ zip n p++  -- Look up, and calculate the partial derivative+  -- of the operation with respect to each operand+  pde <- lookupPDs op $ map (`LeafExp` opReturnType op) n+  mapM (evalPrimExp m) pde++-- VJP / Reverse mode automatic differentiation--+-- In reverse mode AD, the entire computation+-- leading up to a variable must be saved+-- This is represented as a Tape+newtype VJPValue = VJPValue Tape+  deriving (Show)++-- | Represents a computation tree, as well as every intermediate+-- value in its evaluation. TODO: make this a graph.+data Tape+  = -- | This represents a variable. Each variable is given a unique ID,+    -- and has an initial value+    TapeID Int ADValue+  | -- | This represents a constant.+    TapeConst ADValue+  | -- | This represents the application of a mathematical operation.+    -- Each parameter is given by its Tape, and the return value of+    -- the operation is saved+    TapeOp Op [Tape] ADValue+  deriving (Show)++-- | Returns the primal value of a Tape.+tapePrimal :: Tape -> ADValue+tapePrimal (TapeID _ v) = v+tapePrimal (TapeConst v) = v+tapePrimal (TapeOp _ _ v) = v++-- This updates Tape of a VJPValue with a new operation,+-- treating all operands of a lower depth as constants+vjpHandleOp :: Op -> [Either ADValue VJPValue] -> ADValue -> Tape+vjpHandleOp op p v = do+  TapeOp op (map toTape p) v+  where+    toTape (Left v') = TapeConst v'+    toTape (Right (VJPValue t)) = t++-- | This calculates every partial derivative of a 'Tape'. The result+-- is a map of the partial derivatives, each key corresponding to the+-- ID of a free variable (see TapeID).+deriveTape :: Tape -> ADValue -> Maybe (M.Map Int ADValue)+deriveTape (TapeID i _) s = Just $ M.fromList [(i, s)]+deriveTape (TapeConst _) _ = Just M.empty+deriveTape (TapeOp op p _) s = do+  -- Calculate the new sensitivities+  s'' <- case op of+    OpConv op' -> do+      -- In case of type conversion, simply convert the sensitivity+      s' <- doOp (OpConv $ flipConvOp op') [s]+      Just [s']+    _ -> do+      pds <- calculatePDs op $ map tapePrimal p+      mapM (mul s) pds++  -- Propagate the new sensitivities+  pd <- zipWithM deriveTape p s''+  -- Add up the results+  Just $ foldl (M.unionWith add) M.empty pd+  where+    add x y =+      fromMaybe (error "deriveTape: addition failed") $+        doOp (OpBin $ addFor $ opReturnType op) [x, y]+    mul x y = doOp (OpBin $ mulFor $ opReturnType op) [x, y]++-- JVP / Forward mode automatic differentiation--++-- | In JVP, the derivative of the variable must be saved. This is+-- represented as a second value.+data JVPValue = JVPValue ADValue ADValue+  deriving (Show)++-- | This calculates the derivative part of the JVPValue resulting+-- from the application of a mathematical operation on one or more+-- JVPValues.+jvpHandleFn :: Op -> [Either ADValue JVPValue] -> Maybe ADValue+jvpHandleFn op p = do+  case op of+    OpConv _ ->+      -- In case of type conversion, simply convert+      -- the old derivative+      doOp op [derivative $ head p]+    _ -> do+      -- Calculate the new derivative using the chain+      -- rule+      pds <- calculatePDs op $ map primal' p+      vs <- zipWithM mul pds $ map derivative p+      foldM add (Constant $ blankPrimValue $ opReturnType op) vs+  where+    primal' (Left v) = v+    primal' (Right (JVPValue v _)) = v+    derivative (Left v) = Constant $ blankPrimValue $ primValueType $ primitive v+    derivative (Right (JVPValue _ d)) = d++    add x y = doOp (OpBin $ addFor $ opReturnType op) [x, y]+    mul x y = doOp (OpBin $ mulFor $ opReturnType op) [x, y]

@@ -14,6 +14,8 @@     valueShape,     prettyValue,     valueText,+    valueAccum,+    valueAccumLM,     fromTuple,     arrayLength,     isEmptyArray,@@ -28,6 +30,7 @@ where  import Data.Array+import Data.Bifunctor (Bifunctor (second)) import Data.List (genericLength) import Data.Map qualified as M import Data.Maybe@@ -35,9 +38,10 @@ import Data.Text qualified as T import Data.Vector.Storable qualified as SVec import Futhark.Data qualified as V-import Futhark.Util (chunk)+import Futhark.Util (chunk, mapAccumLM) import Futhark.Util.Pretty import Language.Futhark hiding (Shape, matchDims)+import Language.Futhark.Interpreter.AD qualified as AD import Language.Futhark.Primitive qualified as P import Prelude hiding (break, mod) @@ -106,6 +110,8 @@     ValueSum ValueShape Name [Value m]   | -- The shape, the update function, and the array.     ValueAcc ValueShape (Value m -> Value m -> m (Value m)) !(Array Int (Value m))+  | -- A primitive value with added information used in automatic differentiation+    ValueAD Int AD.ADVariable  instance Show (Value m) where   show (ValuePrim v) = "ValuePrim " <> show v <> ""@@ -114,6 +120,7 @@   show (ValueSum shape c vs) = unwords ["ValueSum", "(" <> show shape <> ")", show c, "(" <> show vs <> ")"]   show ValueFun {} = "ValueFun _"   show ValueAcc {} = "ValueAcc _"+  show (ValueAD d v) = unwords ["ValueAD", show d, show v]  instance Eq (Value m) where   ValuePrim (SignedValue x) == ValuePrim (SignedValue y) =@@ -145,6 +152,8 @@     pprPrec _ ValueAcc {} = "#<acc>"     pprPrec p (ValueSum _ n vs) =       parensIf (p > (0 :: Int)) $ "#" <> sep (pretty n : map (pprPrec 1) vs)+    -- TODO: This could be prettier. Perhaps add pretty printing for ADVariable / ADValues+    pprPrec _ (ValueAD d v) = pretty $ "d[" ++ show d ++ "]" ++ show v     pprElem v@ValueArray {} = pprPrec 0 v     pprElem v = group $ pprPrec 0 v @@ -181,6 +190,40 @@ valueShape (ValueRecord fs) = ShapeRecord $ M.map valueShape fs valueShape (ValueSum shape _ _) = shape valueShape _ = ShapeLeaf++-- TODO: Perhaps there is some clever way to reuse the code between+-- valueAccum and valueAccumLM+valueAccum :: (a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m)+valueAccum f i v@(ValuePrim {}) = f i v+valueAccum f i v@(ValueAD {}) = f i v+valueAccum f i (ValueRecord m) = second ValueRecord $ M.mapAccum (valueAccum f) i m+valueAccum f i (ValueArray s a) = do+  -- TODO: This could probably be better+  -- Transform into a map+  let m = M.fromList $ assocs a+  -- Accumulate over the map+  let (i', m') = M.mapAccum (valueAccum f) i m+  -- Transform back into an array and return+  let a' = array (bounds a) (M.toList m')+  (i', ValueArray s a')+valueAccum _ _ v = error $ "valueAccum not implemented for " ++ show v++valueAccumLM :: (Monad f) => (a -> Value m -> f (a, Value m)) -> a -> Value m -> f (a, Value m)+valueAccumLM f i v@(ValuePrim {}) = f i v+valueAccumLM f i v@(ValueAD {}) = f i v+valueAccumLM f i (ValueRecord m) = do+  (a, b) <- mapAccumLM (valueAccumLM f) i m+  pure (a, ValueRecord b)+valueAccumLM f i (ValueArray s a) = do+  -- TODO: This could probably be better+  -- Transform into a map+  let m = M.fromList $ assocs a+  -- Accumulate over the map+  (i', m') <- mapAccumLM (valueAccumLM f) i m+  -- Transform back into an array and return+  let a' = array (bounds a) (M.toList m')+  pure (i', ValueArray s a')+valueAccumLM _ _ v = error $ "valueAccum not implemented for " ++ show v  -- | Does the value correspond to an empty array? isEmptyArray :: Value m -> Bool