diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -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
 
diff --git a/docs/language-reference.rst b/docs/language-reference.rst
--- a/docs/language-reference.rst
+++ b/docs/language-reference.rst
@@ -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
 ~~~~~~~~~~~~~~
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -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
diff --git a/rts/c/atomics.h b/rts/c/atomics.h
--- a/rts/c/atomics.h
+++ b/rts/c/atomics.h
@@ -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
 }
 
diff --git a/rts/c/tuning.h b/rts/c/tuning.h
--- a/rts/c/tuning.h
+++ b/rts/c/tuning.h
@@ -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.
 }
diff --git a/rts/c/values.h b/rts/c/values.h
--- a/rts/c/values.h
+++ b/rts/c/values.h
@@ -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;
diff --git a/rts/python/values.py b/rts/python/values.py
--- a/rts/python/values.py
+++ b/rts/python/values.py
@@ -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
diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs
--- a/src/Futhark/Analysis/HORep/MapNest.hs
+++ b/src/Futhark/Analysis/HORep/MapNest.hs
@@ -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,
diff --git a/src/Futhark/Builder.hs b/src/Futhark/Builder.hs
--- a/src/Futhark/Builder.hs
+++ b/src/Futhark/Builder.hs
@@ -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.
diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU.hs
@@ -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) =>
diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs
--- a/src/Futhark/IR/SOACS/Simplify.hs
+++ b/src/Futhark/IR/SOACS/Simplify.hs
@@ -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,
diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs
--- a/src/Futhark/IR/Syntax.hs
+++ b/src/Futhark/IR/Syntax.hs
@@ -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.
diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs
--- a/src/Futhark/Internalise/Exps.hs
+++ b/src/Futhark/Internalise/Exps.hs
@@ -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)
diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs
--- a/src/Futhark/Internalise/Lambdas.hs
+++ b/src/Futhark/Internalise/Lambdas.hs
@@ -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
diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs
--- a/src/Futhark/Internalise/Monomorphise.hs
+++ b/src/Futhark/Internalise/Monomorphise.hs
@@ -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
diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs
--- a/src/Futhark/Optimise/BlkRegTiling.hs
+++ b/src/Futhark/Optimise/BlkRegTiling.hs
@@ -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
diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs
--- a/src/Futhark/Optimise/Fusion.hs
+++ b/src/Futhark/Optimise/Fusion.hs
@@ -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
diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs
--- a/src/Futhark/Optimise/Fusion/GraphRep.hs
+++ b/src/Futhark/Optimise/Fusion/GraphRep.hs
@@ -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
diff --git a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs
@@ -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
diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs
--- a/src/Futhark/Optimise/Fusion/TryFusion.hs
+++ b/src/Futhark/Optimise/Fusion/TryFusion.hs
@@ -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,
diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs
--- a/src/Futhark/Optimise/Simplify/Engine.hs
+++ b/src/Futhark/Optimise/Simplify/Engine.hs
@@ -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)
diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs
--- a/src/Futhark/Optimise/TileLoops.hs
+++ b/src/Futhark/Optimise/TileLoops.hs
@@ -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
diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs
--- a/src/Futhark/Optimise/TileLoops/Shared.hs
+++ b/src/Futhark/Optimise/TileLoops/Shared.hs
@@ -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
diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs
--- a/src/Futhark/Pass/ExtractKernels/Interchange.hs
+++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs
@@ -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
 
diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs
--- a/src/Futhark/Script.hs
+++ b/src/Futhark/Script.hs
@@ -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
         ]
 
diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs
--- a/src/Futhark/Transform/FirstOrderTransform.hs
+++ b/src/Futhark/Transform/FirstOrderTransform.hs
@@ -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'
diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs
--- a/src/Language/Futhark/Interpreter.hs
+++ b/src/Language/Futhark/Interpreter.hs
@@ -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
diff --git a/src/Language/Futhark/Interpreter/AD.hs b/src/Language/Futhark/Interpreter/AD.hs
new file mode 100644
--- /dev/null
+++ b/src/Language/Futhark/Interpreter/AD.hs
@@ -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]
diff --git a/src/Language/Futhark/Interpreter/Values.hs b/src/Language/Futhark/Interpreter/Values.hs
--- a/src/Language/Futhark/Interpreter/Values.hs
+++ b/src/Language/Futhark/Interpreter/Values.hs
@@ -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