diff --git a/docs/c-api.rst b/docs/c-api.rst
--- a/docs/c-api.rst
+++ b/docs/c-api.rst
@@ -8,6 +8,11 @@
 ``futlib.h``.  The API provided in the ``.h`` file is documented in
 the following.
 
+The ``.h`` file can be included by a C++ source file to access the
+functions (``extern "C"`` is added automatically), but the ``.c`` file
+must be compiled with a proper C compiler and the resulting object
+file linked with the rest of the program.
+
 Using the API requires creating a *configuration object*, which is
 then used to obtain a *context object*, which is then used to perform
 most other operations, such as calling Futhark functions.
diff --git a/docs/error-index.rst b/docs/error-index.rst
--- a/docs/error-index.rst
+++ b/docs/error-index.rst
@@ -582,32 +582,6 @@
 Module errors
 -------------
 
-.. _nested-entry:
-
-"Entry points may not be declared inside modules."
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This occurs when the program uses the ``entry`` keyword inside a
-module:
-
-.. code-block:: futhark
-
-  module m = {
-    entry f x = x + 1
-  }
-
-Entry points can only be declared at the top level of a file.  When we
-wish to make a function from inside a module available as an entry
-point, we must define a wrapper function:
-
-.. code-block:: futhark
-
-  module m = {
-    def f x = x + 1
-  }
-
-  entry f = m.f
-
 .. _module-is-parametric:
 
 "Module *x* is a parametric module
@@ -767,3 +741,123 @@
 .. code-block:: futhark
 
   def f (r : {x:i32}) = r with x = 0
+
+Entry points
+------------
+
+.. _nested-entry:
+
+"Entry points may not be declared inside modules."
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This occurs when the program uses the ``entry`` keyword inside a
+module:
+
+.. code-block:: futhark
+
+  module m = {
+    entry f x = x + 1
+  }
+
+Entry points can only be declared at the top level of a file.  When we
+wish to make a function from inside a module available as an entry
+point, we must define a wrapper function:
+
+.. code-block:: futhark
+
+  module m = {
+    def f x = x + 1
+  }
+
+  entry f = m.f
+
+.. _polymorphic-entry:
+
+"Entry point functions may not be polymorphic."
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Entry points are Futhark functions that can be called from other
+languages, and are therefore limited how advanced their types can be.
+In this case, the problem is that an entry point may not have a
+polymorphic type, for example:
+
+.. code-block:: futhark
+
+   entry dup 't (x: t) : (t,t) = x
+
+This is an invalid entry point because it uses a type parameter
+``'t``.  This error occurs frequently when we want to test a
+polymorphic function.  In such cases, the solution is to define one or
+more *monomorphic* entry points, each for a distinct type.  For
+example, to we can define a variety of monomorphic entry points that
+call the built-in function ``scan``:
+
+.. code-block:: futhark
+
+   entry scan_i32 (xs: []i32) = scan (+) 0 xs
+
+   entry scan_f32 (xs: []i32) = scan (*) 1 xs
+
+.. _higher-order-entry:
+
+"Entry point functions may not be higher-order."
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Entry points are Futhark functions that can be called from other
+languages, and are therefore limited how advanced their types can be.
+In this case, the problem is that an entry point may use functions as
+input or output.  For example:
+
+.. code-block:: futhark
+
+   entry apply (f: i32 -> i32) (x: i32) = f x
+
+There is no simple workaround for such cases.  One option is to
+manually `defunctionalise
+<https://en.wikipedia.org/wiki/Defunctionalization>`_ to use a
+non-functional encoding of the functional values, but this can quickly
+get very elaborate.  Following up on the example above, if we know
+that the only functions that would ever be passed are ``(+y)`` or
+``(*y)`` for some ``y``, we could do something like the following:
+
+.. code-block:: futhark
+
+   type function = #add i32 | #mul i32
+
+   entry apply (f: function) (x: i32) =
+     match f
+     case #add y -> x + y
+     case #mul y -> x + y
+
+Although in many cases, the best solution is simply to define a
+handful of simpler entry points instead of a single complicated one.
+
+.. _size-polymorphic-entry:
+
+"Entry point functions must not be size-polymorphic in their return type."
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This somewhat rare error occurs when an entry point returns an array
+that can have an arbitrary size chosen by its caller.  Contrived example:
+
+.. code-block:: futhark
+
+   -- Entry point taking no parameters.
+   entry f [n] : [0][n]i32 = []
+
+The size ``n`` is chosen by the caller.  Note that the ``n`` might be
+inferred and invisible, as in this example:
+
+.. code-block:: futhark
+
+   entry g : [0][]i32 = []
+
+When calling functions within a Futhark program, size parameters are
+handled by type inference, but entry points are called from the
+outside world, which is not subject to type inference.  If you really
+must have entry points like this, turn the size parameter into an
+ordinary parameter:
+
+.. code-block:: futhark
+
+   entry f (n: i64) : [0][n]i32 = []
diff --git a/docs/installation.rst b/docs/installation.rst
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -38,7 +38,7 @@
 
 To compile Futhark you must first install an appropriate version of
 GHC, either with `ghcup <https://www.haskell.org/ghcup>`_ or a package
-manager.  Any version since GHC 8.10 should work.  You also need the
+manager.  Any version since GHC 9.0 should work.  You also need the
 ``cabal`` command line program, which ghcup will install for you as
 well.
 
@@ -96,6 +96,12 @@
 **Linux (x86_64)**
   `futhark-nightly-linux-x86_64.tar.xz <https://futhark-lang.org/releases/futhark-nightly-linux-x86_64.tar.xz>`_
 
+**macOS (x86_64)**
+  `futhark-nightly-macos-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-x86_64.zip>`_
+
+**Windows (x86_64)**
+  `futhark-nightly-windows-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-windows-x86_64.zip>`_
+
   You will still likely need to make a C compiler (such as GCC) available on your own.
 
 .. _linux-installation:
@@ -103,16 +109,13 @@
 Installing Futhark on Linux
 ---------------------------
 
-* `Linuxbrew`_ is a distribution-agnostic package manager that
-  contains a formula for Futhark.  If Linuxbrew is installed (which
-  does not require ``root`` access), installation is as easy as::
+* `Homebrew`_ is a distribution-agnostic package manager for macOS and
+  Linux that contains a formula for Futhark. If Homebrew is already
+  installed (which does not require ``root`` access), installation is
+  as easy as::
 
     $ brew install futhark
 
-  Note that as of this writing, Linuxbrew is hampered by limited
-  compute resources for building packages, so the Futhark version may
-  be a bit behind.
-
 * Arch Linux users can use a `futhark-nightly package
   <https://aur.archlinux.org/packages/futhark-nightly/>`_ or a
   `regular futhark package
@@ -179,13 +182,14 @@
 Setting up Futhark on Windows
 -----------------------------
 
-Due to limited maintenance and testing resources, Futhark is not
-directly supported on Windows.  Install `WSL
+Due to limited maintenance and testing resources, Futhark is only
+partially supported on Windows.  A precompiled nightly snapshot is
+available above.
+
+In most cases, it is better to install `WSL
 <https://docs.microsoft.com/en-us/windows/wsl/install>`_ and follow
 the Linux instructions above.  The C code generated by the Futhark
-compiler should work on Windows.
-
-In the future, we may support Windows directly again.
+compiler should work on Windows, except for the ``multicore`` backend.
 
 Futhark with Nix
 ----------------
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.22.4
+version:        0.22.5
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -150,6 +150,7 @@
       Futhark.AD.Fwd
       Futhark.AD.Rev
       Futhark.AD.Rev.Loop
+      Futhark.AD.Rev.Hist
       Futhark.AD.Rev.Map
       Futhark.AD.Rev.Monad
       Futhark.AD.Rev.Reduce
diff --git a/rts/c/cuda.h b/rts/c/cuda.h
--- a/rts/c/cuda.h
+++ b/rts/c/cuda.h
@@ -349,6 +349,9 @@
   } else {
     opts[i++] = strdup("--disable-warnings");
   }
+  opts[i++] = msgprintf("-D%s=%d",
+                        "max_group_size",
+                        (int)ctx->max_block_size);
   for (size_t j = 0; j < ctx->cfg.num_sizes; j++) {
     opts[i++] = msgprintf("-D%s=%zu", ctx->cfg.size_vars[j],
                           ctx->cfg.size_values[j]);
@@ -427,6 +430,7 @@
   }
   cache_hash(h, src, strlen(src));
   size_t ptxsize;
+  errno = 0;
   if (cache_restore(cache_fname, h, (unsigned char**)ptx, &ptxsize) != 0) {
     if (ctx->cfg.logging) {
       fprintf(stderr, "Failed to restore cache (errno: %s)\n", strerror(errno));
diff --git a/rts/c/opencl.h b/rts/c/opencl.h
--- a/rts/c/opencl.h
+++ b/rts/c/opencl.h
@@ -517,6 +517,11 @@
                    "-DLOCKSTEP_WIDTH=%d ",
                    (int)ctx->lockstep_width);
 
+  w += snprintf(compile_opts+w, compile_opts_size-w,
+                "-D%s=%d ",
+                "max_group_size",
+                (int)ctx->max_group_size);
+
   for (int i = 0; i < ctx->cfg.num_sizes; i++) {
     w += snprintf(compile_opts+w, compile_opts_size-w,
                   "-D%s=%d ",
@@ -742,6 +747,7 @@
 
       unsigned char *buf;
       size_t bufsize;
+      errno = 0;
       if (cache_restore(cache_fname, &h, &buf, &bufsize) != 0) {
         if (ctx->cfg.logging) {
           fprintf(stderr, "Failed to restore cache (errno: %s)\n", strerror(errno));
diff --git a/rts/c/server.h b/rts/c/server.h
--- a/rts/c/server.h
+++ b/rts/c/server.h
@@ -388,6 +388,7 @@
         return;
       }
 
+      errno = 0;
       if (t->restore(t->aux, f, s->ctx, value_ptr(&v->value)) != 0) {
         failure();
         printf("Failed to restore variable %s.\n"
@@ -852,11 +853,10 @@
   if (arr == NULL) {
     return 1;
   }
-  assert(futhark_context_sync(ctx) == 0);
-
+  int err = futhark_context_sync(ctx);
   *(void**)p = arr;
   free(data);
-  return 0;
+  return err;
 }
 
 void store_array(const struct array_aux *aux, FILE *f,
diff --git a/rts/python/opencl.py b/rts/python/opencl.py
--- a/rts/python/opencl.py
+++ b/rts/python/opencl.py
@@ -217,6 +217,8 @@
     if (len(program_src) >= 0):
         build_options += ["-DLOCKSTEP_WIDTH={}".format(lockstep_width)]
 
+        build_options += ["-D{}={}".format('max_group_size', max_group_size)]
+
         build_options += ["-D{}={}".format(s.
                                            replace('z', 'zz').
                                            replace('.', 'zi').
diff --git a/src/Futhark/AD/Rev/Hist.hs b/src/Futhark/AD/Rev/Hist.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/AD/Rev/Hist.hs
@@ -0,0 +1,889 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Futhark.AD.Rev.Hist
+  ( diffMinMaxHist,
+    diffMulHist,
+    diffAddHist,
+    diffHist,
+  )
+where
+
+import Control.Monad
+import Futhark.AD.Rev.Monad
+import Futhark.Analysis.PrimExp.Convert
+import Futhark.Builder
+import Futhark.IR.SOACS
+import Futhark.Tools
+import Futhark.Transform.Rename
+
+getBinOpPlus :: PrimType -> BinOp
+getBinOpPlus (IntType x) = Add x OverflowUndef
+getBinOpPlus (FloatType f) = FAdd f
+getBinOpPlus _ = error "In getBinOpMul, Hist.hs: input not supported"
+
+getBinOpDiv :: PrimType -> BinOp
+getBinOpDiv (IntType t) = SDiv t Unsafe
+getBinOpDiv (FloatType t) = FDiv t
+getBinOpDiv _ = error "In getBinOpDiv, Hist.hs: input not supported"
+
+withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName
+withinBounds [] = TPrimExp $ ValueExp (BoolValue True)
+withinBounds [(q, i)] = (le64 i .<. pe64 q) .&&. (pe64 (intConst Int64 (-1)) .<. le64 i)
+withinBounds (qi : qis) = withinBounds [qi] .&&. withinBounds qis
+
+elseIf :: PrimType -> [(ADM (Exp SOACS), ADM (Exp SOACS))] -> [ADM (Body SOACS)] -> ADM (Exp SOACS)
+elseIf t [(c1, c2)] [bt, bf] =
+  eIf
+    (eCmpOp (CmpEq t) c1 c2)
+    bt
+    bf
+elseIf t ((c1, c2) : cs) (bt : bs) =
+  eIf
+    (eCmpOp (CmpEq t) c1 c2)
+    bt
+    $ eBody
+    $ pure
+    $ elseIf t cs bs
+elseIf _ _ _ = error "In elseIf, Hist.hs: input not supported"
+
+bindSubExpRes :: String -> [SubExpRes] -> ADM [VName]
+bindSubExpRes s =
+  traverse
+    ( \(SubExpRes cs se) -> do
+        bn <- newVName s
+        certifying cs $ letBindNames [bn] $ BasicOp $ SubExp se
+        pure bn
+    )
+
+nestedmap :: [SubExp] -> [PrimType] -> Lambda SOACS -> ADM (Lambda SOACS)
+nestedmap [] _ lam = pure lam
+nestedmap s@(h : r) pt lam = do
+  params <- traverse (\tp -> newParam "x" $ Array tp (Shape s) NoUniqueness) pt
+  body <- nestedmap r pt lam
+  mkLambda params $
+    fmap varsRes . letTupExp "res" . Op $
+      Screma h (map paramName params) (mapSOAC body)
+
+-- \ds hs -> map2 lam ds hs
+mkF' :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], [VName], Lambda SOACS)
+mkF' lam tps n = do
+  lam' <- renameLambda lam
+
+  ds_params <- traverse (newParam "ds_param") tps
+  hs_params <- traverse (newParam "hs_param") tps
+  let ds_pars = fmap paramName ds_params
+  let hs_pars = fmap paramName hs_params
+  lam_map <-
+    mkLambda (ds_params <> hs_params) $
+      fmap varsRes . letTupExp "map_f'" . Op $
+        Screma n (ds_pars <> hs_pars) (mapSOAC lam')
+
+  pure (ds_pars, hs_pars, lam_map)
+
+-- \ls as rs -> map3 (\li ai ri -> li `lam` ai `lam` ri) ls as rs
+mkF :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], Lambda SOACS)
+mkF lam tps n = do
+  lam_l <- renameLambda lam
+  lam_r <- renameLambda lam
+  let q = length $ lambdaReturnType lam
+      (lps, aps) = splitAt q $ lambdaParams lam_l
+      (ips, rps) = splitAt q $ lambdaParams lam_r
+  lam' <- mkLambda (lps <> aps <> rps) $ do
+    lam_l_res <- bodyBind $ lambdaBody lam_l
+    forM_ (zip ips lam_l_res) $ \(ip, SubExpRes cs se) ->
+      certifying cs $ letBindNames [paramName ip] $ BasicOp $ SubExp se
+    bodyBind $ lambdaBody lam_r
+
+  ls_params <- traverse (newParam "ls_param") tps
+  as_params <- traverse (newParam "as_param") tps
+  rs_params <- traverse (newParam "rs_param") tps
+  let map_params = ls_params <> as_params <> rs_params
+  lam_map <-
+    mkLambda map_params $
+      fmap varsRes . letTupExp "map_f" $
+        Op $
+          Screma n (map paramName map_params) $
+            mapSOAC lam'
+
+  pure (map paramName as_params, lam_map)
+
+mapout :: VName -> SubExp -> SubExp -> ADM VName
+mapout is n w = do
+  par_is <- newParam "is" $ Prim int64
+  is'_lam <-
+    mkLambda [par_is] $
+      fmap varsRes . letTupExp "is'"
+        =<< eIf
+          (toExp $ withinBounds $ pure (w, paramName par_is))
+          (eBody $ pure $ eParam par_is)
+          (eBody $ pure $ eSubExp w)
+
+  letExp "is'" $ Op $ Screma n (pure is) $ mapSOAC is'_lam
+
+multiScatter :: SubExp -> [VName] -> VName -> [VName] -> ADM [VName]
+multiScatter n dst is vs = do
+  tps <- traverse lookupType vs
+  par_i <- newParam "i" $ Prim int64
+  scatter_params <- traverse (newParam "scatter_param" . rowType) tps
+  scatter_lam <-
+    mkLambda (par_i : scatter_params) $
+      fmap subExpsRes . mapM (letSubExp "scatter_map_res") =<< do
+        p1 <- replicateM (length scatter_params) $ eParam par_i
+        p2 <- traverse eParam scatter_params
+        pure $ p1 <> p2
+
+  letTupExp "scatter_res" . Op $
+    Scatter n (is : vs) scatter_lam $
+      zipWith (\t -> (,,) (Shape $ pure $ arraySize 0 t) 1) tps dst
+
+multiIndex :: [VName] -> [DimIndex SubExp] -> ADM [VName]
+multiIndex vs s = do
+  traverse
+    ( \x -> do
+        t <- lookupType x
+        letExp "sorted" $ BasicOp $ Index x (fullSlice t s)
+    )
+    vs
+
+--
+-- special case of histogram with min/max as operator.
+-- Original, assuming `is: [n]i64` and `dst: [w]btp`
+--     let x = reduce_by_index dst minmax ne is vs
+-- Forward sweep:
+--     need to copy dst: reverse sweep might use it 7
+--       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).
+--     let dst_cpy = copy dst
+--     let (x, x_inds) = zip vs (iota n)
+--                       |> reduce_by_index (dst_cpy,-1s) argminmax (ne,-1) is
+--
+-- Reverse sweep:
+--     dst_bar += map2 (\i b -> if i == -1
+--                              then b
+--                              else 0
+--                     ) x_inds x_bar
+
+--     vs_ctrbs = map2 (\i b -> if i == -1
+--                              then 0
+--                              else vs_bar[i] + b
+--                     ) x_inds x_bar
+--     vs_bar <- scatter vs_bar x_inds vs_ctrbs
+diffMinMaxHist ::
+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
+diffMinMaxHist _ops x aux n minmax ne is vs w rf dst m = do
+  let t = binOpType minmax
+  vs_type <- lookupType vs
+  let vs_elm_type = elemType vs_type
+  let vs_dims = arrayDims vs_type
+  let inner_dims = tail vs_dims
+  let nr_dims = length vs_dims
+  dst_type <- lookupType dst
+  let dst_dims = arrayDims dst_type
+
+  dst_cpy <- letExp (baseString dst <> "_copy") $ BasicOp $ Copy dst
+
+  acc_v_p <- newParam "acc_v" $ Prim t
+  acc_i_p <- newParam "acc_i" $ Prim int64
+  v_p <- newParam "v" $ Prim t
+  i_p <- newParam "i" $ Prim int64
+  hist_lam_inner <-
+    mkLambda [acc_v_p, acc_i_p, v_p, i_p] $
+      fmap varsRes . letTupExp "idx_res"
+        =<< eIf
+          (eCmpOp (CmpEq t) (eParam acc_v_p) (eParam v_p))
+          ( eBody
+              [ eParam acc_v_p,
+                eBinOp (SMin Int64) (eParam acc_i_p) (eParam i_p)
+              ]
+          )
+          ( eBody
+              [ eIf
+                  ( eCmpOp
+                      (CmpEq t)
+                      (eParam acc_v_p)
+                      (eBinOp minmax (eParam acc_v_p) (eParam v_p))
+                  )
+                  (eBody [eParam acc_v_p, eParam acc_i_p])
+                  (eBody [eParam v_p, eParam i_p])
+              ]
+          )
+  hist_lam <- nestedmap inner_dims [vs_elm_type, int64, vs_elm_type, int64] hist_lam_inner
+
+  dst_minus_ones <-
+    letExp "minus_ones" . BasicOp $
+      Replicate (Shape dst_dims) (intConst Int64 (-1))
+  ne_minus_ones <-
+    letSubExp "minus_ones" . BasicOp $
+      Replicate (Shape inner_dims) (intConst Int64 (-1))
+  iota_n <-
+    letExp "red_iota" . BasicOp $
+      Iota n (intConst Int64 0) (intConst Int64 1) Int64
+
+  inp_iota <- do
+    if nr_dims == 1
+      then pure iota_n
+      else do
+        i <- newParam "i" $ Prim int64
+        lam <-
+          mkLambda [i] $
+            fmap varsRes . letTupExp "res" =<< do
+              pure $ BasicOp $ Replicate (Shape inner_dims) $ Var $ paramName i
+
+        letExp "res" $ Op $ Screma n [iota_n] $ mapSOAC lam
+
+  let hist_op = HistOp (Shape [w]) rf [dst_cpy, dst_minus_ones] [ne, if nr_dims == 1 then intConst Int64 (-1) else ne_minus_ones] hist_lam
+  f' <- mkIdentityLambda [Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]
+  x_inds <- newVName (baseString x <> "_inds")
+  auxing aux $
+    letBindNames [x, x_inds] $
+      Op $
+        Hist n [is, vs, inp_iota] [hist_op] f'
+
+  m
+
+  x_bar <- lookupAdjVal x
+
+  x_ind_dst <- newParam (baseString x <> "_ind_param") $ Prim int64
+  x_bar_dst <- newParam (baseString x <> "_bar_param") $ Prim t
+  dst_lam_inner <-
+    mkLambda [x_ind_dst, x_bar_dst] $
+      fmap varsRes . letTupExp "dst_bar"
+        =<< eIf
+          (toExp $ le64 (paramName x_ind_dst) .==. -1)
+          (eBody $ pure $ eParam x_bar_dst)
+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
+  dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner
+
+  dst_bar <-
+    letExp (baseString dst <> "_bar") . Op $
+      Screma w [x_inds, x_bar] (mapSOAC dst_lam)
+
+  updateAdj dst dst_bar
+
+  vs_bar <- lookupAdjVal vs
+
+  inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []
+  let inds = x_inds : inds'
+
+  par_x_ind_vs <- replicateM nr_dims $ newParam (baseString x <> "_ind_param") $ Prim int64
+  par_x_bar_vs <- newParam (baseString x <> "_bar_param") $ Prim t
+  vs_lam_inner <-
+    mkLambda (par_x_bar_vs : par_x_ind_vs) $
+      fmap varsRes . letTupExp "res"
+        =<< eIf
+          (toExp $ le64 (paramName $ head par_x_ind_vs) .==. -1)
+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
+          ( eBody $
+              pure $ do
+                vs_bar_i <-
+                  letSubExp (baseString vs_bar <> "_el") . BasicOp $
+                    Index vs_bar . Slice $
+                      fmap (DimFix . Var . paramName) par_x_ind_vs
+                eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)
+          )
+  vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner
+
+  vs_bar_p <-
+    letExp (baseString vs <> "_partial") . Op $
+      Screma w (x_bar : inds) (mapSOAC vs_lam)
+
+  q <-
+    letSubExp "q"
+      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) dst_dims
+
+  scatter_inps <- do
+    -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]
+    -- ToDo: Cosmin asks: is the below the correct translation of the line above?
+    traverse (letExp "flat" . BasicOp . Reshape ReshapeArbitrary (Shape [q])) $
+      inds ++ [vs_bar_p]
+
+  f'' <- mkIdentityLambda $ replicate nr_dims (Prim int64) ++ [Prim t]
+  vs_bar' <-
+    letExp (baseString vs <> "_bar") . Op $
+      Scatter q scatter_inps f'' [(Shape vs_dims, 1, vs_bar)]
+  insAdj vs vs_bar'
+  where
+    mk_indices :: [SubExp] -> [SubExp] -> ADM [VName]
+    mk_indices [] _ = pure []
+    mk_indices [d] iotas = do
+      reps <- traverse (letExp "rep" . BasicOp . Replicate (Shape [d])) iotas
+      iota_d <-
+        letExp "red_iota" . BasicOp $
+          Iota d (intConst Int64 0) (intConst Int64 1) Int64
+      pure $ reps ++ [iota_d]
+    mk_indices (d : dims) iotas = do
+      iota_d <-
+        letExp "red_iota" . BasicOp $
+          Iota d (intConst Int64 0) (intConst Int64 1) Int64
+
+      i_param <- newParam "i" $ Prim int64
+      lam <-
+        mkLambda [i_param] $
+          fmap varsRes $
+            mk_indices dims $
+              iotas ++ [Var $ paramName i_param]
+
+      letTupExp "res" $ Op $ Screma d [iota_d] $ mapSOAC lam
+
+--
+-- special case of histogram with multiplication as operator.
+-- Original, assuming `is: [n]i64` and `dst: [w]btp`
+--     let x = reduce_by_index dst (*) ne is vs
+-- Forward sweep:
+--     dst does not need to be copied: dst is not overwritten
+--     let (ps, zs) = map (\v -> if v == 0 then (1,1) else (v,0)) vs
+--     let non_zero_prod = reduce_by_index nes (*) ne is ps
+--     let zero_count = reduce_by_index 0s (+) 0 is zs
+--     let h_part = map2 (\p c -> if c == 0 then p else 0
+--                       ) non_zero_prod zero_count
+--     let x = map2 (*) dst h_part
+--
+-- Reverse sweep:
+--     dst_bar += map2 (*) h_part x_bar
+
+--     let part_bar = map2 (*) dst x_bar
+--     vs_bar += map2 (\i v -> let zr_cts = zero_count[i]
+--                             let pr_bar = part_bar[i]
+--                             let nz_prd = non_zero_prod[i]
+--                             in if zr_cts == 0
+--                             then pr_bar * (nz_prd / v)
+--                             else if zr_cts == 1 and v == 0
+--                             then nz_prd * pr_bar
+--                             else 0
+--                    ) is vs
+diffMulHist ::
+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
+diffMulHist _ops x aux n mul ne is vs w rf dst m = do
+  let t = binOpType mul
+  vs_type <- lookupType vs
+  let vs_dims = arrayDims vs_type
+  let vs_elm_type = elemType vs_type
+  dst_type <- lookupType dst
+  let dst_dims = arrayDims dst_type
+  let inner_dims = tail vs_dims
+
+  v_param <- newParam "v" $ Prim t
+  lam_ps_zs_inner <-
+    mkLambda [v_param] $
+      fmap varsRes . letTupExp "map_res"
+        =<< eIf
+          (eCmpOp (CmpEq t) (eParam v_param) (eSubExp $ Constant $ blankPrimValue t))
+          (eBody $ fmap eSubExp [Constant $ onePrimValue t, intConst Int64 1])
+          (eBody [eParam v_param, eSubExp $ intConst Int64 0])
+  lam_ps_zs <- nestedmap vs_dims [vs_elm_type] lam_ps_zs_inner
+  ps_zs_res <- eLambda lam_ps_zs [eSubExp $ Var vs]
+  ps_zs <- bindSubExpRes "ps_zs" ps_zs_res
+  let [ps, zs] = ps_zs
+
+  lam_mul_inner <- binOpLambda mul t
+  lam_mul <- nestedmap inner_dims [vs_elm_type, vs_elm_type] lam_mul_inner
+  nz_prods0 <- letExp "nz_prd" $ BasicOp $ Replicate (Shape [w]) ne
+  let hist_nzp = HistOp (Shape [w]) rf [nz_prods0] [ne] lam_mul
+
+  lam_add_inner <- binOpLambda (Add Int64 OverflowUndef) int64
+  lam_add <- nestedmap inner_dims [int64, int64] lam_add_inner
+  zr_counts0 <- letExp "zr_cts" $ BasicOp $ Replicate (Shape dst_dims) (intConst Int64 0)
+  zrn_ne <- letSubExp "zr_ne" $ BasicOp $ Replicate (Shape inner_dims) (intConst Int64 0)
+  let hist_zrn = HistOp (Shape [w]) rf [zr_counts0] [if length vs_dims == 1 then intConst Int64 0 else zrn_ne] lam_add
+
+  f' <- mkIdentityLambda [Prim int64, Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]
+  nz_prods <- newVName "non_zero_prod"
+  zr_counts <- newVName "zero_count"
+  auxing aux $
+    letBindNames [nz_prods, zr_counts] $
+      Op $
+        Hist n [is, is, ps, zs] [hist_nzp, hist_zrn] f'
+
+  p_param <- newParam "prod" $ Prim t
+  c_param <- newParam "count" $ Prim int64
+  lam_h_part_inner <-
+    mkLambda [p_param, c_param] $
+      fmap varsRes . letTupExp "h_part"
+        =<< eIf
+          (toExp $ 0 .==. le64 (paramName c_param))
+          (eBody $ pure $ eParam p_param)
+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
+  lam_h_part <- nestedmap dst_dims [vs_elm_type, int64] lam_h_part_inner
+  h_part_res <- eLambda lam_h_part $ map (eSubExp . Var) [nz_prods, zr_counts]
+  h_part' <- bindSubExpRes "h_part" h_part_res
+  let [h_part] = h_part'
+
+  lam_mul_inner' <- binOpLambda mul t
+  lam_mul' <- nestedmap dst_dims [vs_elm_type, vs_elm_type] lam_mul_inner'
+  x_res <- eLambda lam_mul' $ map (eSubExp . Var) [dst, h_part]
+  x' <- bindSubExpRes "x" x_res
+  auxing aux $ letBindNames [x] $ BasicOp $ SubExp $ Var $ head x'
+
+  m
+
+  x_bar <- lookupAdjVal x
+
+  lam_mul'' <- renameLambda lam_mul'
+  dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]
+  dst_bar <- bindSubExpRes (baseString dst <> "_bar") dst_bar_res
+  updateAdj dst $ head dst_bar
+
+  lam_mul''' <- renameLambda lam_mul'
+  part_bar_res <- eLambda lam_mul''' $ map (eSubExp . Var) [dst, x_bar]
+  part_bar' <- bindSubExpRes "part_bar" part_bar_res
+  let [part_bar] = part_bar'
+
+  inner_params <- zipWithM newParam ["zr_cts", "pr_bar", "nz_prd", "a"] $ map Prim [int64, t, t, t]
+  let [zr_cts, pr_bar, nz_prd, a_param] = inner_params
+  lam_vsbar_inner <-
+    mkLambda inner_params $
+      fmap varsRes . letTupExp "vs_bar" =<< do
+        eIf
+          (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 0) (eParam zr_cts))
+          (eBody $ pure $ eBinOp mul (eParam pr_bar) $ eBinOp (getBinOpDiv t) (eParam nz_prd) $ eParam a_param)
+          ( eBody $
+              pure $
+                eIf
+                  ( eBinOp
+                      LogAnd
+                      (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 1) (eParam zr_cts))
+                      (eCmpOp (CmpEq t) (eSubExp $ Constant $ blankPrimValue t) $ eParam a_param)
+                  )
+                  (eBody $ pure $ eBinOp mul (eParam nz_prd) (eParam pr_bar))
+                  (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)
+          )
+
+  lam_vsbar_middle <- nestedmap inner_dims [int64, t, t, t] lam_vsbar_inner
+
+  i_param <- newParam "i" $ Prim int64
+  a_param' <- newParam "a" $ rowType vs_type
+  lam_vsbar <-
+    mkLambda [i_param, a_param'] $
+      fmap varsRes . letTupExp "vs_bar"
+        =<< eIf
+          (toExp $ withinBounds $ pure (w, paramName i_param))
+          ( buildBody_ $ do
+              let i = fullSlice vs_type [DimFix $ Var $ paramName i_param]
+              names <- traverse newVName ["zr_cts", "pr_bar", "nz_prd"]
+              zipWithM_ (\name -> letBindNames [name] . BasicOp . flip Index i) names [zr_counts, part_bar, nz_prods]
+              eLambda lam_vsbar_middle $ map (eSubExp . Var) names <> [eParam a_param']
+          )
+          (eBody $ pure $ pure $ zeroExp $ rowType dst_type)
+
+  vs_bar <-
+    letExp (baseString vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar
+
+  updateAdj vs vs_bar
+
+--
+-- special case of histogram with add as operator.
+-- Original, assuming `is: [n]i64` and `dst: [w]btp`
+--     let x = reduce_by_index dst (+) ne is vs
+-- Forward sweep:
+--     need to copy dst: reverse sweep might use it 7
+--       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).
+--     let dst_cpy = copy dst
+--     let x = reduce_by_index dst_cpy (+) ne is vs
+--
+-- Reverse sweep:
+--     dst_bar += x_bar
+--
+--     vs_bar += map (\i -> x_bar[i]) is
+diffAddHist ::
+  VjpOps -> VName -> StmAux () -> SubExp -> Lambda SOACS -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()
+diffAddHist _ops x aux n add ne is vs w rf dst m = do
+  let t = paramDec $ head $ lambdaParams add
+
+  dst_cpy <- letExp (baseString dst <> "_copy") $ BasicOp $ Copy dst
+
+  f <- mkIdentityLambda [Prim int64, t]
+  auxing aux . letBindNames [x] . Op $
+    Hist n [is, vs] [HistOp (Shape [w]) rf [dst_cpy] [ne] add] f
+
+  m
+
+  x_bar <- lookupAdjVal x
+
+  updateAdj dst x_bar
+
+  x_type <- lookupType x
+  i_param <- newParam (baseString vs <> "_i") $ Prim int64
+  let i = paramName i_param
+  lam_vsbar <-
+    mkLambda [i_param] $
+      fmap varsRes . letTupExp "vs_bar"
+        =<< eIf
+          (toExp $ withinBounds $ pure (w, i))
+          (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i])
+          (eBody $ pure $ eSubExp ne)
+
+  vs_bar <- letExp (baseString vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar
+  updateAdj vs vs_bar
+
+--
+-- a step in the radix sort implementation
+-- it assumes the key we are sorting
+-- after is [n]i64 and it is the first VName
+--
+-- local def radix_sort_step [n] 't (xs: [n]t) (get_bit: i32 -> t -> i32)
+--                                  (digit_n: i32): [n]t =
+--   let num x = get_bit (digit_n+1) x * 2 + get_bit digit_n x
+--   let pairwise op (a1,b1,c1,d1) (a2,b2,c2,d2) =
+--     (a1 `op` a2, b1 `op` b2, c1 `op` c2, d1 `op` d2)
+--   let bins = xs |> map num
+--   let flags = bins |> map (\x -> if x == 0 then (1,0,0,0)
+--                                  else if x == 1 then (0,1,0,0)
+--                                  else if x == 2 then (0,0,1,0)
+--                                  else (0,0,0,1))
+--   let offsets = scan (pairwise (+)) (0,0,0,0) flags
+--   let (na,nb,nc,_nd) = last offsets
+--   let f bin (a,b,c,d) = match bin
+--                         case 0 -> a-1
+--                         case 1 -> na+b-1
+--                         case 2 -> na+nb+c-1
+--                         case _ -> na+nb+nc+d-1
+--   let is = map2 f bins offsets
+--   in scatter scratch is xs
+radixSortStep :: [VName] -> [Type] -> SubExp -> SubExp -> SubExp -> ADM [VName]
+radixSortStep xs tps bit n w = do
+  -- let is = head xs
+  is <- mapout (head xs) n w
+
+  num_param <- newParam "num" $ Prim int64
+  num_lam <-
+    mkLambda [num_param] $
+      fmap varsRes . letTupExp "num_res"
+        =<< eBinOp
+          (Add Int64 OverflowUndef)
+          ( eBinOp
+              (And Int64)
+              (eBinOp (AShr Int64) (eParam num_param) (eSubExp bit))
+              (iConst 1)
+          )
+          ( eBinOp
+              (Mul Int64 OverflowUndef)
+              (iConst 2)
+              ( eBinOp
+                  (And Int64)
+                  (eBinOp (AShr Int64) (eParam num_param) (eBinOp (Add Int64 OverflowUndef) (eSubExp bit) (iConst 1)))
+                  (iConst 1)
+              )
+          )
+
+  bins <- letExp "bins" $ Op $ Screma n [is] $ mapSOAC num_lam
+  flag_param <- newParam "flag" $ Prim int64
+  flag_lam <-
+    mkLambda [flag_param] $
+      fmap varsRes . letTupExp "flag_res"
+        =<< elseIf
+          int64
+          (map ((,) (eParam flag_param) . iConst) [0 .. 2])
+          (map (eBody . fmap iConst . (\i -> map (\j -> if i == j then 1 else 0) [0 .. 3])) ([0 .. 3] :: [Integer]))
+
+  flags <- letTupExp "flags" $ Op $ Screma n [bins] $ mapSOAC flag_lam
+
+  scan_params <- traverse (flip newParam $ Prim int64) ["a1", "b1", "c1", "d1", "a2", "b2", "c2", "d2"]
+  scan_lam <-
+    mkLambda scan_params $
+      fmap subExpsRes . mapM (letSubExp "scan_res") =<< do
+        uncurry (zipWithM (eBinOp $ Add Int64 OverflowUndef)) $ splitAt 4 $ map eParam scan_params
+
+  scan <- scanSOAC $ pure $ Scan scan_lam $ map (intConst Int64) [0, 0, 0, 0]
+  offsets <- letTupExp "offsets" $ Op $ Screma n flags scan
+
+  ind <- letSubExp "ind_last" =<< eBinOp (Sub Int64 OverflowUndef) (eSubExp n) (iConst 1)
+  let i = Slice [DimFix ind]
+  nabcd <- traverse newVName ["na", "nb", "nc", "nd"]
+  zipWithM_ (\abcd -> letBindNames [abcd] . BasicOp . flip Index i) nabcd offsets
+
+  let vars = map Var nabcd
+  map_params <- traverse (flip newParam $ Prim int64) ["bin", "a", "b", "c", "d"]
+  map_lam <-
+    mkLambda map_params $
+      fmap varsRes . letTupExp "map_res"
+        =<< elseIf
+          int64
+          (map ((,) (eParam $ head map_params) . iConst) [0 .. 2])
+          ( zipWith
+              ( \j p ->
+                  eBody $
+                    pure $ do
+                      t <- letSubExp "t" =<< eBinOp (Sub Int64 OverflowUndef) (eParam p) (iConst 1)
+                      foldBinOp (Add Int64 OverflowUndef) (intConst Int64 0) (t : take j vars)
+              )
+              [0 .. 3]
+              (tail map_params)
+          )
+
+  nis <- letExp "nis" $ Op $ Screma n (bins : offsets) $ mapSOAC map_lam
+
+  scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) tps
+  multiScatter n scatter_dst nis xs
+  where
+    iConst c = eSubExp $ intConst Int64 c
+
+--
+-- the radix sort implementation
+-- def radix_sort [n] 't (xs: [n]i64) =
+--   let iters = if n == 0 then 0 else 32
+--   in loop xs for i < iters do radix_sort_step xs i64.get_bit (i*2)
+radixSort :: [VName] -> SubExp -> SubExp -> ADM [VName]
+radixSort xs n w = do
+  logw <- log2 =<< letSubExp "w1" =<< toExp (pe64 w + 1)
+  -- ceil logw by (logw + 1) / 2
+  iters <- letSubExp "iters" =<< toExp (untyped (pe64 logw + 1) ~/~ untyped (pe64 (intConst Int64 2)))
+
+  types <- traverse lookupType xs
+  params <- zipWithM (\x -> newParam (baseString x) . flip toDecl Nonunique) xs types
+  i <- newVName "i"
+  loopbody <- buildBody_ . localScope (scopeOfFParams params) $
+    fmap varsRes $ do
+      bit <- letSubExp "bit" =<< toExp (le64 i * 2)
+      radixSortStep (map paramName params) types bit n w
+
+  letTupExp "sorted" $
+    DoLoop
+      (zip params $ map Var xs)
+      (ForLoop i Int64 iters [])
+      loopbody
+  where
+    log2 :: SubExp -> ADM SubExp
+    log2 m = do
+      params <- zipWithM newParam ["cond", "r", "i"] $ map Prim [Bool, int64, int64]
+      let [cond, r, i] = params
+
+      body <- buildBody_ . localScope (scopeOfFParams params) $ do
+        r' <- letSubExp "r'" =<< toExp (le64 (paramName r) .>>. 1)
+        cond' <- letSubExp "cond'" =<< toExp (bNot $ pe64 r' .==. 0)
+        i' <- letSubExp "i'" =<< toExp (le64 (paramName i) + 1)
+        pure $ subExpsRes [cond', r', i']
+
+      cond_init <- letSubExp "test" =<< toExp (bNot $ pe64 m .==. 0)
+
+      l <-
+        letTupExp' "log2res" $
+          DoLoop
+            (zip params [cond_init, m, Constant $ blankPrimValue int64])
+            (WhileLoop $ paramName cond)
+            body
+
+      let [_, _, res] = l
+      pure res
+
+radixSort' :: [VName] -> SubExp -> SubExp -> ADM [VName]
+radixSort' xs n w = do
+  iota_n <-
+    letExp "red_iota" . BasicOp $
+      Iota n (intConst Int64 0) (intConst Int64 1) Int64
+
+  radres <- radixSort [head xs, iota_n] n w
+  let [is', iota'] = radres
+
+  i_param <- newParam "i" $ Prim int64
+  let slice = [DimFix $ Var $ paramName i_param]
+  map_lam <- mkLambda [i_param] $ varsRes <$> multiIndex (tail xs) slice
+
+  sorted <- letTupExp "sorted" $ Op $ Screma n [iota'] $ mapSOAC map_lam
+  pure $ iota' : is' : sorted
+
+--
+-- generic case of histogram.
+-- Original, assuming `is: [n]i64` and `dst: [w]btp`
+--   let xs = reduce_by_index dst odot ne is as
+-- Forward sweep:
+-- let h_part = reduce_by_index (replicate w ne) odot ne is as
+-- let xs = map2 odot dst h_part
+-- Reverse sweep:
+-- h_part_bar += f'' dst h_part
+-- dst_bar += f' dst h_part
+
+-- let flag = map (\i -> i == 0 || sis[i] != sis[i-1]) (iota n)
+-- let flag_rev = map (\i -> i==0 || flag[n-i]) (iota n)
+-- let ls = seg_scan_exc odot ne flag sas
+-- let rs = reverse sas |>
+--          seg_scan_exc odot ne flag_rev |> reverse
+-- let f_bar = map (\i -> if i < w && -1 < w
+--                        then h_part_bar[i]
+--                        else 0s
+--                 ) sis
+-- let sas_bar = f f_dst ls sas rs
+-- as_bar += scatter (Scratch alpha n) siota sas_bar
+-- Where:
+--  siota: 'iota n' sorted wrt 'is'
+--  sis: 'is' sorted wrt 'is'
+--  sas: 'as' sorted wrt 'is'
+--  f'' = vjpLambda xs_bar h_part (map2 odot)
+--  f' = vjpLambda xs_bar dst (map2 odot)
+--  f  = vjpLambda f_bar sas (map4 (\di li ai ri -> di odot li odot ai odot ri))
+--  0s is an alpha-dimensional array with 0 (possibly 0-dim)
+diffHist :: VjpOps -> [VName] -> StmAux () -> SubExp -> Lambda SOACS -> [SubExp] -> [VName] -> [SubExp] -> SubExp -> [VName] -> ADM () -> ADM ()
+diffHist ops xs aux n lam0 ne as w rf dst m = do
+  as_type <- traverse lookupType $ tail as
+  dst_type <- traverse lookupType dst
+
+  nes <- traverse (letExp "new_dst" . BasicOp . Replicate (Shape $ pure $ head w)) ne
+
+  h_map <- mkIdentityLambda $ Prim int64 : map rowType as_type
+  h_part <- traverse (newVName . flip (<>) "_h_part" . baseString) xs
+  auxing aux . letBindNames h_part . Op $
+    Hist n as [HistOp (Shape w) rf nes ne lam0] h_map
+
+  lam0' <- renameLambda lam0
+  auxing aux . letBindNames xs . Op $
+    Screma (head w) (dst <> h_part) (mapSOAC lam0')
+
+  m
+
+  xs_bar <- traverse lookupAdjVal xs
+
+  (dst_params, hp_params, f') <- mkF' lam0 dst_type $ head w
+  f'_adj_dst <- vjpLambda ops (map adjFromVar xs_bar) dst_params f'
+  f'_adj_hp <- vjpLambda ops (map adjFromVar xs_bar) hp_params f'
+
+  dst_bar' <- eLambda f'_adj_dst $ map (eSubExp . Var) $ dst <> h_part
+  dst_bar <- bindSubExpRes "dst_bar" dst_bar'
+  zipWithM_ updateAdj dst dst_bar
+
+  h_part_bar' <- eLambda f'_adj_hp $ map (eSubExp . Var) $ dst <> h_part
+  h_part_bar <- bindSubExpRes "h_part_bar" h_part_bar'
+
+  lam <- renameLambda lam0
+  lam' <- renameLambda lam0
+
+  -- is' <- mapout (head as) n (head w)
+  -- sorted <- radixSort' (is' : tail as) n $ head w
+  sorted <- radixSort' as n $ head w
+  let siota = head sorted
+  let sis = head $ tail sorted
+  let sas = drop 2 sorted
+
+  iota_n <-
+    letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64
+
+  par_i <- newParam "i" $ Prim int64
+  flag_lam <- mkFlagLam par_i sis
+  flag <- letExp "flag" $ Op $ Screma n [iota_n] $ mapSOAC flag_lam
+
+  -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n)
+  par_i' <- newParam "i" $ Prim int64
+  let i' = paramName par_i'
+  g_lam <-
+    mkLambda [par_i'] $
+      fmap subExpsRes . mapM (letSubExp "scan_inps") =<< do
+        im1 <- letSubExp "i_1" =<< toExp (le64 i' - 1)
+        nmi <- letSubExp "n_i" =<< toExp (pe64 n - le64 i')
+        let s1 = [DimFix im1]
+        let s2 = [DimFix nmi]
+
+        -- flag array for left scan
+        f1 <- letSubExp "f1" $ BasicOp $ Index flag $ Slice [DimFix $ Var i']
+
+        -- array for left scan
+        r1 <-
+          letTupExp' "r1"
+            =<< eIf
+              (eSubExp f1)
+              (eBody $ fmap eSubExp ne)
+              (eBody . fmap (eSubExp . Var) =<< multiIndex sas s1)
+
+        -- array for right scan inc flag
+        r2 <-
+          letTupExp' "r2"
+            =<< eIf
+              (toExp $ le64 i' .==. 0)
+              (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)
+              ( eBody $
+                  pure $ do
+                    eIf
+                      (pure $ BasicOp $ Index flag $ Slice s2)
+                      (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)
+                      ( eBody . fmap eSubExp . (Constant (blankPrimValue Bool) :) . fmap Var
+                          =<< multiIndex sas s2
+                      )
+              )
+
+        traverse eSubExp $ f1 : r1 ++ r2
+
+  -- scan (\(f1,v1) (f2,v2) ->
+  --   let f = f1 || f2
+  --   let v = if f2 then v2 else g v1 v2
+  --   in (f,v) ) (false,ne) (zip flags vals)
+  scan_lams <-
+    traverse
+      ( \l -> do
+          f1 <- newParam "f1" $ Prim Bool
+          f2 <- newParam "f2" $ Prim Bool
+          ps <- lambdaParams <$> renameLambda lam0
+          let (p1, p2) = splitAt (length ne) ps
+
+          mkLambda (f1 : p1 ++ f2 : p2) $
+            fmap varsRes . letTupExp "scan_res" =<< do
+              let f = eBinOp LogOr (eParam f1) (eParam f2)
+              eIf
+                (eParam f2)
+                (eBody $ f : fmap eParam p2)
+                ( eBody . (f :) . fmap (eSubExp . Var)
+                    =<< bindSubExpRes "gres"
+                    =<< eLambda l (fmap eParam ps)
+                )
+      )
+      [lam, lam']
+
+  let ne' = Constant (BoolValue False) : ne
+
+  scansres <-
+    letTupExp "adj_ctrb_scan" . Op $
+      Screma n [iota_n] (scanomapSOAC (map (`Scan` ne') scan_lams) g_lam)
+
+  let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres
+
+  -- map (\i -> if i < w && -1 < w then (xs_bar[i], dst[i]) else (0,ne)) sis
+  par_i'' <- newParam "i" $ Prim int64
+  let i'' = paramName par_i''
+  map_lam <-
+    mkLambda [par_i''] $
+      fmap varsRes . letTupExp "scan_res"
+        =<< eIf
+          (toExp $ withinBounds $ pure (head w, i''))
+          (eBody . fmap (eSubExp . Var) =<< multiIndex h_part_bar [DimFix $ Var i''])
+          ( eBody $ do
+              map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type
+          )
+
+  f_bar <- letTupExp "f_bar" $ Op $ Screma n [sis] $ mapSOAC map_lam
+
+  (as_params, f) <- mkF lam0 as_type n
+  f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f
+
+  -- map (\i -> rs_arr_rev[n-i-1]) (iota n)
+  par_i''' <- newParam "i" $ Prim int64
+  let i''' = paramName par_i'''
+  rev_lam <- mkLambda [par_i'''] $ do
+    nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1)
+    varsRes <$> multiIndex rs_arr_rev [DimFix nmim1]
+
+  rs_arr <- letTupExp "rs_arr" $ Op $ Screma n [iota_n] $ mapSOAC rev_lam
+
+  sas_bar <-
+    bindSubExpRes "sas_bar"
+      =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr)
+
+  scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type
+  as_bar <- multiScatter n scatter_dst siota sas_bar
+
+  zipWithM_ updateAdj (tail as) as_bar
+  where
+    -- map (\i -> if i == 0 then true else is[i] != is[i-1]) (iota n)
+    mkFlagLam :: LParam SOACS -> VName -> ADM (Lambda SOACS)
+    mkFlagLam par_i sis =
+      mkLambda [par_i] $
+        fmap varsRes . letTupExp "flag" =<< do
+          let i = paramName par_i
+          eIf
+            (toExp (le64 i .==. 0))
+            (eBody $ pure $ eSubExp $ Constant $ onePrimValue Bool)
+            ( eBody $
+                pure $ do
+                  i_p <- letExp "i_p" =<< toExp (le64 i - 1)
+                  vs <- traverse (letExp "vs" . BasicOp . Index sis . Slice . pure . DimFix . Var) [i, i_p]
+                  let [vs_i, vs_p] = vs
+                  toExp $ bNot $ le64 vs_i .==. le64 vs_p
+            )
diff --git a/src/Futhark/AD/Rev/Reduce.hs b/src/Futhark/AD/Rev/Reduce.hs
--- a/src/Futhark/AD/Rev/Reduce.hs
+++ b/src/Futhark/AD/Rev/Reduce.hs
@@ -3,6 +3,8 @@
 module Futhark.AD.Rev.Reduce
   ( diffReduce,
     diffMinMaxReduce,
+    diffVecReduce,
+    diffMulReduce,
   )
 where
 
@@ -201,3 +203,140 @@
     letSubExp "minmax_in_bounds" . BasicOp $
       CmpOp (CmpSlt Int64) (intConst Int64 0) w
   updateAdjIndex as (CheckBounds (Just in_bounds), Var x_ind) (Var x_adj)
+
+--
+-- Special case of vectorised reduce:
+--    let x = reduce (map2 op) nes as
+-- Idea:
+--    rewrite to
+--      let x = map2 (\as ne -> reduce op ne as) (transpose as) nes
+--    and diff
+diffVecReduce ::
+  VjpOps -> Pat Type -> StmAux () -> SubExp -> Commutativity -> Lambda SOACS -> VName -> VName -> ADM () -> ADM ()
+diffVecReduce ops x aux w iscomm lam ne as m = do
+  stms <- collectStms_ $ do
+    rank <- arrayRank <$> lookupType as
+    let rear = [1, 0] ++ drop 2 [0 .. rank - 1]
+
+    tran_as <- letExp "tran_as" $ BasicOp $ Rearrange rear as
+    ts <- lookupType tran_as
+    t_ne <- lookupType ne
+
+    as_param <- newParam "as_param" $ rowType ts
+    ne_param <- newParam "ne_param" $ rowType t_ne
+
+    reduce_form <- reduceSOAC [Reduce iscomm lam [Var $ paramName ne_param]]
+
+    map_lam <-
+      mkLambda [as_param, ne_param] $
+        fmap varsRes . letTupExp "idx_res" $
+          Op $
+            Screma w [paramName as_param] reduce_form
+    addStm $ Let x aux $ Op $ Screma (arraySize 0 ts) [tran_as, ne] $ mapSOAC map_lam
+
+  foldr (vjpStm ops) m stms
+
+--
+-- Special case of reduce with mul:
+--    let x = reduce (*) ne as
+-- Forward trace (assuming w = length as):
+--    let (p, z) = map (\a -> if a == 0 then (1, 1) else (a, 0)) as
+--    non_zero_prod = reduce (*) ne p
+--    zr_count = reduce (+) 0 z
+--    let x =
+--      if 0 == zr_count
+--      then non_zero_prod
+--      else 0
+-- Reverse trace:
+--    as_bar = map2
+--      (\a a_bar ->
+--        if zr_count == 0
+--        then a_bar + non_zero_prod/a * x_bar
+--        else if zr_count == 1
+--        then a_bar + (if a == 0 then non_zero_prod * x_bar else 0)
+--        else as_bar
+--      ) as as_bar
+diffMulReduce ::
+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> ADM () -> ADM ()
+diffMulReduce _ops x aux w mul ne as m = do
+  let t = binOpType mul
+  let const_zero = eSubExp $ Constant $ blankPrimValue t
+
+  a_param <- newParam "a" $ Prim t
+  map_lam <-
+    mkLambda [a_param] $
+      fmap varsRes . letTupExp "map_res"
+        =<< eIf
+          (eCmpOp (CmpEq t) (eParam a_param) const_zero)
+          (eBody $ fmap eSubExp [Constant $ onePrimValue t, intConst Int64 1])
+          (eBody [eParam a_param, eSubExp $ intConst Int64 0])
+
+  ps <- newVName "ps"
+  zs <- newVName "zs"
+  auxing aux $
+    letBindNames [ps, zs] $
+      Op $
+        Screma w [as] $
+          mapSOAC map_lam
+
+  red_lam_mul <- binOpLambda mul t
+  red_lam_add <- binOpLambda (Add Int64 OverflowUndef) int64
+
+  red_form_mul <- reduceSOAC $ pure $ Reduce Commutative red_lam_mul $ pure ne
+  red_form_add <- reduceSOAC $ pure $ Reduce Commutative red_lam_add $ pure $ intConst Int64 0
+
+  nz_prods <- newVName "non_zero_prod"
+  zr_count <- newVName "zero_count"
+  auxing aux $ letBindNames [nz_prods] $ Op $ Screma w [ps] red_form_mul
+  auxing aux $ letBindNames [zr_count] $ Op $ Screma w [zs] red_form_add
+
+  auxing aux $
+    letBindNames [x]
+      =<< eIf
+        (toExp $ 0 .==. le64 zr_count)
+        (eBody $ pure $ eSubExp $ Var nz_prods)
+        (eBody $ pure const_zero)
+
+  m
+
+  x_adj <- lookupAdjVal x
+
+  a_param_rev <- newParam "a" $ Prim t
+  map_lam_rev <-
+    mkLambda [a_param_rev] $
+      fmap varsRes . letTupExp "adj_res"
+        =<< eIf
+          (toExp $ 0 .==. le64 zr_count)
+          ( eBody $
+              pure $
+                eBinOp mul (eSubExp $ Var x_adj) $
+                  eBinOp (getDiv t) (eSubExp $ Var nz_prods) $
+                    eParam a_param_rev
+          )
+          ( eBody $
+              pure $
+                eIf
+                  (toExp $ 1 .==. le64 zr_count)
+                  ( eBody $
+                      pure $
+                        eIf
+                          (eCmpOp (CmpEq t) (eParam a_param_rev) const_zero)
+                          ( eBody $
+                              pure $
+                                eBinOp mul (eSubExp $ Var x_adj) $
+                                  eSubExp $
+                                    Var nz_prods
+                          )
+                          (eBody $ pure const_zero)
+                  )
+                  (eBody $ pure const_zero)
+          )
+
+  as_adjup <- letExp "adjs" $ Op $ Screma w [as] $ mapSOAC map_lam_rev
+
+  updateAdj as as_adjup
+  where
+    getDiv :: PrimType -> BinOp
+    getDiv (IntType t) = SDiv t Unsafe
+    getDiv (FloatType t) = FDiv t
+    getDiv _ = error "In getDiv, Reduce.hs: input not supported"
diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs
--- a/src/Futhark/AD/Rev/SOAC.hs
+++ b/src/Futhark/AD/Rev/SOAC.hs
@@ -3,6 +3,7 @@
 module Futhark.AD.Rev.SOAC (vjpSOAC) where
 
 import Control.Monad
+import Futhark.AD.Rev.Hist
 import Futhark.AD.Rev.Map
 import Futhark.AD.Rev.Monad
 import Futhark.AD.Rev.Reduce
@@ -34,14 +35,47 @@
       onOps _ _ _ = m
   onOps ops pat_per_op as_per_op
 
+-- We split multi-op histograms into multiple operations so we
+-- can take advantage of special cases. Post-AD, the result may
+-- be fused again.
+splitHist :: VjpOps -> Pat Type -> StmAux () -> [HistOp SOACS] -> SubExp -> [VName] -> [VName] -> ADM () -> ADM ()
+splitHist vjpops pat aux ops w is as m = do
+  let ks = map (length . histNeutral) ops
+      pat_per_op = map Pat $ chunks ks $ patElems pat
+      as_per_op = chunks ks as
+      onOps (op : ops') (op_pat : op_pats') (op_is : op_is') (op_as : op_as') = do
+        f <- mkIdentityLambda . (Prim int64 :) =<< traverse lookupType op_as
+        vjpSOAC vjpops op_pat aux (Hist w (op_is : op_as) [op] f) $
+          onOps ops' op_pats' op_is' op_as'
+      onOps _ _ _ _ = m
+  onOps ops pat_per_op is as_per_op
+
+-- unfusing a map-histogram construct into a map and a histogram.
+histomapToMapAndHist :: Pat Type -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName]) -> ADM (Stm SOACS, Stm SOACS)
+histomapToMapAndHist (Pat pes) (w, histops, map_lam, as) = do
+  map_pat <- traverse accMapPatElem $ lambdaReturnType map_lam
+  let map_stm = mkLet map_pat $ Op $ Screma w as $ mapSOAC map_lam
+  new_lam <- mkIdentityLambda $ lambdaReturnType map_lam
+  let hist_stm = Let (Pat pes) (defAux ()) $ Op $ Hist w (map identName map_pat) histops new_lam
+  pure (map_stm, hist_stm)
+  where
+    accMapPatElem =
+      newIdent "hist_map_res" . (`arrayOfRow` w)
+
 commonSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
 commonSOAC pat aux soac m = do
   addStm $ Let pat aux $ Op soac
   m
   returnSweepCode $ mapM lookupAdj $ patNames pat
 
+-- Reverse-mode differentiation of SOACs
 vjpSOAC :: VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()
+-- Differentiating Reduces
 vjpSOAC ops pat aux soac@(Screma w as form) m
+  | Just [Reduce iscomm lam [Var ne]] <- isReduceSOAC form,
+    [a] <- as,
+    Just op <- mapOp lam =
+      diffVecReduce ops pat aux w iscomm op ne a m
   | Just reds <- isReduceSOAC form,
     length reds > 1 =
       splitScanRed ops (reduceSOAC, redNeutral) (pat, aux, reds, w, as) m
@@ -52,35 +86,133 @@
     Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,
     isMinMaxOp op =
       diffMinMaxReduce ops x aux w op ne a m
+  | Just [red] <- isReduceSOAC form,
+    [x] <- patNames pat,
+    [ne] <- redNeutral red,
+    [a] <- as,
+    Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,
+    isMulOp op =
+      diffMulReduce ops x aux w op ne a m
   | Just red <- singleReduce <$> isReduceSOAC form = do
       pat_adj <- mapM adjVal =<< commonSOAC pat aux soac m
       diffReduce ops pat_adj w as red
-  where
-    isMinMaxOp (SMin _) = True
-    isMinMaxOp (UMin _) = True
-    isMinMaxOp (FMin _) = True
-    isMinMaxOp (SMax _) = True
-    isMinMaxOp (UMax _) = True
-    isMinMaxOp (FMax _) = True
-    isMinMaxOp _ = False
+
+-- Differentiating Scans
 vjpSOAC ops pat aux soac@(Screma w as form) m
+  | Just [Scan lam [ne]] <- isScanSOAC form,
+    [x] <- patNames pat,
+    [a] <- as,
+    Just [(op, _, _, _)] <- lamIsBinOp lam,
+    isAddOp op = do
+      void $ commonSOAC pat aux soac m
+      diffScanAdd ops x w lam ne a
+  | Just [Scan lam ne] <- isScanSOAC form,
+    Just op <- mapOp lam = do
+      diffScanVec ops (patNames pat) aux w op ne as m
   | Just scans <- isScanSOAC form,
     length scans > 1 =
       splitScanRed ops (scanSOAC, scanNeutral) (pat, aux, scans, w, as) m
   | Just red <- singleScan <$> isScanSOAC form = do
       void $ commonSOAC pat aux soac m
       diffScan ops (patNames pat) w as red
+
+-- Differentiating Maps
 vjpSOAC ops pat aux soac@(Screma w as form) m
   | Just lam <- isMapSOAC form = do
       pat_adj <- commonSOAC pat aux soac m
       vjpMap ops pat_adj aux w lam as
+
+-- Differentiating Redomaps
 vjpSOAC ops pat _aux (Screma w as form) m
   | Just (reds, map_lam) <-
       isRedomapSOAC form = do
       (mapstm, redstm) <-
         redomapToMapAndReduce pat (w, reds, map_lam, as)
       vjpStm ops mapstm $ vjpStm ops redstm m
+
+-- Differentiating Scatter
 vjpSOAC ops pat aux (Scatter w lam ass written_info) m =
   vjpScatter ops pat aux (w, lam, ass, written_info) m
+-- Differentiating Histograms
+vjpSOAC ops pat aux (Hist n as histops f) m
+  | isIdentityLambda f,
+    length histops > 1 = do
+      let (is, vs) = splitAt (length histops) as
+      splitHist ops pat aux histops n is vs m
+vjpSOAC ops pat aux (Hist n [is, vs] [histop] f) m
+  | isIdentityLambda f,
+    [x] <- patNames pat,
+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,
+    lam' <- nestedMapOp lam,
+    Just [(op, _, _, _)] <- lamIsBinOp lam',
+    isMinMaxOp op =
+      diffMinMaxHist ops x aux n op ne is vs w rf dst m
+  | isIdentityLambda f,
+    [x] <- patNames pat,
+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,
+    lam' <- nestedMapOp lam,
+    Just [(op, _, _, _)] <- lamIsBinOp lam',
+    isMulOp op =
+      diffMulHist ops x aux n op ne is vs w rf dst m
+  | isIdentityLambda f,
+    [x] <- patNames pat,
+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,
+    lam' <- nestedMapOp lam,
+    Just [(op, _, _, _)] <- lamIsBinOp lam',
+    isAddOp op =
+      diffAddHist ops x aux n lam ne is vs w rf dst m
+vjpSOAC ops pat aux (Hist n as [histop] f) m
+  | isIdentityLambda f,
+    HistOp (Shape w) rf dst ne lam <- histop = do
+      diffHist ops (patNames pat) aux n lam ne as w rf dst m
+vjpSOAC ops pat _aux (Hist n as histops f) m
+  | not (isIdentityLambda f) = do
+      (mapstm, redstm) <-
+        histomapToMapAndHist pat (n, histops, f, as)
+      vjpStm ops mapstm $ vjpStm ops redstm m
 vjpSOAC _ _ _ soac _ =
   error $ "vjpSOAC unhandled:\n" ++ prettyString soac
+
+---------------
+--- Helpers ---
+---------------
+
+isMinMaxOp :: BinOp -> Bool
+isMinMaxOp (SMin _) = True
+isMinMaxOp (UMin _) = True
+isMinMaxOp (FMin _) = True
+isMinMaxOp (SMax _) = True
+isMinMaxOp (UMax _) = True
+isMinMaxOp (FMax _) = True
+isMinMaxOp _ = False
+
+isMulOp :: BinOp -> Bool
+isMulOp (Mul _ _) = True
+isMulOp (FMul _) = True
+isMulOp _ = False
+
+isAddOp :: BinOp -> Bool
+isAddOp (Add _ _) = True
+isAddOp (FAdd _) = True
+isAddOp _ = False
+
+-- Identifies vectorized operators (lambdas):
+--   if the lambda argument is a map, then returns
+--   just the map's lambda; otherwise nothing.
+mapOp :: Lambda SOACS -> Maybe (Lambda SOACS)
+mapOp (Lambda [pa1, pa2] lam_body _)
+  | [SubExpRes cs r] <- bodyResult lam_body,
+    cs == mempty,
+    [map_stm] <- stmsToList (bodyStms lam_body),
+    (Let (Pat [pe]) _ (Op scrm)) <- map_stm,
+    (Screma _ [a1, a2] (ScremaForm [] [] map_lam)) <- scrm,
+    (a1 == paramName pa1 && a2 == paramName pa2) || (a1 == paramName pa2 && a2 == paramName pa1),
+    r == Var (patElemName pe) =
+      Just map_lam
+mapOp _ = Nothing
+
+-- getting the innermost lambda of a perfect-map nest
+--   (i.e., the first lambda that does not consists of exactly a map)
+nestedMapOp :: Lambda SOACS -> Lambda SOACS
+nestedMapOp lam =
+  maybe lam nestedMapOp (mapOp lam)
diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs
--- a/src/Futhark/AD/Rev/Scan.hs
+++ b/src/Futhark/AD/Rev/Scan.hs
@@ -1,39 +1,76 @@
-module Futhark.AD.Rev.Scan (diffScan) where
+module Futhark.AD.Rev.Scan (diffScan, diffScanVec, diffScanAdd) where
 
 import Control.Monad
+import Data.List (transpose)
 import Futhark.AD.Rev.Monad
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.Builder
 import Futhark.IR.SOACS
+import Futhark.IR.SOACS.Simplify (simplifyLambda)
 import Futhark.Tools
 import Futhark.Transform.Rename
-import Futhark.Util (pairs, unpairs)
+import Futhark.Util (chunk)
 
 data FirstOrSecond = WrtFirst | WrtSecond
 
+identityM :: Int -> Type -> ADM [[SubExp]]
+identityM n t =
+  traverse
+    (traverse (letSubExp "id"))
+    [[if i == j then oneExp t else zeroExp t | i <- [1 .. n]] | j <- [1 .. n]]
+
+matrixMul :: [[PrimExp VName]] -> [[PrimExp VName]] -> PrimType -> [[PrimExp VName]]
+matrixMul m1 m2 t =
+  let zero = primExpFromSubExp t $ Constant $ blankPrimValue t
+   in [[foldl (~+~) zero $ zipWith (~*~) r q | q <- transpose m2] | r <- m1]
+
+matrixVecMul :: [[PrimExp VName]] -> [PrimExp VName] -> PrimType -> [PrimExp VName]
+matrixVecMul m v t =
+  let zero = primExpFromSubExp t $ Constant $ blankPrimValue t
+   in [foldl (~+~) zero $ zipWith (~*~) v r | r <- m]
+
+vectorAdd :: [PrimExp VName] -> [PrimExp VName] -> [PrimExp VName]
+vectorAdd = zipWith (~+~)
+
+orderArgs :: Special -> [a] -> [[a]]
+orderArgs s lst =
+  let d = div (length lst) $ specialScans s
+   in chunk d lst
+
 -- computes `d(x op y)/dx` or d(x op y)/dy
-mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> ADM (Lambda SOACS)
-mkScanAdjointLam ops lam0 which = do
+mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> [SubExp] -> ADM (Lambda SOACS)
+mkScanAdjointLam ops lam0 which adjs = do
   let len = length $ lambdaReturnType lam0
   lam <- renameLambda lam0
   let p2diff =
         case which of
           WrtFirst -> take len $ lambdaParams lam
           WrtSecond -> drop len $ lambdaParams lam
-  p_adjs <- mapM unitAdjOfType (lambdaReturnType lam)
-  vjpLambda ops p_adjs (map paramName p2diff) lam
+  vjpLambda ops (fmap AdjVal adjs) (map paramName p2diff) lam
 
 -- Should generate something like:
 -- `\ j -> let i = n - 1 - j
---         if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (0,1) )`
+--         if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (ys_adj[i],1) )`
 -- where `ys` is  the result of scan
 --       `xs` is  the input  of scan
 --       `ys_adj` is the known adjoint of ys
 --       `j` draw values from `iota n`
-mkScanFusedMapLam :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM (Lambda SOACS)
-mkScanFusedMapLam ops w scn_lam xs ys ys_adj = do
-  lam <- mkScanAdjointLam ops scn_lam WrtFirst
-  ys_ts <- mapM lookupType ys
+mkScanFusedMapLam ::
+  VjpOps ->
+  SubExp ->
+  Lambda SOACS ->
+  [VName] ->
+  [VName] ->
+  [VName] ->
+  Special ->
+  Int ->
+  ADM (Lambda SOACS)
+mkScanFusedMapLam ops w scn_lam xs ys ys_adj s d = do
+  let sc = specialCase s
+  let k = specialSubSize s
+  ys_ts <- traverse lookupType ys
+  idmat <- identityM (length ys) $ rowType $ head ys_ts
+  lams <- traverse (mkScanAdjointLam ops scn_lam WrtFirst) idmat
   par_i <- newParam "i" $ Prim int64
   let i = paramName par_i
   mkLambda [par_i] $
@@ -41,110 +78,268 @@
       =<< eIf
         (toExp $ le64 i .==. 0)
         ( buildBody_ $ do
-            zs <- mapM (letSubExp "ct_zero" . zeroExp . rowType) ys_ts
-            os <- mapM (letSubExp "ct_one" . oneExp . rowType) ys_ts
-            pure $ subExpsRes $ unpairs $ zip zs os
+            j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
+            y_s <- forM ys_adj $ \y_ ->
+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)
+            let zso = orderArgs s y_s
+            let ido = orderArgs s $ case_jac k sc idmat
+            pure $ subExpsRes $ concat $ zipWith (++) zso $ fmap concat ido
         )
         ( buildBody_ $ do
             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
             j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)
-            let index idx arr t = BasicOp $ Index arr $ fullSlice t [DimFix idx]
-            y_s <- forM (zip ys_adj ys_ts) $ \(y_, t) ->
-              letSubExp (baseString y_ ++ "_j") $ index j y_ t
-            lam_rs <-
-              eLambda lam . map pure $
-                zipWith (index j) ys ys_ts ++ zipWith (index j1) xs ys_ts
-            pure $ unpairs $ zip (subExpsRes y_s) lam_rs
+            y_s <- forM ys_adj $ \y_ ->
+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)
+
+            let args =
+                  map (`eIndex` eSubExp j) ys ++ map (`eIndex` eSubExp j1) xs
+            lam_rs <- traverse (`eLambda` args) lams
+
+            let yso = orderArgs s $ subExpsRes y_s
+            let jaco = orderArgs s $ case_jac k sc $ transpose lam_rs
+
+            pure $ concat $ zipWith (++) yso $ fmap concat jaco
         )
+  where
+    case_jac :: Int -> SpecialCase -> [[a]] -> [[a]]
+    case_jac _ Generic jac = jac
+    case_jac k ZeroQuadrant jac =
+      concat
+        $ zipWith
+          (\i -> map (take k . drop (i * k)))
+          [0 .. d `div` k]
+        $ chunk k jac
+    case_jac k MatrixMul jac =
+      take k <$> take k jac
 
--- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b1 * b2)
-mkScanLinFunO :: Type -> ADM (Scan SOACS)
-mkScanLinFunO t = do
+-- a1 a2 b -> a2 + b * a1
+linFunT0 :: [PrimExp VName] -> [PrimExp VName] -> [[PrimExp VName]] -> Special -> PrimType -> [PrimExp VName]
+linFunT0 a1 a2 b s pt =
+  let t = case specialCase s of
+        MatrixMul ->
+          concatMap (\v -> matrixVecMul b v pt) $ chunk (specialSubSize s) a1
+        _ -> matrixVecMul b a1 pt
+   in a2 `vectorAdd` t
+
+-- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b2 * b1)
+mkScanLinFunO :: Type -> Special -> ADM (Scan SOACS)
+mkScanLinFunO t s = do
   let pt = elemType t
-  zero <- letSubExp "zeros" $ zeroExp t
-  one <- letSubExp "ones" $ oneExp t
-  tmp <- mapM newVName ["a1", "b1", "a2", "b2"]
-  let [a1, b1, a2, b2] = tmp
-      pet = primExpFromSubExp pt . Var
-  lam <- mkLambda (map (\v -> Param mempty v t) [a1, b1, a2, b2]) . fmap varsRes $
-    tabNest (arrayRank t) [a1, b1, a2, b2] $ \_ [a1', b1', a2', b2'] -> do
-      x <- letExp "x" <=< toExp $ pet a2' ~+~ pet b2' ~*~ pet a1'
-      y <- letExp "y" <=< toExp $ pet b1' ~*~ pet b2'
-      pure [x, y]
-  pure $ Scan lam [zero, one]
+  neu_elm <- mkNeutral $ specialNeutral s
+  let (as, bs) = specialParams s
+  (a1s, b1s, a2s, b2s) <- mkParams (as, bs)
+  let pet = primExpFromSubExp pt . Var
+  let (_, n) = specialNeutral s
 
--- build the map following the scan with linear-function-composition:
--- for each (ds,cs) length-n array results of the scan, combine them as:
---    `let rs = map2 (\ d_i c_i -> d_i + c_i * y_adj[n-1]) d c |> reverse`
--- but insert explicit indexing to reverse inside the map.
-mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName
-mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do
-  y_adj_last <- letExp (baseString y_adj <> "_last") =<< eLast y_adj
+  lam <- mkLambda (map (\v -> Param mempty v (rowType t)) (a1s ++ b1s ++ a2s ++ b2s)) . fmap subExpsRes $ do
+    let [a1s', b1s', a2s', b2s'] = (fmap . fmap) pet [a1s, b1s, a2s, b2s]
+    let (b1sm, b2sm) = (chunk n b1s', chunk n b2s')
 
-  par_i <- newParam "i" $ Prim int64
-  lam <- mkLambda [par_i] $ do
-    let i = paramName par_i
-    j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
-    dj <- letExp (baseString ds ++ "_dj") $ BasicOp $ Index ds $ fullSlice arr_tp [DimFix j]
-    cj <- letExp (baseString cs ++ "_cj") $ BasicOp $ Index cs $ fullSlice arr_tp [DimFix j]
+    let t0 = linFunT0 a1s' a2s' b2sm s pt
+    let t1 = concat $ matrixMul b2sm b1sm pt
+    traverse (letSubExp "r" <=< toExp) $ t0 ++ t1
 
-    let pet = primExpFromSubExp (elemType arr_tp) . Var
-    fmap varsRes . tabNest (arrayRank (rowType arr_tp)) [y_adj_last, dj, cj] $ \_ [y_adj_last', dj', cj'] ->
-      letTupExp "res" <=< toExp $ pet dj' ~+~ pet cj' ~*~ pet y_adj_last'
+  pure $ Scan lam neu_elm
+  where
+    mkNeutral (a, b) = do
+      zeros <- replicateM a $ letSubExp "zeros" $ zeroExp $ rowType t
+      idmat <- identityM b $ Prim $ elemType t
+      pure $ zeros ++ concat idmat
 
-  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
-  letExp "after_scan" $ Op (Screma w [iota] (ScremaForm [] [] lam))
+    mkParams (a, b) = do
+      a1s <- replicateM a $ newVName "a1"
+      b1s <- replicateM b $ newVName "b1"
+      a2s <- replicateM a $ newVName "a2"
+      b2s <- replicateM b $ newVName "b2"
+      pure (a1s, b1s, a2s, b2s)
 
--- perform the final map, which is fusable with the maps obtained from `mkScan2ndMaps`
+-- perform the final map
 -- let xs_contribs =
 --    map3 (\ i a r -> if i==0 then r else (df2dy (ys[i-1]) a) \bar{*} r)
---         (iota n) xs rs
+--         (iota n) xs (reverse ds)
 mkScanFinalMap :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM [VName]
-mkScanFinalMap ops w scan_lam xs ys rs = do
+mkScanFinalMap ops w scan_lam xs ys ds = do
   let eltps = lambdaReturnType scan_lam
-  lam <- mkScanAdjointLam ops scan_lam WrtSecond
+
   par_i <- newParam "i" $ Prim int64
   let i = paramName par_i
-  par_x <- mapM (\(x, t) -> newParam (baseString x ++ "_par_x") t) $ zip xs eltps
-  par_r <- mapM (\(r, t) -> newParam (baseString r ++ "_par_r") t) $ zip rs eltps
+  par_x <- zipWithM (\x -> newParam (baseString x ++ "_par_x")) xs eltps
 
   map_lam <-
-    mkLambda (par_i : par_x ++ par_r) $
+    mkLambda (par_i : par_x) $ do
+      j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
+
+      dj <-
+        traverse
+          (\dd -> letExp (baseString dd ++ "_dj") =<< eIndex dd (eSubExp j))
+          ds
+
       fmap varsRes . letTupExp "scan_contribs"
         =<< eIf
           (toExp $ le64 i .==. 0)
-          (resultBodyM $ map (Var . paramName) par_r)
+          (resultBodyM $ fmap Var dj)
           ( buildBody_ $ do
-              im1 <- letSubExp "im1" =<< toExp (le64 i - 1)
-              ys_im1 <- forM ys $ \y ->
-                letSubExp (baseString y <> "_im1") =<< eIndex y (eSubExp im1)
-
-              lam_res <-
-                mapM (letExp "const" . BasicOp . SubExp . resSubExp)
-                  =<< eLambda lam (map eSubExp $ ys_im1 ++ map (Var . paramName) par_x)
+              lam <- mkScanAdjointLam ops scan_lam WrtSecond $ fmap Var dj
 
-              fmap (varsRes . mconcat) . forM (zip3 lam_res (map paramName par_r) eltps) $
-                \(lam_r, r, eltp) -> do
-                  let pet = primExpFromSubExp (elemType eltp) . Var
+              im1 <- letSubExp "im1" =<< toExp (le64 i - 1)
+              ys_im1 <- forM ys $ \y -> do
+                y_t <- lookupType y
+                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1]
 
-                  tabNest (arrayRank eltp) [lam_r, r] $ \_ [lam_r', r'] ->
-                    letTupExp "res" <=< toExp $ pet lam_r' ~*~ pet r'
+              let args = map eSubExp $ ys_im1 ++ map (Var . paramName) par_x
+              eLambda lam args
           )
 
   iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
-  letTupExp "scan_contribs" $ Op (Screma w (iota : xs ++ rs) (ScremaForm [] [] map_lam))
+  letTupExp "scan_contribs" $ Op $ Screma w (iota : xs) $ mapSOAC map_lam
 
+data SpecialCase
+  = Generic
+  | ZeroQuadrant
+  | MatrixMul
+  deriving (Show)
+
+data Special = Special
+  { specialNeutral :: (Int, Int),
+    specialParams :: (Int, Int),
+    specialScans :: Int,
+    specialSubSize :: Int,
+    specialCase :: SpecialCase
+  }
+  deriving (Show)
+
+subMats :: Int -> [[Exp SOACS]] -> Exp SOACS -> Maybe Int
+subMats d mat zero =
+  let sub_d = filter (\x -> d `mod` x == 0) [1 .. (d `div` 2)]
+      poss = map (\m -> all (ok m) $ zip mat [0 .. d - 1]) sub_d
+      tmp = filter fst (zip poss sub_d)
+   in if null tmp then Nothing else Just $ snd $ head tmp
+  where
+    ok m (row, i) =
+      all (\(v, j) -> v == zero || (i `div` m == j `div` m)) $
+        zip row [0 .. d - 1]
+
+cases :: Int -> Type -> [[Exp SOACS]] -> Special
+cases d t mat
+  | Just k <- subMats d mat $ zeroExp t =
+      let nonZeros = zipWith (\i -> map (take k . drop (i * k))) [0 .. d `div` k] $ chunk k mat
+       in if all (== head nonZeros) $ tail nonZeros
+            then Special (d, k) (d, k * k) 1 k MatrixMul
+            else Special (k, k) (k, k * k) (d `div` k) k ZeroQuadrant
+cases d _ _ = Special (d, d) (d, d * d) 1 d Generic
+
+identifyCase :: VjpOps -> Lambda SOACS -> ADM Special
+identifyCase ops lam = do
+  let t = lambdaReturnType lam
+  let d = length t
+  idmat <- identityM d $ head t
+  lams <- traverse (mkScanAdjointLam ops lam WrtFirst) idmat
+
+  par1 <- traverse (newParam "tmp1") t
+  par2 <- traverse (newParam "tmp2") t
+  jac_lam <- mkLambda (par1 ++ par2) $ do
+    let args = fmap eParam $ par1 ++ par2
+    lam_rs <- traverse (`eLambda` args) lams
+
+    pure $ concat (transpose lam_rs)
+
+  simp <- simplifyLambda jac_lam
+  let jac = chunk d $ fmap (BasicOp . SubExp . resSubExp) $ bodyResult $ lambdaBody simp
+  pure $ cases d (head t) jac
+
 diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM ()
 diffScan ops ys w as scan = do
+  sc <- identifyCase ops (scanLambda scan)
+  let d = length as
   ys_adj <- mapM lookupAdjVal ys
   as_ts <- mapM lookupType as
-  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj
-  scans_lin_fun_o <- mapM mkScanLinFunO $ lambdaReturnType $ scanLambda scan
+  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d
+  scans_lin_fun_o <- mkScanLinFunO (head as_ts) sc
+  scan_lams <- mkScans (specialScans sc) scans_lin_fun_o
   iota <-
     letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
   r_scan <-
-    letTupExp "adj_ctrb_scan" $
-      Op (Screma w [iota] (ScremaForm scans_lin_fun_o [] map1_lam))
-  red_nms <- mapM (mkScan2ndMaps w) (zip3 as_ts ys_adj (pairs r_scan))
-  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys red_nms
+    letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $
+      scanomapSOAC scan_lams map1_lam
+
+  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d)
   zipWithM_ updateAdj as as_contribs
+  where
+    mkScans :: Int -> Scan SOACS -> ADM [Scan SOACS]
+    mkScans d s =
+      replicateM d $ do
+        lam' <- renameLambda $ scanLambda s
+        pure $ Scan lam' $ scanNeutral s
+
+    splitScanRes sc res d =
+      concat $ take (div d $ specialScans sc) <$> orderArgs sc res
+
+diffScanVec ::
+  VjpOps ->
+  [VName] ->
+  StmAux () ->
+  SubExp ->
+  Lambda SOACS ->
+  [SubExp] ->
+  [VName] ->
+  ADM () ->
+  ADM ()
+diffScanVec ops ys aux w lam ne as m = do
+  stmts <- collectStms_ $ do
+    rank <- arrayRank <$> lookupType (head as)
+    let rear = [1, 0] ++ drop 2 [0 .. rank - 1]
+
+    transp_as <-
+      traverse
+        (\a -> letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange rear a)
+        as
+
+    ts <- traverse lookupType transp_as
+    let n = arraysSize 0 ts
+
+    as_par <- traverse (newParam "as_par" . rowType) ts
+    ne_par <- traverse (newParam "ne_par") $ lambdaReturnType lam
+
+    scan_form <- scanSOAC [Scan lam (map (Var . paramName) ne_par)]
+
+    map_lam <-
+      mkLambda (as_par ++ ne_par) . fmap varsRes . letTupExp "map_res" . Op $
+        Screma w (map paramName as_par) scan_form
+
+    transp_ys <-
+      letTupExp "trans_ys" . Op $
+        Screma n (transp_as ++ subExpVars ne) (mapSOAC map_lam)
+
+    zipWithM
+      (\y x -> auxing aux $ letBindNames [y] $ BasicOp $ Rearrange rear x)
+      ys
+      transp_ys
+
+  foldr (vjpStm ops) m stmts
+
+diffScanAdd :: VjpOps -> VName -> SubExp -> Lambda SOACS -> SubExp -> VName -> ADM ()
+diffScanAdd _ops ys n lam' ne as = do
+  lam <- renameLambda lam'
+  ys_bar <- lookupAdjVal ys
+
+  map_scan <- rev_arr_lam ys_bar
+
+  iota <-
+    letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64
+
+  scan_res <-
+    letExp "res_rev" $ Op $ Screma n [iota] $ scanomapSOAC [Scan lam [ne]] map_scan
+
+  rev_lam <- rev_arr_lam scan_res
+  contrb <- letExp "contrb" $ Op $ Screma n [iota] $ mapSOAC rev_lam
+
+  updateAdj as contrb
+  where
+    rev_arr_lam :: VName -> ADM (Lambda SOACS)
+    rev_arr_lam arr = do
+      par_i <- newParam "i" $ Prim int64
+      mkLambda [par_i] $ do
+        a <-
+          letExp "ys_bar_rev"
+            =<< eIndex arr (toExp (pe64 n - le64 (paramName par_i) - 1))
+        pure [varRes a]
diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs
--- a/src/Futhark/Actions.hs
+++ b/src/Futhark/Actions.hs
@@ -4,6 +4,7 @@
   ( printAction,
     printAliasesAction,
     printLastUseGPU,
+    printLastUseGPUSS,
     printFusionGraph,
     printInterferenceGPU,
     printMemAliasGPU,
@@ -26,7 +27,9 @@
 
 import Control.Monad
 import Control.Monad.IO.Class
+import Data.Bifunctor
 import Data.List (intercalate)
+import Data.Map qualified as M
 import Data.Maybe (fromMaybe)
 import Data.Text qualified as T
 import Data.Text.IO qualified as T
@@ -55,6 +58,7 @@
 import Futhark.IR.Prop.Aliases
 import Futhark.IR.SOACS (SOACS)
 import Futhark.IR.SeqMem (SeqMem)
+import Futhark.Optimise.ArrayShortCircuiting.LastUse qualified as LastUseSS
 import Futhark.Optimise.Fusion.GraphRep qualified
 import Futhark.Util (runProgramWithExitCode, unixEnvironment)
 import Futhark.Version (versionString)
@@ -87,9 +91,24 @@
   Action
     { actionName = "print last use gpu",
       actionDescription = "Print last use information on gpu.",
-      actionProcedure = liftIO . print . LastUse.analyseGPUMem
+      actionProcedure = liftIO . putStrLn . prettyString . M.toList . LastUse.analyseGPUMem
     }
 
+-- | Print last use information to stdout.
+printLastUseGPUSS :: Action GPUMem
+printLastUseGPUSS =
+  Action
+    { actionName = "print last use ss gpu",
+      actionDescription = "Print last use ss information on gpu.",
+      actionProcedure =
+        liftIO
+          . putStrLn
+          . prettyString
+          . bimap M.toList (M.toList . fmap M.toList)
+          . LastUseSS.lastUseGPUMem
+          . aliasAnalysis
+    }
+
 -- | Print fusion graph to stdout.
 printFusionGraph :: Action SOACS
 printFusionGraph =
@@ -225,7 +244,7 @@
   ret_ispc <-
     liftIO $
       runProgramWithExitCode
-        "ispc"
+        cmdISPC
         ( [ispcpath, "-o", ispcbase `addExtension` "o"]
             ++ ["--addressing=64", "--pic"]
             ++ cmdISPCFLAGS ispc_flags -- These flags are always needed
@@ -245,24 +264,25 @@
         mempty
   case ret_ispc of
     Left err ->
-      externalErrorS $ "Failed to run " ++ cmdCC ++ ": " ++ show err
-    Right (ExitFailure code, _, gccerr) -> throwError code gccerr
+      externalErrorS $ "Failed to run " ++ cmdISPC ++ ": " ++ show err
+    Right (ExitFailure code, _, ispcerr) -> throwError cmdISPC code ispcerr
     Right (ExitSuccess, _, _) ->
       case ret of
         Left err ->
           externalErrorS $ "Failed to run ispc: " ++ show err
-        Right (ExitFailure code, _, gccerr) -> throwError code gccerr
+        Right (ExitFailure code, _, gccerr) -> throwError cmdCC code gccerr
         Right (ExitSuccess, _, _) ->
           pure ()
   where
+    cmdISPC = "ispc"
     ispcbase = outpath <> ispcextension
-    throwError code gccerr =
+    throwError prog code err =
       externalErrorS $
-        cmdCC
+        prog
           ++ " failed with code "
           ++ show code
           ++ ":\n"
-          ++ gccerr
+          ++ err
 
 -- | The @futhark c@ action.
 compileCAction :: FutharkConfig -> CompilerMode -> FilePath -> Action SeqMem
diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs
--- a/src/Futhark/Analysis/Interference.hs
+++ b/src/Futhark/Analysis/Interference.hs
@@ -214,7 +214,7 @@
   applyAliases (MemAlias.analyzeGPUMem prog) $
     onConsts (progConsts prog) <> foldMap onFun (progFuns prog)
   where
-    (lumap, _) = LastUse.analyseGPUMem prog
+    lumap = LastUse.analyseGPUMem prog
     onFun f =
       runReader (analyseGPU lumap $ bodyStms $ funDefBody f) $ scopeOf f
     onConsts stms =
diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs
--- a/src/Futhark/Analysis/LastUse.hs
+++ b/src/Futhark/Analysis/LastUse.hs
@@ -41,10 +41,10 @@
 
 type LastUseM rep = Reader (Env rep) (LastUse, Used)
 
-analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)
+analyseGPUMem :: Prog GPUMem -> LastUseMap
 analyseGPUMem = analyseProg analyseGPUOp
 
-analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)
+analyseSeqMem :: Prog SeqMem -> LastUseMap
 analyseSeqMem = analyseProg analyseSeqOp
 
 analyseGPUOp :: LastUseOp GPUMem
@@ -100,7 +100,7 @@
 -- | Analyses a program to return a last-use map, mapping each simple statement
 -- in the program to the values that were last used within that statement, and
 -- the set of all `VName` that were used inside.
-analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> (LastUseMap, Used)
+analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> LastUseMap
 analyseProg onOp prog =
   runReader helper (Env onOp)
   where
@@ -112,9 +112,9 @@
           prog_alias = aliasAnalysis prog
           consts = progConsts prog_alias
           funs = progFuns prog_alias
-      (consts_lu, consts_used) <- analyseStms mempty mempty consts
-      (lus, used) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs
-      pure (flipMap $ consts_lu <> lus, consts_used <> used)
+      (consts_lu, _) <- analyseStms mempty mempty consts
+      (lus, _) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs
+      pure $ flipMap $ consts_lu <> lus
 
 analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep
 analyseFun lumap used fun = do
diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs
--- a/src/Futhark/Analysis/PrimExp.hs
+++ b/src/Futhark/Analysis/PrimExp.hs
@@ -731,13 +731,13 @@
   S.unions $ map leafExpTypes pes
 
 -- | Multiplication of untyped 'PrimExp's, which must have the same
--- type.
+-- type.  Uses 'OverflowWrap' for integer operations.
 (~*~) :: PrimExp v -> PrimExp v -> PrimExp v
 x ~*~ y = BinOpExp op x y
   where
     t = primExpType x
     op = case t of
-      IntType it -> Mul it OverflowUndef
+      IntType it -> Mul it OverflowWrap
       FloatType ft -> FMul ft
       Bool -> LogAnd
       Unit -> LogAnd
@@ -755,23 +755,25 @@
       Unit -> LogAnd
 
 -- | Addition of untyped 'PrimExp's, which must have the same type.
+-- Uses 'OverflowWrap' for integer operations.
 (~+~) :: PrimExp v -> PrimExp v -> PrimExp v
 x ~+~ y = BinOpExp op x y
   where
     t = primExpType x
     op = case t of
-      IntType it -> Add it OverflowUndef
+      IntType it -> Add it OverflowWrap
       FloatType ft -> FAdd ft
       Bool -> LogOr
       Unit -> LogOr
 
 -- | Subtraction of untyped 'PrimExp's, which must have the same type.
+-- Uses 'OverflowWrap' for integer operations.
 (~-~) :: PrimExp v -> PrimExp v -> PrimExp v
 x ~-~ y = BinOpExp op x y
   where
     t = primExpType x
     op = case t of
-      IntType it -> Sub it OverflowUndef
+      IntType it -> Sub it OverflowWrap
       FloatType ft -> FSub ft
       Bool -> LogOr
       Unit -> LogOr
diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs
--- a/src/Futhark/CLI/Dev.hs
+++ b/src/Futhark/CLI/Dev.hs
@@ -44,6 +44,7 @@
 import Futhark.Pass.AD
 import Futhark.Pass.ExpandAllocations
 import Futhark.Pass.ExplicitAllocations.GPU qualified as GPU
+import Futhark.Pass.ExplicitAllocations.MC qualified as MC
 import Futhark.Pass.ExplicitAllocations.Seq qualified as Seq
 import Futhark.Pass.ExtractKernels
 import Futhark.Pass.ExtractMulticore
@@ -229,6 +230,13 @@
   externalErrorS $
     "Pass " ++ name ++ " expects SeqMem representation, but got " ++ representation rep
 
+mcMemProg :: String -> UntypedPassState -> FutharkM (Prog MCMem.MCMem)
+mcMemProg _ (MCMem prog) =
+  pure prog
+mcMemProg name rep =
+  externalErrorS $
+    "Pass " ++ name ++ " expects MCMem representation, but got " ++ representation rep
+
 typedPassOption ::
   Checkable torep =>
   (String -> UntypedPassState -> FutharkM (Prog fromrep)) ->
@@ -263,6 +271,13 @@
 seqMemPassOption =
   typedPassOption seqMemProg SeqMem
 
+mcMemPassOption ::
+  Pass MCMem.MCMem MCMem.MCMem ->
+  String ->
+  FutharkOption
+mcMemPassOption =
+  typedPassOption mcMemProg MCMem
+
 kernelsMemPassOption ::
   Pass GPUMem.GPUMem GPUMem.GPUMem ->
   String ->
@@ -302,6 +317,9 @@
     perform (Seq prog) config =
       SeqMem
         <$> runPipeline (onePass Seq.explicitAllocations) config prog
+    perform (MC prog) config =
+      MCMem
+        <$> runPipeline (onePass MC.explicitAllocations) config prog
     perform s _ =
       externalErrorS $
         "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s
@@ -492,6 +510,14 @@
       "Print last use information.",
     Option
       []
+      ["print-last-use-gpu-ss"]
+      ( NoArg $
+          Right $ \opts ->
+            opts {futharkAction = GPUMemAction $ \_ _ _ -> printLastUseGPUSS}
+      )
+      "Print last use information ss.",
+    Option
+      []
       ["print-interference-gpu"]
       ( NoArg $
           Right $ \opts ->
@@ -599,6 +625,7 @@
     seqMemPassOption LowerAllocations.lowerAllocationsSeqMem [],
     kernelsMemPassOption LowerAllocations.lowerAllocationsGPUMem [],
     seqMemPassOption ArrayShortCircuiting.optimiseSeqMem [],
+    mcMemPassOption ArrayShortCircuiting.optimiseMCMem [],
     kernelsMemPassOption ArrayShortCircuiting.optimiseGPUMem [],
     cseOption [],
     simplifyOption "e",
@@ -641,9 +668,17 @@
       ["seq-mem"],
     pipelineOption
       getSOACSProg
+      "MC"
+      MC
+      "Run the multicore compilation pipeline"
+      mcPipeline
+      []
+      ["mc"],
+    pipelineOption
+      getSOACSProg
       "MCMem"
       MCMem
-      "Run the multicore compilation pipeline"
+      "Run the multicore+memory compilation pipeline"
       multicorePipeline
       []
       ["mc-mem"]
diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs
--- a/src/Futhark/CLI/Test.hs
+++ b/src/Futhark/CLI/Test.hs
@@ -138,10 +138,14 @@
       check ["-s"]
     GpuPipeline ->
       check ["--gpu"]
+    MCPipeline ->
+      check ["--mc"]
     SeqMemPipeline ->
       check ["--seq-mem"]
     GpuMemPipeline ->
       check ["--gpu-mem"]
+    MCMemPipeline ->
+      check ["--mc-mem"]
     NoPipeline ->
       check []
   where
@@ -165,9 +169,11 @@
   where
     maybePipeline :: StructurePipeline -> T.Text
     maybePipeline SOACSPipeline = "(soacs) "
-    maybePipeline GpuPipeline = "(kernels) "
+    maybePipeline GpuPipeline = "(gpu) "
+    maybePipeline MCPipeline = "(mc) "
     maybePipeline SeqMemPipeline = "(seq-mem) "
     maybePipeline GpuMemPipeline = "(gpu-mem) "
+    maybePipeline MCMemPipeline = "(mc-mem) "
     maybePipeline NoPipeline = ""
 
     ok (AstMetrics metrics) (name, expected_occurences) =
diff --git a/src/Futhark/CodeGen/Backends/CCUDA.hs b/src/Futhark/CodeGen/Backends/CCUDA.hs
--- a/src/Futhark/CodeGen/Backends/CCUDA.hs
+++ b/src/Futhark/CodeGen/Backends/CCUDA.hs
@@ -279,17 +279,13 @@
 cudaMemoryType space =
   error $ "CUDA backend does not support '" ++ space ++ "' memory space."
 
-callKernel :: GC.OpCompiler OpenCL ()
-callKernel (GetSize v key) =
-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key;|]
-callKernel (CmpSizeLe v key x) = do
-  x' <- GC.compileExp x
-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key <= $exp:x';|]
-  sizeLoggingCode v key x'
-callKernel (GetSizeMax v size_class) =
-  let field = "max_" ++ cudaSizeClass size_class
-   in GC.stm [C.cstm|$id:v = ctx->cuda.$id:field;|]
+kernelConstToExp :: KernelConst -> C.Exp
+kernelConstToExp (SizeConst key) =
+  [C.cexp|*ctx->tuning_params.$id:key|]
+kernelConstToExp (SizeMaxConst size_class) =
+  [C.cexp|ctx->cuda.$id:field|]
   where
+    field = "max_" <> cudaSizeClass size_class
     cudaSizeClass SizeThreshold {} = "threshold"
     cudaSizeClass SizeGroup = "block_size"
     cudaSizeClass SizeNumGroups = "grid_size"
@@ -297,6 +293,23 @@
     cudaSizeClass SizeRegTile = "reg_tile_size"
     cudaSizeClass SizeLocalMemory = "shared_memory"
     cudaSizeClass (SizeBespoke x _) = prettyString x
+
+compileGroupDim :: GroupDim -> GC.CompilerM op s C.Exp
+compileGroupDim (Left e) = GC.compileExp e
+compileGroupDim (Right kc) = pure $ kernelConstToExp kc
+
+callKernel :: GC.OpCompiler OpenCL ()
+callKernel (GetSize v key) = do
+  let e = kernelConstToExp $ SizeConst key
+  GC.stm [C.cstm|$id:v = $exp:e;|]
+callKernel (CmpSizeLe v key x) = do
+  let e = kernelConstToExp $ SizeConst key
+  x' <- GC.compileExp x
+  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]
+  sizeLoggingCode v key x'
+callKernel (GetSizeMax v size_class) = do
+  let e = kernelConstToExp $ SizeMaxConst size_class
+  GC.stm [C.cstm|$id:v = $exp:e;|]
 callKernel (LaunchKernel safety kernel_name args num_blocks block_size) = do
   args_arr <- newVName "kernel_args"
   time_start <- newVName "time_start"
@@ -310,7 +323,7 @@
     GC.decl [C.cdecl|unsigned int $id:arg = $exp:offset;|]
 
   (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_blocks
-  (block_x, block_y, block_z) <- mkDims <$> mapM GC.compileExp block_size
+  (block_x, block_y, block_z) <- mkDims <$> mapM compileGroupDim block_size
   let perm_args
         | length num_blocks == 3 = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]
         | otherwise = []
diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs
--- a/src/Futhark/CodeGen/Backends/COpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/COpenCL.hs
@@ -326,16 +326,30 @@
 staticOpenCLArray _ space _ _ =
   error $ "OpenCL backend cannot create static array in memory space '" ++ space ++ "'"
 
+kernelConstToExp :: KernelConst -> C.Exp
+kernelConstToExp (SizeConst key) =
+  [C.cexp|*ctx->tuning_params.$id:key|]
+kernelConstToExp (SizeMaxConst size_class) =
+  [C.cexp|ctx->opencl.$id:field|]
+  where
+    field = "max_" <> prettyString size_class
+
+compileGroupDim :: GroupDim -> GC.CompilerM op s C.Exp
+compileGroupDim (Left e) = GC.compileExp e
+compileGroupDim (Right kc) = pure $ kernelConstToExp kc
+
 callKernel :: GC.OpCompiler OpenCL ()
-callKernel (GetSize v key) =
-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key;|]
+callKernel (GetSize v key) = do
+  let e = kernelConstToExp $ SizeConst key
+  GC.stm [C.cstm|$id:v = $exp:e;|]
 callKernel (CmpSizeLe v key x) = do
+  let e = kernelConstToExp $ SizeConst key
   x' <- GC.compileExp x
-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key <= $exp:x';|]
+  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]
   sizeLoggingCode v key x'
-callKernel (GetSizeMax v size_class) =
-  let field = "max_" ++ prettyString size_class
-   in GC.stm [C.cstm|$id:v = ctx->opencl.$id:field;|]
+callKernel (GetSizeMax v size_class) = do
+  let e = kernelConstToExp $ SizeMaxConst size_class
+  GC.stm [C.cstm|$id:v = $exp:e;|]
 callKernel (LaunchKernel safety name args num_workgroups workgroup_size) = do
   -- The other failure args are set automatically when the kernel is
   -- first created.
@@ -349,7 +363,7 @@
 
   zipWithM_ setKernelArg [numFailureParams safety ..] args
   num_workgroups' <- mapM GC.compileExp num_workgroups
-  workgroup_size' <- mapM GC.compileExp workgroup_size
+  workgroup_size' <- mapM compileGroupDim workgroup_size
   local_bytes <- foldM localBytes [C.cexp|0|] args
 
   launchKernel name num_workgroups' workgroup_size' local_bytes
@@ -416,10 +430,11 @@
         fprintf(ctx->log, $string:debug_str, $args:debug_args);
         $id:time_start = get_wall_time();
       }
+      typename cl_event *pevent = $exp:(profilingEvent kernel_name);
       OPENCL_SUCCEED_OR_RETURN(
         clEnqueueNDRangeKernel(ctx->opencl.queue, ctx->$id:kernel_name, $int:kernel_rank, NULL,
                                $id:global_work_size, $id:local_work_size,
-                               0, NULL, $exp:(profilingEvent kernel_name)));
+                               0, NULL, pevent));
       if (ctx->debugging) {
         OPENCL_SUCCEED_FATAL(clFinish(ctx->opencl.queue));
         $id:time_end = get_wall_time();
diff --git a/src/Futhark/CodeGen/Backends/GenericC/Pretty.hs b/src/Futhark/CodeGen/Backends/GenericC/Pretty.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/Pretty.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/Pretty.hs
@@ -7,6 +7,7 @@
     definitionsText,
     typeText,
     idText,
+    funcText,
     funcsText,
   )
 where
@@ -29,5 +30,8 @@
 idText :: C.Id -> T.Text
 idText = T.pack . MPP.pretty 8000 . MPP.ppr
 
+funcText :: C.Func -> T.Text
+funcText = T.pack . MPP.pretty 8000 . MPP.ppr
+
 funcsText :: [C.Func] -> T.Text
-funcsText = T.unlines . map (T.pack . MPP.pretty 8000 . MPP.ppr)
+funcsText = T.unlines . map funcText
diff --git a/src/Futhark/CodeGen/Backends/MulticoreISPC.hs b/src/Futhark/CodeGen/Backends/MulticoreISPC.hs
--- a/src/Futhark/CodeGen/Backends/MulticoreISPC.hs
+++ b/src/Futhark/CodeGen/Backends/MulticoreISPC.hs
@@ -598,16 +598,26 @@
 compileCode (Free name space) = do
   cached <- isJust <$> GC.cacheMem name
   unless cached $ unRefMem name space
-compileCode (For i bound body) = do
-  let i' = C.toIdent i
-      t = GC.primTypeToCType $ primExpType bound
-  bound' <- compileExp bound
-  body' <- GC.collect $ compileCode body
-  quals <- getVariabilityQuals i
-  GC.stm
-    [C.cstm|for ($tyquals:quals $ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) {
+compileCode (For i bound body)
+  -- The special-case here is to avoid certain pathological/contrived
+  -- programs that construct statically known zero-element arrays.
+  -- Due to the way we do constant-fold index functions, this produces
+  -- code that looks like it has uniform/varying mismatches (i.e. race
+  -- conditions) to ISPC, even though that code is never actually run.
+  | isZero bound = pure ()
+  | otherwise = do
+      let i' = C.toIdent i
+          t = GC.primTypeToCType $ primExpType bound
+      bound' <- compileExp bound
+      body' <- GC.collect $ compileCode body
+      quals <- getVariabilityQuals i
+      GC.stm
+        [C.cstm|for ($tyquals:quals $ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) {
             $items:body'
           }|]
+  where
+    isZero (ValueExp v) = zeroIsh v
+    isZero _ = False
 compileCode (While cond body) = do
   cond' <- compileExp $ untyped cond
   body' <- GC.collect $ compileCode body
@@ -937,6 +947,10 @@
         $items:body'
       }
     }|]
+compileOp (ExtractLane dest (ValueExp v) _) =
+  -- extract() on constants is not allowed (type is uniform, not
+  -- varying), so just turn them into an assignment.
+  GC.stm [C.cstm|$id:dest = $exp:v;|]
 compileOp (ExtractLane dest tar lane) = do
   tar' <- compileExp tar
   lane' <- compileExp lane
diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
--- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
@@ -199,27 +199,32 @@
 asLong :: PyExp -> PyExp
 asLong x = Py.simpleCall "np.int64" [x]
 
+kernelConstToExp :: Imp.KernelConst -> PyExp
+kernelConstToExp (Imp.SizeConst key) =
+  Index (Var "self.sizes") (IdxExp $ String $ prettyString key)
+kernelConstToExp (Imp.SizeMaxConst size_class) =
+  Var $ "self.max_" <> prettyString size_class
+
+compileGroupDim :: Imp.GroupDim -> Py.CompilerM op s PyExp
+compileGroupDim (Left e) = asLong <$> Py.compileExp e
+compileGroupDim (Right kc) = pure $ kernelConstToExp kc
+
 callKernel :: Py.OpCompiler Imp.OpenCL ()
 callKernel (Imp.GetSize v key) = do
   v' <- Py.compileVar v
-  Py.stm $
-    Assign v' $
-      Index (Var "self.sizes") (IdxExp $ String $ prettyString key)
+  Py.stm $ Assign v' $ kernelConstToExp $ Imp.SizeConst key
 callKernel (Imp.CmpSizeLe v key x) = do
   v' <- Py.compileVar v
   x' <- Py.compileExp x
   Py.stm $
     Assign v' $
-      BinOp "<=" (Index (Var "self.sizes") (IdxExp $ String $ prettyString key)) x'
+      BinOp "<=" (kernelConstToExp (Imp.SizeConst key)) x'
 callKernel (Imp.GetSizeMax v size_class) = do
   v' <- Py.compileVar v
-  Py.stm $
-    Assign v' $
-      Var $
-        "self.max_" ++ prettyString size_class
+  Py.stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class
 callKernel (Imp.LaunchKernel safety name args num_workgroups workgroup_size) = do
   num_workgroups' <- mapM (fmap asLong . Py.compileExp) num_workgroups
-  workgroup_size' <- mapM (fmap asLong . Py.compileExp) workgroup_size
+  workgroup_size' <- mapM compileGroupDim workgroup_size
   let kernel_size = zipWith mult_exp num_workgroups' workgroup_size'
       total_elements = foldl mult_exp (Integer 1) kernel_size
       cond = BinOp "!=" total_elements (Integer 0)
diff --git a/src/Futhark/CodeGen/ImpCode/GPU.hs b/src/Futhark/CodeGen/ImpCode/GPU.hs
--- a/src/Futhark/CodeGen/ImpCode/GPU.hs
+++ b/src/Futhark/CodeGen/ImpCode/GPU.hs
@@ -10,6 +10,7 @@
     KernelOp (..),
     Fence (..),
     AtomicOp (..),
+    GroupDim,
     Kernel (..),
     KernelUse (..),
     module Futhark.CodeGen.ImpCode,
@@ -32,7 +33,9 @@
 type KernelCode = Code KernelOp
 
 -- | A run-time constant related to kernels.
-newtype KernelConst = SizeConst Name
+data KernelConst
+  = SizeConst Name
+  | SizeMaxConst SizeClass
   deriving (Eq, Ord, Show)
 
 -- | An expression whose variables are kernel constants.
@@ -46,13 +49,16 @@
   | GetSizeMax VName SizeClass
   deriving (Show)
 
+-- | The size of one dimension of a group.
+type GroupDim = Either Exp KernelConst
+
 -- | A generic kernel containing arbitrary kernel code.
 data Kernel = Kernel
   { kernelBody :: Code KernelOp,
     -- | The host variables referenced by the kernel.
     kernelUses :: [KernelUse],
     kernelNumGroups :: [Exp],
-    kernelGroupSize :: [Exp],
+    kernelGroupSize :: [GroupDim],
     -- | A short descriptive and _unique_ name - should be
     -- alphanumeric and without spaces.
     kernelName :: Name,
@@ -81,7 +87,12 @@
 
 instance Pretty KernelConst where
   pretty (SizeConst key) = "get_size" <> parens (pretty key)
+  pretty (SizeMaxConst size_class) = "get_max_size" <> parens (pretty size_class)
 
+instance FreeIn KernelConst where
+  freeIn' (SizeConst _) = mempty
+  freeIn' (SizeMaxConst _) = mempty
+
 instance Pretty KernelUse where
   pretty (ScalarUse name t) =
     oneLine $ "scalar_copy" <> parens (commasep [pretty name, pretty t])
@@ -118,8 +129,11 @@
 
 instance FreeIn Kernel where
   freeIn' kernel =
-    freeIn' (kernelBody kernel)
-      <> freeIn' [kernelNumGroups kernel, kernelGroupSize kernel]
+    freeIn'
+      ( kernelBody kernel,
+        kernelNumGroups kernel,
+        kernelGroupSize kernel
+      )
 
 instance Pretty Kernel where
   pretty kernel =
@@ -128,7 +142,7 @@
         ( "groups"
             <+> brace (pretty $ kernelNumGroups kernel)
             </> "group_size"
-            <+> brace (pretty $ kernelGroupSize kernel)
+            <+> brace (list $ map (either pretty pretty) $ kernelGroupSize kernel)
             </> "uses"
             <+> brace (commasep $ map pretty $ kernelUses kernel)
             </> "failure_tolerant"
diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs
--- a/src/Futhark/CodeGen/ImpCode/OpenCL.hs
+++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs
@@ -16,6 +16,8 @@
     numFailureParams,
     KernelTarget (..),
     FailureMsg (..),
+    GroupDim,
+    KernelConst (..),
     module Futhark.CodeGen.ImpCode,
     module Futhark.IR.GPU.Sizes,
   )
@@ -24,6 +26,7 @@
 import Data.Map qualified as M
 import Data.Text qualified as T
 import Futhark.CodeGen.ImpCode
+import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..))
 import Futhark.IR.GPU.Sizes
 import Futhark.Util.Pretty
 
@@ -92,7 +95,7 @@
 
 -- | Host-level OpenCL operation.
 data OpenCL
-  = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [Exp]
+  = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [GroupDim]
   | GetSize VName Name
   | CmpSizeLe VName Name Exp
   | GetSizeMax VName SizeClass
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
@@ -874,6 +874,8 @@
         constExp =<< hasExp =<< M.lookup name vtable
       constExp (Op (Inner (SizeOp (GetSize key _)))) =
         Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32
+      constExp (Op (Inner (SizeOp (GetSizeMax size_class)))) =
+        Just $ LeafExp (Imp.SizeMaxConst size_class) int32
       constExp e = primExpFromExp lookupConstExp e
   pure $ replaceInPrimExpM onLeaf size
   where
@@ -1141,16 +1143,27 @@
   HostEnv atomics _ locks <- askEnv
   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m
   uses <- computeKernelUses body mempty
+  group_size <- onGroupSize $ kernelGroupSize constants
   emit . Imp.Op . Imp.CallKernel $
     Imp.Kernel
       { Imp.kernelBody = body,
         Imp.kernelUses = uses,
         Imp.kernelNumGroups = [untyped $ kernelNumGroups constants],
-        Imp.kernelGroupSize = [untyped $ kernelGroupSize constants],
+        Imp.kernelGroupSize = [group_size],
         Imp.kernelName = name,
         Imp.kernelFailureTolerant = kAttrFailureTolerant attrs,
         Imp.kernelCheckLocalMemory = kAttrCheckLocalMemory attrs
       }
+  where
+    -- Figure out if this expression actually corresponds to a
+    -- KernelConst.
+    onGroupSize e = do
+      vtable <- getVTable
+      x <- isConstExp vtable $ untyped e
+      pure $
+        case x of
+          Just (LeafExp kc _) -> Right kc
+          _ -> Left $ untyped e
 
 sKernelFailureTolerant ::
   Bool ->
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
@@ -39,6 +39,7 @@
 
 import Control.Monad.Except
 import Data.List (foldl', genericLength, zip5)
+import Data.Map qualified as M
 import Data.Maybe
 import Futhark.CodeGen.ImpCode.GPU qualified as Imp
 import Futhark.CodeGen.ImpGen
@@ -542,10 +543,7 @@
           AtomicCAS f -> pure $ const $ pure f
           AtomicLocking f -> pure $ \hist_H_chk -> do
             let lock_shape =
-                  Shape $
-                    tvSize num_subhistos_per_group
-                      : shapeDims (histOpShape op)
-                      ++ [hist_H_chk]
+                  Shape [tvSize num_subhistos_per_group, hist_H_chk]
 
             let dims = map pe64 $ shapeDims lock_shape
 
@@ -902,6 +900,19 @@
 
   max_group_size <- dPrim "max_group_size" int32
   sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup
+
+  -- XXX: we need to record for later use that max_group_size is the
+  -- result of GetSizeMax.  This is an ugly hack that reflects our
+  -- inability to track which variables are actually constants.
+  let withSizeMax vtable =
+        case M.lookup (tvVar max_group_size) vtable of
+          Just (ScalarVar _ se) ->
+            M.insert
+              (tvVar max_group_size)
+              (ScalarVar (Just (Op (Inner (SizeOp (GetSizeMax SizeGroup))))) se)
+              vtable
+          _ -> vtable
+
   let group_size = Imp.Count $ Var $ tvVar max_group_size
   num_groups <-
     fmap (Imp.Count . tvSize) $
@@ -1030,16 +1041,17 @@
               Just $
                 untyped $
                   unCount groups_per_segment
-        histKernelLocal
-          hist_M
-          groups_per_segment
-          map_pes
-          num_groups
-          group_size
-          space
-          hist_S
-          slugs
-          kbody
+        localVTable withSizeMax $
+          histKernelLocal
+            hist_M
+            groups_per_segment
+            map_pes
+            num_groups
+            group_size
+            space
+            hist_S
+            slugs
+            kbody
 
   pure (pick_local, run)
 
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
@@ -11,6 +11,7 @@
 import Control.Monad.Identity
 import Control.Monad.Reader
 import Control.Monad.State
+import Data.Bifunctor (second)
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Data.Set qualified as S
@@ -64,13 +65,13 @@
 
       (device_prototypes, device_defs) = unzip $ M.elems device_funs
       kernels' = M.map fst kernels
-      opencl_code = openClCode $ map snd $ M.elems kernels
+      opencl_code = T.unlines $ map snd $ M.elems kernels
 
       opencl_prelude =
         T.unlines
           [ genPrelude target used_types,
             definitionsText device_prototypes,
-            funcsText device_defs
+            T.unlines device_defs
           ]
    in ImpOpenCL.Program
         opencl_code
@@ -125,8 +126,8 @@
 errorLabel = ("error_" ++) . show . kernelNextSync
 
 data ToOpenCL = ToOpenCL
-  { clGPU :: M.Map KernelName (KernelSafety, C.Func),
-    clDevFuns :: M.Map Name (C.Definition, C.Func),
+  { clGPU :: M.Map KernelName (KernelSafety, T.Text),
+    clDevFuns :: M.Map Name (C.Definition, T.Text),
     clUsedTypes :: S.Set PrimType,
     clSizes :: M.Map Name SizeClass,
     clFailures :: [FailureMsg]
@@ -237,7 +238,7 @@
   modify $ \s ->
     s
       { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s,
-        clDevFuns = M.insert fname func $ clDevFuns s,
+        clDevFuns = M.insert fname (second funcText func) $ clDevFuns s,
         clFailures = kernelFailures kstate
       }
 
@@ -284,6 +285,15 @@
     toDevice :: HostOp -> KernelOp
     toDevice _ = bad
 
+isConst :: GroupDim -> Maybe T.Text
+isConst (Left (ValueExp (IntValue x))) =
+  Just $ prettyText $ intToInt64 x
+isConst (Right (SizeConst v)) =
+  Just $ T.pack $ zEncodeString $ nameToString v
+isConst (Right (SizeMaxConst size_class)) =
+  Just $ T.pack $ "max_" <> prettyString size_class
+isConst _ = Nothing
+
 onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL
 onKernel target kernel = do
   called <- ensureDeviceFuns $ kernelBody kernel
@@ -389,18 +399,30 @@
           ++ catMaybes local_memory_params
           ++ use_params
 
+      attribute =
+        case (target, mapM isConst $ kernelGroupSize kernel) of
+          (TargetOpenCL, Just [x, y, z]) ->
+            "__attribute__((reqd_work_group_size" <> prettyText (x, y, z) <> "))\n"
+          (TargetOpenCL, Just [x, y]) ->
+            "__attribute__((reqd_work_group_size" <> prettyText (x, y, 1 :: Int) <> "))\n"
+          (TargetOpenCL, Just [x]) ->
+            "__attribute__((reqd_work_group_size" <> prettyText (x, 1 :: Int, 1 :: Int) <> "))\n"
+          _ -> ""
+
       kernel_fun =
-        [C.cfun|__kernel void $id:name ($params:params) {
-                  $items:(mconcat unpack_params)
-                  $items:const_defs
-                  $items:block_dim_init
-                  $items:local_memory_init
-                  $items:error_init
-                  $items:kernel_body
+        attribute
+          <> funcText
+            [C.cfun|__kernel void $id:name ($params:params) {
+                    $items:(mconcat unpack_params)
+                    $items:const_defs
+                    $items:block_dim_init
+                    $items:local_memory_init
+                    $items:error_init
+                    $items:kernel_body
 
-                  $id:(errorLabel kstate): return;
+                    $id:(errorLabel kstate): return;
 
-                  $items:const_undefs
+                    $items:const_undefs
                 }|]
   modify $ \s ->
     s
@@ -411,9 +433,7 @@
 
   -- The argument corresponding to the global_failure parameters is
   -- added automatically later.
-  let args =
-        catMaybes local_memory_args
-          ++ kernelArgs kernel
+  let args = catMaybes local_memory_args ++ kernelArgs kernel
 
   pure $ LaunchKernel safety name args num_groups group_size
   where
@@ -472,15 +492,6 @@
     undef = "#undef " <> idText (C.toIdent v mempty)
 constDef _ = Nothing
 
-openClCode :: [C.Func] -> T.Text
-openClCode kernels =
-  definitionsText [C.cunit|$edecls:funcs|]
-  where
-    funcs =
-      [ [C.cedecl|$func:kernel_func|]
-        | kernel_func <- kernels
-      ]
-
 genOpenClPrelude :: S.Set PrimType -> T.Text
 genOpenClPrelude ts =
   [untrimming|
@@ -637,6 +648,8 @@
   where
     compileKernelConst (SizeConst key) =
       pure [C.cexp|$id:(zEncodeString (prettyString key))|]
+    compileKernelConst (SizeMaxConst size_class) =
+      pure [C.cexp|$id:("max_" <> prettyString size_class)|]
 
 kernelArgs :: Kernel -> [KernelArg]
 kernelArgs = mapMaybe useToArg . kernelUses
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs b/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs
@@ -327,7 +327,7 @@
           <> mapTranspose block_dim args t kind,
       kernelUses = uses,
       kernelNumGroups = map untyped num_groups,
-      kernelGroupSize = map untyped group_size,
+      kernelGroupSize = map (Left . untyped) group_size,
       kernelName = nameFromString name,
       kernelFailureTolerant = True,
       kernelCheckLocalMemory = False
diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs
--- a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs
+++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs
@@ -36,12 +36,15 @@
 lamBody = lambdaBody . segBinOpLambda
 
 -- Arrays for storing worker results.
-resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]]
-resultArrays s segops =
+carryArrays :: String -> TV Int32 -> [SegBinOp MCMem] -> MulticoreGen [[VName]]
+carryArrays s nsubtasks segops =
   forM segops $ \(SegBinOp _ lam _ shape) ->
     forM (lambdaReturnType lam) $ \t -> do
       let pt = elemType t
-          full_shape = shape <> arrayShape t
+          full_shape =
+            Shape [Var (tvVar nsubtasks)]
+              <> shape
+              <> arrayShape t
       sAllocArray s pt full_shape DefaultSpace
 
 nonsegmentedScan ::
@@ -69,14 +72,17 @@
           | vectorize && no_array_param = (scanStage1Nested, scanStage3Nested)
           | otherwise = (scanStage1Fallback, scanStage3Fallback)
 
+    emit $ Imp.DebugPrint "Scan stage 1" Nothing
     scanStage1 pat space kbody scan_ops
 
     let nsubtasks' = tvExp nsubtasks
     sWhen (nsubtasks' .>. 1) $ do
       scan_ops2 <- renameSegBinOp scan_ops
-      scanStage2 pat nsubtasks space scan_ops2 kbody
+      emit $ Imp.DebugPrint "Scan stage 2" Nothing
+      carries <- scanStage2 pat nsubtasks space scan_ops2
       scan_ops3 <- renameSegBinOp scan_ops
-      scanStage3 pat space kbody scan_ops3
+      emit $ Imp.DebugPrint "Scan stage 3" Nothing
+      scanStage3 pat space scan_ops3 carries
 
 -- Different ways to generate code for a scan loop
 data ScanLoopType
@@ -98,7 +104,10 @@
 getExtract _ = extractVectorLane
 
 genBinOpParams :: [SegBinOp MCMem] -> MulticoreGen ()
-genBinOpParams scan_ops = dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops
+genBinOpParams scan_ops =
+  dScope Nothing $
+    scopeOfLParams $
+      concatMap (lambdaParams . segBinOpLambda) scan_ops
 
 genLocalAccsStage1 :: [SegBinOp MCMem] -> MulticoreGen [[VName]]
 genLocalAccsStage1 scan_ops = do
@@ -127,7 +136,41 @@
 getNestLoop ScanNested = sLoopNestVectorized
 getNestLoop _ = sLoopNest
 
--- Generate a loop which performs a potentially vectorized scan.
+applyScanOps ::
+  ScanLoopType ->
+  Pat LetDecMem ->
+  SegSpace ->
+  [SubExp] ->
+  [SegBinOp MCMem] ->
+  [[VName]] ->
+  ImpM MCMem HostEnv Imp.Multicore ()
+applyScanOps typ pat space all_scan_res scan_ops local_accs = do
+  let per_scan_res = segBinOpChunks scan_ops all_scan_res
+      per_scan_pes = segBinOpChunks scan_ops $ patElems pat
+  let (is, _) = unzip $ unSegSpace space
+
+  -- Potential vector load and then do sequential scan
+  getScanLoop typ $ \j ->
+    forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) ->
+      getNestLoop typ (segBinOpShape scan_op) $ \vec_is -> do
+        sComment "Read accumulator" $
+          forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> do
+            copyDWIMFix (paramName p) [] (Var acc') vec_is
+        sComment "Read next values" $
+          forM_ (zip (yParams scan_op) scan_res) $ \(p, se) ->
+            getExtract typ j $
+              collect $
+                copyDWIMFix (paramName p) [] se vec_is
+        -- Scan body
+        sComment "Scan op body" $
+          compileStms mempty (bodyStms $ lamBody scan_op) $
+            forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $
+              \(acc', pe, se) -> do
+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []
+                copyDWIMFix acc' vec_is se []
+
+-- Generate a loop which performs a potentially vectorized scan on the
+-- result of a kernel body.
 genScanLoop ::
   ScanLoopType ->
   Pat LetDecMem ->
@@ -138,37 +181,18 @@
   Imp.TExp Int64 ->
   ImpM MCMem HostEnv Imp.Multicore ()
 genScanLoop typ pat space kbody scan_ops local_accs i = do
-  let (all_scan_res, map_res) = splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody
-      per_scan_res = segBinOpChunks scan_ops all_scan_res
-      per_scan_pes = segBinOpChunks scan_ops $ patElems pat
+  let (all_scan_res, map_res) =
+        splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody
   let (is, ns) = unzip $ unSegSpace space
       ns' = map pe64 ns
 
   zipWithM_ dPrimV_ is $ unflattenIndex ns' i
   compileStms mempty (kernelBodyStms kbody) $ do
-    -- Potential vector load and then do sequential scan
-    getScanLoop typ $ \j -> do
-      sComment "write mapped values results to memory" $ do
-        let map_arrs = drop (segBinOpResults scan_ops) $ patElems pat
-        zipWithM_ (compileThreadResult space) map_arrs map_res
-      forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) ->
-        getNestLoop typ (segBinOpShape scan_op) $ \vec_is -> do
-          -- Read accum value
-          forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> do
-            copyDWIMFix (paramName p) [] (Var acc') vec_is
-          -- Read next value
-          sComment "Read next values" $
-            forM_ (zip (yParams scan_op) scan_res) $ \(p, se) ->
-              getExtract typ j $
-                collect $
-                  copyDWIMFix (paramName p) [] (kernelResultSubExp se) vec_is
-          -- Scan body
-          sComment "Scan body" $
-            compileStms mempty (bodyStms $ lamBody scan_op) $
-              forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $
-                \(acc', pe, se) -> do
-                  copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []
-                  copyDWIMFix acc' vec_is se []
+    let map_arrs = drop (segBinOpResults scan_ops) $ patElems pat
+    sComment "write mapped values results to memory" $
+      zipWithM_ (compileThreadResult space) map_arrs map_res
+    sComment "Apply scan op" $
+      applyScanOps typ pat space (map kernelResultSubExp all_scan_res) scan_ops local_accs
 
 scanStage1Scalar ::
   Pat LetDecMem ->
@@ -200,11 +224,10 @@
     dPrim_ (segFlat space) int64
     sOp $ Imp.GetTaskId (segFlat space)
 
-    lparams <- collect $ genBinOpParams scan_ops
     local_accs <- genLocalAccsStage1 scan_ops
 
     inISPC $ do
-      emit lparams
+      genBinOpParams scan_ops
       generateChunkLoop "SegScan" Scalar $ \i -> do
         genScanLoop ScanNested pat space kbody scan_ops local_accs i
 
@@ -237,14 +260,12 @@
   TV Int32 ->
   SegSpace ->
   [SegBinOp MCMem] ->
-  KernelBody MCMem ->
-  MulticoreGen ()
-scanStage2 pat nsubtasks space scan_ops kbody = do
-  emit $ Imp.DebugPrint "nonsegmentedScan stage 2" Nothing
+  MulticoreGen [[VName]]
+scanStage2 pat nsubtasks space scan_ops = do
   let (is, ns) = unzip $ unSegSpace space
       ns_64 = map pe64 ns
       per_scan_pes = segBinOpChunks scan_ops $ patElems pat
-      nsubtasks' = tvExp nsubtasks
+      nsubtasks' = sExt64 $ tvExp nsubtasks
 
   dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops
   offset <- dPrimV "offset" (0 :: Imp.TExp Int64)
@@ -255,103 +276,106 @@
   -- Parameters used to find the chunk sizes
   -- Perhaps get this information from ``scheduling information``
   -- instead of computing it manually here.
-  let iter_pr_subtask = product ns_64 `quot` sExt64 nsubtasks'
-      remainder = product ns_64 `rem` sExt64 nsubtasks'
+  let iter_pr_subtask = product ns_64 `quot` nsubtasks'
+      remainder = product ns_64 `rem` nsubtasks'
 
-  accs <- resultArrays "scan_stage_2_accum" scan_ops
-  forM_ (zip scan_ops accs) $ \(scan_op, acc) ->
-    sLoopNest (segBinOpShape scan_op) $ \vec_is ->
-      forM_ (zip acc $ segBinOpNeutral scan_op) $ \(acc', ne) ->
-        copyDWIMFix acc' vec_is ne []
+  carries <- carryArrays "scan_stage_2_carry" nsubtasks scan_ops
+  sComment "carry-in for first chunk is neutral" $
+    forM_ (zip scan_ops carries) $ \(scan_op, carry) ->
+      sLoopNest (segBinOpShape scan_op) $ \vec_is ->
+        forM_ (zip carry $ segBinOpNeutral scan_op) $ \(carry', ne) ->
+          copyDWIMFix carry' (0 : vec_is) ne []
 
   -- Perform sequential scan over the last element of each chunk
-  sFor "i" (nsubtasks' - 1) $ \i -> do
+  sComment "scan carries" $ sFor "i" (nsubtasks' - 1) $ \i -> do
     offset <-- iter_pr_subtask
     sWhen (sExt64 i .<. remainder) (offset <-- offset' + 1)
     offset_index <-- offset_index' + offset'
     zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 offset_index'
 
-    compileStms mempty (kernelBodyStms kbody) $
-      forM_ (zip3 per_scan_pes scan_ops accs) $ \(pes, scan_op, acc) ->
-        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do
-          sComment "Read carry in" $
-            forM_ (zip (xParams scan_op) acc) $ \(p, acc') ->
-              copyDWIMFix (paramName p) [] (Var acc') vec_is
-
-          sComment "Read next values" $
-            forM_ (zip (yParams scan_op) pes) $ \(p, pe) ->
-              copyDWIMFix (paramName p) [] (Var $ patElemName pe) ((offset_index' - 1) : vec_is)
+    forM_ (zip3 per_scan_pes scan_ops carries) $ \(pes, scan_op, carry) ->
+      sLoopNest (segBinOpShape scan_op) $ \vec_is -> do
+        sComment "Read carry" $
+          forM_ (zip (xParams scan_op) carry) $ \(p, carry') ->
+            copyDWIMFix (paramName p) [] (Var carry') (i : vec_is)
 
-          compileStms mempty (bodyStms $ lamBody scan_op) $
-            forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $
-              \(acc', pe, se) -> do
-                copyDWIMFix (patElemName pe) ((offset_index' - 1) : vec_is) se []
-                copyDWIMFix acc' vec_is se []
+        sComment "Read next values" $
+          forM_ (zip (yParams scan_op) pes) $ \(p, pe) ->
+            copyDWIMFix (paramName p) [] (Var $ patElemName pe) ((offset_index' - 1) : vec_is)
 
-genLocalAccsStage3 :: [SegBinOp MCMem] -> [[PatElem LetDecMem]] -> MulticoreGen [[VName]]
-genLocalAccsStage3 scan_ops per_scan_pes =
-  forM (zip scan_ops per_scan_pes) $ \(scan_op, pes) -> do
-    let shape = segBinOpShape scan_op
-        ts = lambdaReturnType $ segBinOpLambda scan_op
-    forM (zip4 (xParams scan_op) pes ts $ segBinOpNeutral scan_op) $ \(p, pe, t, ne) -> do
-      acc <-
-        case shapeDims shape of
-          [] -> pure $ paramName p
-          _ -> do
-            let pt = elemType t
-            sAllocArray "local_acc" pt (shape <> arrayShape t) DefaultSpace
+        compileStms mempty (bodyStms $ lamBody scan_op) $
+          forM_ (zip carry $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(carry', se) -> do
+            copyDWIMFix carry' ((i + 1) : vec_is) se []
 
-      -- Initialise the accumulator with neutral from previous chunk.
-      -- or read neutral if first ``iter``
-      (start, _end) <- getLoopBounds
-      sLoopNest (segBinOpShape scan_op) $ \vec_is -> do
-        let read_carry_in =
-              copyDWIMFix acc vec_is (Var $ patElemName pe) (start - 1 : vec_is)
-            read_neutral =
-              copyDWIMFix acc vec_is ne []
-        sIf (start .==. 0) read_neutral read_carry_in
-      pure acc
+  -- Return the array of carries for each chunk.
+  pure carries
 
 scanStage3Scalar ::
   Pat LetDecMem ->
   SegSpace ->
-  KernelBody MCMem ->
   [SegBinOp MCMem] ->
+  [[VName]] ->
   MulticoreGen ()
-scanStage3Scalar pat space kbody scan_ops = do
+scanStage3Scalar pat space scan_ops per_scan_carries = do
   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat
+      (is, ns) = unzip $ unSegSpace space
+      ns' = map pe64 ns
+
   body <- collect $ do
     dPrim_ (segFlat space) int64
-    sOp $ Imp.GetTaskId (segFlat space)
+    sOp $ Imp.GetTaskId $ segFlat space
 
-    genBinOpParams scan_ops
-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes
+    inISPC $ do
+      genBinOpParams scan_ops
+      sComment "load carry-in" $
+        forM_ (zip per_scan_carries scan_ops) $ \(op_carries, scan_op) ->
+          forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->
+            copyDWIMFix (paramName p) [] (Var carries) [le64 (segFlat space)]
+      generateChunkLoop "SegScan" Vectorized $ \i -> do
+        zipWithM_ dPrimV_ is $ unflattenIndex ns' i
+        sComment "load partial result" $
+          forM_ (zip per_scan_pes scan_ops) $ \(scan_pes, scan_op) ->
+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->
+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is)
+        sComment "combine carry with partial result" $
+          forM_ (zip per_scan_pes scan_ops) $ \(scan_pes, scan_op) ->
+            compileStms mempty (bodyStms $ lamBody scan_op) $
+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->
+                copyDWIMFix (patElemName pe) (map Imp.le64 is) se []
 
-    inISPC $
-      generateChunkLoop "SegScan" Vectorized $
-        genScanLoop ScanScalar pat space kbody scan_ops local_accs
   free_params <- freeParams body
   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params
 
 scanStage3Nested ::
   Pat LetDecMem ->
   SegSpace ->
-  KernelBody MCMem ->
   [SegBinOp MCMem] ->
+  [[VName]] ->
   MulticoreGen ()
-scanStage3Nested pat space kbody scan_ops = do
+scanStage3Nested pat space scan_ops per_scan_carries = do
   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat
+      (is, ns) = unzip $ unSegSpace space
+      ns' = map pe64 ns
   body <- collect $ do
     dPrim_ (segFlat space) int64
     sOp $ Imp.GetTaskId (segFlat space)
 
-    lparams <- collect $ genBinOpParams scan_ops
-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes
+    generateChunkLoop "SegScan" Scalar $ \i -> do
+      genBinOpParams scan_ops
+      zipWithM_ dPrimV_ is $ unflattenIndex ns' i
+      forM_ (zip3 per_scan_pes per_scan_carries scan_ops) $ \(scan_pes, op_carries, scan_op) -> do
+        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do
+          sComment "load carry-in" $
+            forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->
+              copyDWIMFix (paramName p) [] (Var carries) (le64 (segFlat space) : vec_is)
 
-    inISPC $ do
-      emit lparams
-      generateChunkLoop "SegScan" Scalar $ \i -> do
-        genScanLoop ScanNested pat space kbody scan_ops local_accs i
+          sComment "load partial result" $
+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->
+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is ++ vec_is)
+          sComment "combine carry with partial result" $
+            compileStms mempty (bodyStms $ lamBody scan_op) $
+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->
+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []
 
   free_params <- freeParams body
   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params
@@ -359,20 +383,34 @@
 scanStage3Fallback ::
   Pat LetDecMem ->
   SegSpace ->
-  KernelBody MCMem ->
   [SegBinOp MCMem] ->
+  [[VName]] ->
   MulticoreGen ()
-scanStage3Fallback pat space kbody scan_ops = do
+scanStage3Fallback pat space scan_ops per_scan_carries = do
   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat
+      (is, ns) = unzip $ unSegSpace space
+      ns' = map pe64 ns
   body <- collect $ do
     dPrim_ (segFlat space) int64
     sOp $ Imp.GetTaskId (segFlat space)
 
     genBinOpParams scan_ops
-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes
 
-    generateChunkLoop "SegScan" Scalar $
-      genScanLoop ScanSeq pat space kbody scan_ops local_accs
+    generateChunkLoop "SegScan" Scalar $ \i -> do
+      zipWithM_ dPrimV_ is $ unflattenIndex ns' i
+      forM_ (zip3 per_scan_pes per_scan_carries scan_ops) $ \(scan_pes, op_carries, scan_op) -> do
+        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do
+          sComment "load carry-in" $
+            forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->
+              copyDWIMFix (paramName p) [] (Var carries) (le64 (segFlat space) : vec_is)
+
+          sComment "load partial result" $
+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->
+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is ++ vec_is)
+          sComment "combine carry with partial result" $
+            compileStms mempty (bodyStms $ lamBody scan_op) $
+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->
+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []
   free_params <- freeParams body
   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params
 
diff --git a/src/Futhark/IR/Mem/IxFun.hs b/src/Futhark/IR/Mem/IxFun.hs
--- a/src/Futhark/IR/Mem/IxFun.hs
+++ b/src/Futhark/IR/Mem/IxFun.hs
@@ -141,10 +141,13 @@
 -- contiguous, i.e., if we instantiate all the points of the current index
 -- function, do we get a contiguous memory interval?
 --
--- By definition, the LMAD denotes the set of points (simplified):
+-- By definition, the LMAD \( \sigma + \{ (n_1, s_1), \ldots, (n_k, s_k) \} \),
+-- where \(n\) and \(s\) denote the shape and stride of each dimension, denotes
+-- the set of points:
 --
---   \{ o + \Sigma_{j=0}^{k} ((i_j+r_j) `mod` n_j)*s_j,
---      \forall i_j such that 0<=i_j<n_j, j=1..k \}
+-- \[
+--    \{ ~ \sigma + i_1 * s_1 + \ldots + i_m * s_m ~ | ~ 0 \leq i_1 < n_1, \ldots, 0 \leq i_m < n_m ~ \}
+-- \]
 data LMAD num = LMAD
   { lmadOffset :: num,
     lmadDims :: [LMADDim num]
diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs
--- a/src/Futhark/IR/Pretty.hs
+++ b/src/Futhark/IR/Pretty.hs
@@ -358,7 +358,7 @@
   pretty (Lambda [] (Body _ stms []) []) | stms == mempty = "nilFn"
   pretty (Lambda params body rettype) =
     "\\"
-      <+> ppTuple' (map pretty params)
+      <+> braces (commastack $ map pretty params)
       </> indent 2 (colon <+> ppTupleLines' (map pretty rettype) <+> "->")
       </> indent 2 (pretty body)
 
diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs
--- a/src/Futhark/IR/Prop/Names.hs
+++ b/src/Futhark/IR/Prop/Names.hs
@@ -213,6 +213,9 @@
 instance (FreeIn a, FreeIn b, FreeIn c, FreeIn d) => FreeIn (a, b, c, d) where
   freeIn' (a, b, c, d) = freeIn' a <> freeIn' b <> freeIn' c <> freeIn' d
 
+instance (FreeIn a, FreeIn b) => FreeIn (Either a b) where
+  freeIn' = either freeIn' freeIn'
+
 instance FreeIn a => FreeIn [a] where
   freeIn' = foldMap freeIn'
 
diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs
--- a/src/Futhark/Internalise/Defunctionalise.hs
+++ b/src/Futhark/Internalise/Defunctionalise.hs
@@ -1,7 +1,6 @@
 -- | Defunctionalization of typed, monomorphic Futhark programs without modules.
 module Futhark.Internalise.Defunctionalise (transformProg) where
 
-import Control.Arrow qualified as Arrow
 import Control.Monad.Identity
 import Control.Monad.Reader
 import Control.Monad.State
@@ -15,21 +14,16 @@
 import Data.Set qualified as S
 import Futhark.IR.Pretty ()
 import Futhark.MonadFreshNames
+import Futhark.Util (mapAccumLM)
 import Language.Futhark
 import Language.Futhark.Traversals
 
--- | An expression or an extended 'Lambda' (with size parameters,
--- which AST lambdas do not support).
-data ExtExp
-  = ExtLambda [Pat] Exp StructRetType SrcLoc
-  | ExtExp Exp
-  deriving (Show)
-
 -- | A static value stores additional information about the result of
 -- defunctionalization of an expression, aside from the residual expression.
 data StaticVal
   = Dynamic PatType
-  | LambdaSV Pat StructRetType ExtExp Env
+  | -- | The Env is the lexical closure of the lambda.
+    LambdaSV Pat StructRetType Exp Env
   | RecordSV [(Name, StaticVal)]
   | -- | The constructor that is actually present, plus
     -- the others that are not.
@@ -56,7 +50,7 @@
 type Env = M.Map VName Binding
 
 localEnv :: Env -> DefM a -> DefM a
-localEnv env = local $ Arrow.second (env <>)
+localEnv env = local $ second (env <>)
 
 -- Even when using a "new" environment (for evaluating closures) we
 -- still ram the global environment of DynamicFuns in there.
@@ -68,7 +62,7 @@
 askEnv = asks snd
 
 areGlobal :: [VName] -> DefM a -> DefM a
-areGlobal vs = local $ Arrow.first (S.fromList vs <>)
+areGlobal vs = local $ first (S.fromList vs <>)
 
 replaceTypeSizes ::
   M.Map VName SizeSubst ->
@@ -98,7 +92,7 @@
        in LambdaSV
             (onAST substs param)
             (RetType t_dims (replaceTypeSizes substs t))
-            (onExtExp substs e)
+            (onExp substs e)
             (onEnv orig_substs closure_env) -- intentional
     Dynamic t ->
       Dynamic $ replaceTypeSizes orig_substs t
@@ -140,6 +134,13 @@
       AppExp (Coerce (onExp substs e) te' loc) (Info (AppRes (replaceTypeSizes substs t) ext))
       where
         te' = onTypeExp substs te
+    onExp substs (Lambda params e ret (Info (als, RetType t_dims t)) loc) =
+      Lambda
+        (map (onAST substs) params)
+        (onExp substs e)
+        ret
+        (Info (als, RetType t_dims (replaceTypeSizes substs t)))
+        loc
     onExp substs e = onAST substs e
 
     onTypeExpDim substs d@(SizeExpNamed v loc) =
@@ -176,15 +177,6 @@
     onTypeExp _ (TEVar v loc) =
       TEVar v loc
 
-    onExtExp substs (ExtExp e) =
-      ExtExp $ onExp substs e
-    onExtExp substs (ExtLambda params e (RetType t_dims t) loc) =
-      ExtLambda
-        (map (onAST substs) params)
-        (onExp substs e)
-        (RetType t_dims (replaceTypeSizes substs t))
-        loc
-
     onEnv substs =
       M.fromList
         . map (second (onBinding substs))
@@ -389,11 +381,11 @@
   -- remaining ones (if there are any) into the body of the lambda.
   let (pat, ret', e0') = case pats of
         [] -> error "Received a lambda with no parameters."
-        [pat'] -> (pat', ret, ExtExp e0)
+        [pat'] -> (pat', ret, e0)
         (pat' : pats') ->
           ( pat',
             RetType [] $ foldFunType (map (toStruct . patternType) pats') ret,
-            ExtLambda pats' e0 ret loc
+            Lambda pats' e0 Nothing (Info (mempty, ret)) loc
           )
 
   -- Construct a record literal that closes over the environment of
@@ -686,11 +678,6 @@
 defuncExp' :: Exp -> DefM Exp
 defuncExp' = fmap fst . defuncExp
 
-defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)
-defuncExtExp (ExtExp e) = defuncExp e
-defuncExtExp (ExtLambda pats e0 ret loc) =
-  defuncFun [] pats e0 ret loc
-
 defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal)
 defuncCase sv (CasePat p e loc) = do
   let p' = updatePat p sv
@@ -722,15 +709,8 @@
 etaExpand :: PatType -> Exp -> DefM ([Pat], Exp, StructRetType)
 etaExpand e_t e = do
   let (ps, ret) = getType $ RetType [] e_t
-  (pats, vars) <- fmap unzip . forM ps $ \(p, t) -> do
-    let t' = fromStruct t
-    x <- case p of
-      Named x -> pure x
-      Unnamed -> newNameFromString "x"
-    pure
-      ( Id x (Info t') mempty,
-        Var (qualName x) (Info t') mempty
-      )
+  -- Some careful hackery to avoid duplicate names.
+  (_, (pats, vars)) <- second unzip <$> mapAccumLM f [] ps
   let e' =
         foldl'
           ( \e1 (e2, t2, argtypes) ->
@@ -746,6 +726,18 @@
       let (ps, r) = getType t2 in ((p, t1) : ps, r)
     getType t = ([], t)
 
+    f prev (p, t) = do
+      let t' = fromStruct t
+      x <- case p of
+        Named x | x `notElem` prev -> pure x
+        _ -> newNameFromString "x"
+      pure
+        ( x : prev,
+          ( Id x (Info t') mempty,
+            Var (qualName x) (Info t') mempty
+          )
+        )
+
 -- | Defunctionalize an indexing of a single array dimension.
 defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
 defuncDimIndex (DimFix e1) = DimFix . fst <$> defuncExp e1
@@ -835,7 +827,7 @@
           dims = mempty
       (e0', sv) <-
         localNewEnv (env' <> closure_env) $
-          defuncExtExp e0
+          defuncExp e0
 
       let closure_pat = buildEnvPat dims closure_env
           pat' = updatePat pat sv2
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
@@ -492,10 +492,10 @@
 
         (loop_initial_cond, init_loop_cond_stms) <- collectStms $ do
           forM_ (zip shapepat shapeinit) $ \(p, se) ->
-            letBindNames [paramName p] $ BasicOp $ SubExp se
+            letBindNames [I.paramName p] $ BasicOp $ SubExp se
           forM_ (zip mergepat' mergeinit) $ \(p, se) ->
-            unless (se == I.Var (paramName p)) $
-              letBindNames [paramName p] $
+            unless (se == I.Var (I.paramName p)) $
+              letBindNames [I.paramName p] $
                 BasicOp $
                   case se of
                     I.Var v
@@ -515,13 +515,13 @@
           -- Careful not to clobber anything.
           loop_end_cond_body <- renameBody <=< buildBody_ $ do
             forM_ (zip shapepat shapeargs) $ \(p, se) ->
-              unless (se == I.Var (paramName p)) $
-                letBindNames [paramName p] $
+              unless (se == I.Var (I.paramName p)) $
+                letBindNames [I.paramName p] $
                   BasicOp $
                     SubExp se
             forM_ (zip mergepat' ses) $ \(p, se) ->
-              unless (se == I.Var (paramName p)) $
-                letBindNames [paramName p] $
+              unless (se == I.Var (I.paramName p)) $
+                letBindNames [I.paramName p] $
                   BasicOp $
                     case se of
                       I.Var v
@@ -1167,7 +1167,7 @@
   img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img'
   let params = bucket_params ++ img_params
       rettype = replicate dim (I.Prim int64) ++ ne_ts
-      body = mkBody mempty $ varsRes $ map paramName params
+      body = mkBody mempty $ varsRes $ map I.paramName params
   lam' <-
     mkLambda params $
       ensureResultShape
@@ -1205,7 +1205,7 @@
     let w = arraysSize 0 bs_ts
     fmap subExpsRes . letValExp' "acc_res" $
       I.Op $
-        I.Screma w (paramName acc_p : bs') (I.mapSOAC lam')
+        I.Screma w (I.paramName acc_p : bs') (I.mapSOAC lam')
 
   op' <-
     case op of
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting.hs b/src/Futhark/Optimise/ArrayShortCircuiting.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting.hs
@@ -3,7 +3,12 @@
 {-# LANGUAGE TypeFamilies #-}
 
 -- | Perform array short circuiting
-module Futhark.Optimise.ArrayShortCircuiting (optimiseSeqMem, optimiseGPUMem) where
+module Futhark.Optimise.ArrayShortCircuiting
+  ( optimiseSeqMem,
+    optimiseGPUMem,
+    optimiseMCMem,
+  )
+where
 
 import Control.Monad.Reader
 import Data.Function ((&))
@@ -12,6 +17,7 @@
 import Futhark.Analysis.Alias qualified as AnlAls
 import Futhark.IR.Aliases
 import Futhark.IR.GPUMem
+import Futhark.IR.MCMem
 import Futhark.IR.Mem.IxFun (substituteInIxFun)
 import Futhark.IR.SeqMem
 import Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing
@@ -20,16 +26,13 @@
 import Futhark.Pass qualified as Pass
 import Futhark.Util
 
-----------------------------------------------------------------
---- Printer/Tester Main Program
-----------------------------------------------------------------
-
 data Env inner = Env
-  { envCoalesceTab :: M.Map VName Coalesced,
-    onInner :: inner -> ReplaceM inner inner
+  { envCoalesceTab :: CoalsTab,
+    onInner :: inner -> UpdateM inner inner,
+    memAllocsToRemove :: Names
   }
 
-type ReplaceM inner a = Reader (Env inner) a
+type UpdateM inner a = Reader (Env inner) a
 
 optimiseSeqMem :: Pass SeqMem SeqMem
 optimiseSeqMem = pass "short-circuit" "Array Short-Circuiting" mkCoalsTab pure replaceInParams
@@ -37,6 +40,9 @@
 optimiseGPUMem :: Pass GPUMem GPUMem
 optimiseGPUMem = pass "short-circuit-gpu" "Array Short-Circuiting (GPU)" mkCoalsTabGPU replaceInHostOp replaceInParams
 
+optimiseMCMem :: Pass MCMem MCMem
+optimiseMCMem = pass "short-circuit-mc" "Array Short-Circuiting (MC)" mkCoalsTabMC replaceInMCOp replaceInParams
+
 replaceInParams :: CoalsTab -> [Param FParamMem] -> (Names, [Param FParamMem])
 replaceInParams coalstab fparams =
   let (mem_allocs_to_remove, fparams') =
@@ -56,47 +62,68 @@
               (to_remove, Param attrs name (MemArray pt shp u $ ArrayIn (dstmem entry) ixf) : acc)
         _ -> (to_remove, Param attrs name dec : acc)
 
-removeStms :: Names -> Body rep -> Body rep
-removeStms to_remove (Body dec stms res) =
-  Body dec (stmsFromList $ filter (not . flip nameIn to_remove . head . patNames . stmPat) $ stmsToList stms) res
+removeAllocsInStms :: Stms rep -> UpdateM inner (Stms rep)
+removeAllocsInStms stms = do
+  to_remove <- asks memAllocsToRemove
+  stmsToList stms
+    & filter (not . flip nameIn to_remove . head . patNames . stmPat)
+    & stmsFromList
+    & pure
 
 pass ::
   (Mem rep inner, LetDec rep ~ LetDecMem, CanBeAliased inner) =>
   String ->
   String ->
-  (FunDef (Aliases rep) -> Pass.PassM CoalsTab) ->
-  (inner -> ReplaceM inner inner) ->
+  (Prog (Aliases rep) -> Pass.PassM (M.Map Name CoalsTab)) ->
+  (inner -> UpdateM inner inner) ->
   (CoalsTab -> [FParam (Aliases rep)] -> (Names, [FParam (Aliases rep)])) ->
   Pass rep rep
 pass flag desc mk on_inner on_fparams =
-  Pass flag desc $
-    Pass.intraproceduralTransformationWithConsts pure $ \_ f -> do
-      coaltab <- mk (AnlAls.analyseFun f)
+  Pass flag desc $ \prog -> do
+    coaltabs <- mk $ AnlAls.aliasAnalysis prog
+    Pass.intraproceduralTransformationWithConsts pure (onFun coaltabs) prog
+  where
+    onFun coaltabs _ f = do
+      let coaltab = coaltabs M.! funDefName f
       let (mem_allocs_to_remove, new_fparams) = on_fparams coaltab $ funDefParams f
       pure $
         f
-          { funDefBody =
-              onBody (foldMap vartab $ M.elems coaltab) $
-                removeStms mem_allocs_to_remove $
-                  funDefBody f,
+          { funDefBody = onBody coaltab mem_allocs_to_remove $ funDefBody f,
             funDefParams = new_fparams
           }
-  where
-    onBody coaltab body =
-      body {bodyStms = runReader (mapM replaceInStm $ bodyStms body) (Env coaltab on_inner)}
 
-replaceInStm :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stm rep -> ReplaceM inner (Stm rep)
+    onBody coaltab mem_allocs_to_remove body =
+      body
+        { bodyStms =
+            runReader
+              (updateStms $ bodyStms body)
+              (Env coaltab on_inner mem_allocs_to_remove),
+          bodyResult = map (replaceResMem coaltab) $ bodyResult body
+        }
+
+replaceResMem :: CoalsTab -> SubExpRes -> SubExpRes
+replaceResMem coaltab res =
+  case flip M.lookup coaltab =<< subExpResVName res of
+    Just entry -> res {resSubExp = Var $ dstmem entry}
+    Nothing -> res
+
+updateStms :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stms rep -> UpdateM inner (Stms rep)
+updateStms stms = do
+  stms' <- mapM replaceInStm stms
+  removeAllocsInStms stms'
+
+replaceInStm :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stm rep -> UpdateM inner (Stm rep)
 replaceInStm (Let (Pat elems) d e) = do
   elems' <- mapM replaceInPatElem elems
   e' <- replaceInExp elems' e
   pure $ Let (Pat elems') d e'
   where
-    replaceInPatElem :: PatElem LetDecMem -> ReplaceM inner (PatElem LetDecMem)
+    replaceInPatElem :: PatElem LetDecMem -> UpdateM inner (PatElem LetDecMem)
     replaceInPatElem p@(PatElem vname (MemArray _ _ u _)) =
       fromMaybe p <$> lookupAndReplace vname PatElem u
     replaceInPatElem p = pure p
 
-replaceInExp :: (Mem rep inner, LetDec rep ~ LetDecMem) => [PatElem LetDecMem] -> Exp rep -> ReplaceM inner (Exp rep)
+replaceInExp :: (Mem rep inner, LetDec rep ~ LetDecMem) => [PatElem LetDecMem] -> Exp rep -> UpdateM inner (Exp rep)
 replaceInExp _ e@(BasicOp _) = pure e
 replaceInExp pat_elems (Match cond_ses cases defbody dec) = do
   defbody' <- replaceInIfBody defbody
@@ -106,38 +133,52 @@
   pure $ Match cond_ses cases' defbody' dec'
 replaceInExp _ (DoLoop loop_inits loop_form (Body dec stms res)) = do
   loop_inits' <- mapM (replaceInFParam . fst) loop_inits
-  stms' <- mapM replaceInStm stms
-  pure $ DoLoop (zip loop_inits' $ map snd loop_inits) loop_form $ Body dec stms' res
-replaceInExp _ e@(Op (Alloc _ _)) = pure e
-replaceInExp _ (Op (Inner i)) = do
-  on_op <- asks onInner
-  Op . Inner <$> on_op i
-replaceInExp _ (Op _) = error "Unreachable" -- This shouldn't be possible?
+  stms' <- updateStms stms
+  coalstab <- asks envCoalesceTab
+  let res' = map (replaceResMem coalstab) res
+  pure $ DoLoop (zip loop_inits' $ map snd loop_inits) loop_form $ Body dec stms' res'
+replaceInExp _ (Op op) =
+  case op of
+    Inner i -> do
+      on_op <- asks onInner
+      Op . Inner <$> on_op i
+    _ -> pure $ Op op
 replaceInExp _ e@WithAcc {} = pure e
 replaceInExp _ e@Apply {} = pure e
 
-replaceInHostOp :: HostOp GPUMem () -> ReplaceM (HostOp GPUMem ()) (HostOp GPUMem ())
-replaceInHostOp (SegOp (SegMap lvl sp tps body)) = do
-  stms <- mapM replaceInStm $ kernelBodyStms body
-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
-replaceInHostOp (SegOp (SegRed lvl sp binops tps body)) = do
-  stms <- mapM replaceInStm $ kernelBodyStms body
-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
-replaceInHostOp (SegOp (SegScan lvl sp binops tps body)) = do
-  stms <- mapM replaceInStm $ kernelBodyStms body
-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
-replaceInHostOp (SegOp (SegHist lvl sp hist_ops tps body)) = do
-  stms <- mapM replaceInStm $ kernelBodyStms body
-  pure $ SegOp $ SegHist lvl sp hist_ops tps $ body {kernelBodyStms = stms}
+replaceInSegOp ::
+  (Mem rep inner, LetDec rep ~ LetDecMem) =>
+  SegOp lvl rep ->
+  UpdateM inner (SegOp lvl rep)
+replaceInSegOp (SegMap lvl sp tps body) = do
+  stms <- updateStms $ kernelBodyStms body
+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
+replaceInSegOp (SegRed lvl sp binops tps body) = do
+  stms <- updateStms $ kernelBodyStms body
+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
+replaceInSegOp (SegScan lvl sp binops tps body) = do
+  stms <- updateStms $ kernelBodyStms body
+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
+replaceInSegOp (SegHist lvl sp hist_ops tps body) = do
+  stms <- updateStms $ kernelBodyStms body
+  pure $ SegHist lvl sp hist_ops tps $ body {kernelBodyStms = stms}
+
+replaceInHostOp :: HostOp GPUMem () -> UpdateM (HostOp GPUMem ()) (HostOp GPUMem ())
+replaceInHostOp (SegOp op) = SegOp <$> replaceInSegOp op
 replaceInHostOp op = pure op
 
-generalizeIxfun :: [PatElem dec] -> PatElem LetDecMem -> BodyReturns -> ReplaceM inner BodyReturns
+replaceInMCOp :: MCOp MCMem () -> UpdateM (MCOp MCMem ()) (MCOp MCMem ())
+replaceInMCOp (ParOp par_op op) =
+  ParOp <$> traverse replaceInSegOp par_op <*> replaceInSegOp op
+replaceInMCOp op = pure op
+
+generalizeIxfun :: [PatElem dec] -> PatElem LetDecMem -> BodyReturns -> UpdateM inner BodyReturns
 generalizeIxfun
   pat_elems
   (PatElem vname (MemArray _ _ _ (ArrayIn mem ixf)))
   m@(MemArray pt shp u _) = do
     coaltab <- asks envCoalesceTab
-    if vname `M.member` coaltab
+    if any (M.member vname . vartab) coaltab
       then
         existentialiseIxFun (map patElemName pat_elems) ixf
           & ReturnsInBlock mem
@@ -146,12 +187,13 @@
       else pure m
 generalizeIxfun _ _ m = pure m
 
-replaceInIfBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> ReplaceM inner (Body rep)
-replaceInIfBody b@(Body _ stms _) = do
-  stms' <- mapM replaceInStm stms
-  pure $ b {bodyStms = stms'}
+replaceInIfBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> UpdateM inner (Body rep)
+replaceInIfBody b@(Body _ stms res) = do
+  coaltab <- asks envCoalesceTab
+  stms' <- updateStms stms
+  pure $ b {bodyStms = stms', bodyResult = map (replaceResMem coaltab) res}
 
-replaceInFParam :: Param FParamMem -> ReplaceM inner (Param FParamMem)
+replaceInFParam :: Param FParamMem -> UpdateM inner (Param FParamMem)
 replaceInFParam p@(Param _ vname (MemArray _ _ u _)) = do
   fromMaybe p <$> lookupAndReplace vname (Param mempty) u
 replaceInFParam p = pure p
@@ -160,10 +202,10 @@
   VName ->
   (VName -> MemBound u -> a) ->
   u ->
-  ReplaceM inner (Maybe a)
+  UpdateM inner (Maybe a)
 lookupAndReplace vname f u = do
   coaltab <- asks envCoalesceTab
-  case M.lookup vname coaltab of
+  case M.lookup vname $ foldMap vartab coaltab of
     Just (Coalesced _ (MemBlock pt shp mem ixf) subs) ->
       ixf
         & fixPoint (substituteInIxFun subs)
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
@@ -1,12 +1,14 @@
-{-# LANGUAGE ConstraintKinds #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeFamilies #-}
 
 -- | The bulk of the short-circuiting implementation.
-module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing (mkCoalsTab, CoalsTab, mkCoalsTabGPU) where
+module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing
+  ( mkCoalsTab,
+    CoalsTab,
+    mkCoalsTabGPU,
+    mkCoalsTabMC,
+  )
+where
 
 import Control.Exception.Base qualified as Exc
 import Control.Monad.Reader
@@ -21,7 +23,8 @@
 import Data.Set qualified as S
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.IR.Aliases
-import Futhark.IR.GPUMem
+import Futhark.IR.GPUMem as GPU
+import Futhark.IR.MCMem as MC
 import Futhark.IR.Mem.IxFun qualified as IxFun
 import Futhark.IR.SeqMem
 import Futhark.MonadFreshNames
@@ -33,7 +36,7 @@
 
 -- | A helper type describing representations that can be short-circuited.
 type Coalesceable rep inner =
-  ( CreatesNewArrOp (OpWithAliases inner),
+  ( Mem rep inner,
     ASTRep rep,
     CanBeAliased inner,
     Op rep ~ MemOp inner,
@@ -49,8 +52,15 @@
 
 type ScalarTableM rep a = Reader (ComputeScalarTableOnOp rep) a
 
-newtype ShortCircuitReader rep = ShortCircuitReader
-  { onOp :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases rep) -> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv
+data ShortCircuitReader rep = ShortCircuitReader
+  { onOp :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases rep) -> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv,
+    ssPointFromOp ::
+      LUTabFun ->
+      TopdownEnv rep ->
+      ScopeTab rep ->
+      Pat (VarAliases, LetDecMem) ->
+      Op (Aliases rep) ->
+      Maybe [SSPointInfo]
   }
 
 newtype ShortCircuitM rep a = ShortCircuitM (ReaderT (ShortCircuitReader rep) (State VNameSource) a)
@@ -87,52 +97,66 @@
 --- Main Coalescing Transformation computes a successful coalescing table    ---
 --------------------------------------------------------------------------------
 
--- | Given a 'FunDef' in 'SegMem' representation, compute the coalescing table
+-- | Given a 'Prog' in 'SegMem' representation, compute the coalescing table
 -- by folding over each function.
-mkCoalsTab :: (MonadFreshNames m) => FunDef (Aliases SeqMem) -> m CoalsTab
-mkCoalsTab =
-  mkCoalsTabFun
-    (snd . lastUseSeqMem)
-    (ShortCircuitReader shortCircuitSeqMem)
+mkCoalsTab :: (MonadFreshNames m) => Prog (Aliases SeqMem) -> m (M.Map Name CoalsTab)
+mkCoalsTab prog =
+  mkCoalsTabProg
+    (lastUseSeqMem prog)
+    (ShortCircuitReader shortCircuitSeqMem genSSPointInfoSeqMem)
     (ComputeScalarTableOnOp $ const $ const $ pure mempty)
+    prog
 
--- | Given a 'FunDef' in 'GPUMem' representation, compute the coalescing table
+-- | Given a 'Prog' in 'GPUMem' representation, compute the coalescing table
 -- by folding over each function.
-mkCoalsTabGPU :: (MonadFreshNames m) => FunDef (Aliases GPUMem) -> m CoalsTab
-mkCoalsTabGPU =
-  mkCoalsTabFun
-    (snd . lastUseGPUMem)
-    (ShortCircuitReader shortCircuitGPUMem)
-    (ComputeScalarTableOnOp computeScalarTableGPUMem)
+mkCoalsTabGPU :: (MonadFreshNames m) => Prog (Aliases GPUMem) -> m (M.Map Name CoalsTab)
+mkCoalsTabGPU prog =
+  mkCoalsTabProg
+    (lastUseGPUMem prog)
+    (ShortCircuitReader shortCircuitGPUMem genSSPointInfoGPUMem)
+    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableGPUMem))
+    prog
 
+-- | Given a 'Prog' in 'MCMem' representation, compute the coalescing table
+-- by folding over each function.
+mkCoalsTabMC :: (MonadFreshNames m) => Prog (Aliases MCMem) -> m (M.Map Name CoalsTab)
+mkCoalsTabMC prog =
+  mkCoalsTabProg
+    (lastUseMCMem prog)
+    (ShortCircuitReader shortCircuitMCMem genSSPointInfoMCMem)
+    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableMCMem))
+    prog
+
 -- | Given a function, compute the coalescing table
-mkCoalsTabFun ::
-  (MonadFreshNames m, Coalesceable rep inner, FParamInfo rep ~ FParamMem) =>
-  (FunDef (Aliases rep) -> LUTabFun) ->
+mkCoalsTabProg ::
+  (MonadFreshNames m, Coalesceable rep inner) =>
+  LUTabProg ->
   ShortCircuitReader rep ->
   ComputeScalarTableOnOp rep ->
-  FunDef (Aliases rep) ->
-  m CoalsTab
-mkCoalsTabFun lufun r computeScalarOnOp fun@(FunDef _ _ _ _ fpars body) = do
-  -- First compute last-use information
-  let lutab = lufun fun
-      unique_mems = getUniqueMemFParam fpars
-      scalar_table =
-        runReader
-          ( concatMapM
-              (computeScalarTable $ scopeOf fun <> scopeOf (bodyStms body))
-              (stmsToList $ bodyStms body)
-          )
-          computeScalarOnOp
-      topenv =
-        emptyTopdownEnv
-          { scope = scopeOfFParams fpars,
-            alloc = unique_mems,
-            scalarTable = scalar_table,
-            nonNegatives = foldMap paramSizes fpars
-          }
-      ShortCircuitM m = fixPointCoalesce lutab fpars body topenv
-  modifyNameSource $ runState (runReaderT m r)
+  Prog (Aliases rep) ->
+  m (M.Map Name CoalsTab)
+mkCoalsTabProg (_, lutab_prog) r computeScalarOnOp = fmap M.fromList . mapM onFun . progFuns
+  where
+    onFun fun@(FunDef _ _ fname _ fpars body) = do
+      -- First compute last-use information
+      let unique_mems = getUniqueMemFParam fpars
+          lutab = lutab_prog M.! fname
+          scalar_table =
+            runReader
+              ( concatMapM
+                  (computeScalarTable $ scopeOf fun <> scopeOf (bodyStms body))
+                  (stmsToList $ bodyStms body)
+              )
+              computeScalarOnOp
+          topenv =
+            emptyTopdownEnv
+              { scope = scopeOfFParams fpars,
+                alloc = unique_mems,
+                scalarTable = scalar_table,
+                nonNegatives = foldMap paramSizes fpars
+              }
+          ShortCircuitM m = fixPointCoalesce lutab fpars body topenv
+      (fname,) <$> modifyNameSource (runState (runReaderT m r))
 
 paramSizes :: Param FParamMem -> Names
 paramSizes (Param _ _ (MemArray _ shp _ _)) = freeIn shp
@@ -145,28 +169,20 @@
 shortCircuitSeqMem :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases SeqMem) -> TopdownEnv SeqMem -> BotUpEnv -> ShortCircuitM SeqMem BotUpEnv
 shortCircuitSeqMem _ _ _ _ = pure
 
--- | Short-circuit handler for 'GPUMem' 'Op'.
---
--- When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do
--- nothing.
-shortCircuitGPUMem ::
+-- | Short-circuit handler for SegOp.
+shortCircuitSegOp ::
+  Coalesceable rep inner =>
+  (lvl -> Bool) ->
   LUTabFun ->
   Pat (VarAliases, LetDecMem) ->
-  Op (Aliases GPUMem) ->
-  TopdownEnv GPUMem ->
+  SegOp lvl (Aliases rep) ->
+  TopdownEnv rep ->
   BotUpEnv ->
-  ShortCircuitM GPUMem BotUpEnv
-shortCircuitGPUMem _ _ (Alloc _ _) _ bu_env = pure bu_env
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThread {} space _ kernel_body))) td_env bu_env =
-  -- No special handling necessary for 'SegMap'. Just call the helper-function.
-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegGroup {} space _ kernel_body))) td_env bu_env =
-  -- No special handling necessary for 'SegMap'. Just call the helper-function.
-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThreadInGroup {} space _ kernel_body))) td_env bu_env =
+  ShortCircuitM rep BotUpEnv
+shortCircuitSegOp lvlOK lutab pat (SegMap lvl space _ kernel_body) td_env bu_env =
   -- No special handling necessary for 'SegMap'. Just call the helper-function.
-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegRed lvl space binops _ kernel_body))) td_env bu_env =
+  shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env
+shortCircuitSegOp lvlOK lutab pat (SegRed lvl space binops _ kernel_body) td_env bu_env =
   -- When handling 'SegRed', we we first invalidate all active coalesce-entries
   -- where any of the variables in 'vartab' are also free in the list of
   -- 'SegBinOp'. In other words, anything that is used as part of the reduction
@@ -176,26 +192,26 @@
         foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
       bu_env' = bu_env {activeCoals = active, inhibit = inh}
       num_reds = length red_ts
-   in shortCircuitGPUMemHelper num_reds lvl lutab pat space kernel_body td_env bu_env'
+   in shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat space kernel_body td_env bu_env'
   where
     segment_dims = init $ segSpaceDims space
     red_ts = do
       op <- binops
       let shp = Shape segment_dims <> segBinOpShape op
       map (`arrayOfShape` shp) (lambdaReturnType $ segBinOpLambda op)
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegScan lvl space binops _ kernel_body))) td_env bu_env =
+shortCircuitSegOp lvlOK lutab pat (SegScan lvl space binops _ kernel_body) td_env bu_env =
   -- Like in the handling of 'SegRed', we do not want to coalesce anything that
   -- is used in the 'SegBinOp'
   let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env
       (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
       bu_env' = bu_env {activeCoals = active, inhibit = inh}
-   in shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'
-shortCircuitGPUMem lutab pat (Inner (SegOp (SegHist lvl space histops _ kernel_body))) td_env bu_env = do
+   in shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env'
+shortCircuitSegOp lvlOK lutab pat (SegHist lvl space histops _ kernel_body) td_env bu_env = do
   -- Need to take zipped patterns and histDest (flattened) and insert transitive coalesces
   let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . histOp) histops) $ activeCoals bu_env
       (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail
       bu_env' = bu_env {activeCoals = active, inhibit = inh}
-  bu_env'' <- shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'
+  bu_env'' <- shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env'
   pure $
     foldl insertHistCoals bu_env'' $
       zip (patElems pat) $
@@ -216,17 +232,33 @@
                     }
             Nothing -> acc
         _ -> acc
-shortCircuitGPUMem lutab pat (Inner (GPUBody _ body)) td_env bu_env = do
+
+-- | Short-circuit handler for 'GPUMem' 'Op'.
+--
+-- When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do
+-- nothing.
+shortCircuitGPUMem ::
+  LUTabFun ->
+  Pat (VarAliases, LetDecMem) ->
+  Op (Aliases GPUMem) ->
+  TopdownEnv GPUMem ->
+  BotUpEnv ->
+  ShortCircuitM GPUMem BotUpEnv
+shortCircuitGPUMem _ _ (Alloc _ _) _ bu_env = pure bu_env
+shortCircuitGPUMem lutab pat (Inner (GPU.SegOp op)) td_env bu_env =
+  shortCircuitSegOp isSegThread lutab pat op td_env bu_env
+shortCircuitGPUMem lutab pat (Inner (GPU.GPUBody _ body)) td_env bu_env = do
   fresh1 <- newNameFromString "gpubody"
   fresh2 <- newNameFromString "gpubody"
-  shortCircuitGPUMemHelper
+  shortCircuitSegOpHelper
     0
+    isSegThread
     -- Construct a 'SegLevel' corresponding to a single thread
-    ( SegThread SegNoVirt $
+    ( GPU.SegThread GPU.SegNoVirt $
         Just $
-          KernelGrid
-            (Count $ Constant $ IntValue $ Int64Value 1)
-            (Count $ Constant $ IntValue $ Int64Value 1)
+          GPU.KernelGrid
+            (GPU.Count $ Constant $ IntValue $ Int64Value 1)
+            (GPU.Count $ Constant $ IntValue $ Int64Value 1)
     )
     lutab
     pat
@@ -234,14 +266,29 @@
     (bodyToKernelBody body)
     td_env
     bu_env
-shortCircuitGPUMem _ _ (Inner (SizeOp _)) _ bu_env = pure bu_env
-shortCircuitGPUMem _ _ (Inner (OtherOp ())) _ bu_env = pure bu_env
+shortCircuitGPUMem _ _ (Inner (GPU.SizeOp _)) _ bu_env = pure bu_env
+shortCircuitGPUMem _ _ (Inner (GPU.OtherOp ())) _ bu_env = pure bu_env
 
+shortCircuitMCMem ::
+  LUTabFun ->
+  Pat (VarAliases, LetDecMem) ->
+  Op (Aliases MCMem) ->
+  TopdownEnv MCMem ->
+  BotUpEnv ->
+  ShortCircuitM MCMem BotUpEnv
+shortCircuitMCMem _ _ (Alloc _ _) _ bu_env = pure bu_env
+shortCircuitMCMem _ _ (Inner (MC.OtherOp ())) _ bu_env = pure bu_env
+shortCircuitMCMem lutab pat (Inner (MC.ParOp (Just par_op) op)) td_env bu_env =
+  shortCircuitSegOp (const True) lutab pat par_op td_env bu_env
+    >>= shortCircuitSegOp (const True) lutab pat op td_env
+shortCircuitMCMem lutab pat (Inner (MC.ParOp Nothing op)) td_env bu_env =
+  shortCircuitSegOp (const True) lutab pat op td_env bu_env
+
 dropLastSegSpace :: SegSpace -> SegSpace
 dropLastSegSpace space = space {unSegSpace = init $ unSegSpace space}
 
-isSegThread :: SegLevel -> Bool
-isSegThread SegThread {} = True
+isSegThread :: GPU.SegLevel -> Bool
+isSegThread GPU.SegThread {} = True
 isSegThread _ = False
 
 -- | Computes the slice written at the end of a thread in a 'SegOp'.
@@ -282,18 +329,21 @@
 --
 -- 4. Mark active coalescings as finished, since a 'SegOp' is an array creation
 -- point.
-shortCircuitGPUMemHelper ::
+shortCircuitSegOpHelper ::
+  Coalesceable rep inner =>
   -- | The number of returns for which we should drop the last seg space
   Int ->
-  SegLevel ->
+  -- | Whether we should look at a segop with this lvl.
+  (lvl -> Bool) ->
+  lvl ->
   LUTabFun ->
   Pat (VarAliases, LetDecMem) ->
   SegSpace ->
-  KernelBody (Aliases GPUMem) ->
-  TopdownEnv GPUMem ->
+  KernelBody (Aliases rep) ->
+  TopdownEnv rep ->
   BotUpEnv ->
-  ShortCircuitM GPUMem BotUpEnv
-shortCircuitGPUMemHelper num_reds lvl lutab pat@(Pat ps0) space0 kernel_body td_env bu_env = do
+  ShortCircuitM rep BotUpEnv
+shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat@(Pat ps0) space0 kernel_body td_env bu_env = do
   -- We need to drop the last element of the 'SegSpace' for pattern elements
   -- that correspond to reductions.
   let ps_space_and_res =
@@ -311,7 +361,7 @@
       (actv_return, inhibit_return) =
         if num_reds > 0
           then (actv0, inhibit0)
-          else foldl (makeSegMapCoals lvl td_env kernel_body) (actv0, inhibit0) ps_space_and_res
+          else foldl (makeSegMapCoals lvlOK lvl td_env kernel_body) (actv0, inhibit0) ps_space_and_res
 
   -- Start from empty references, we'll update with aggregates later.
   let actv0' = M.map (\etry -> etry {memrefs = mempty}) $ actv0 <> actv_return
@@ -442,67 +492,61 @@
 
 -- | Given a pattern element and the corresponding kernel result, try to put the
 -- kernel result directly in the memory block of pattern element
-makeSegMapCoals :: SegLevel -> TopdownEnv GPUMem -> KernelBody (Aliases GPUMem) -> (CoalsTab, InhibitTab) -> (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) -> (CoalsTab, InhibitTab)
-makeSegMapCoals lvl td_env kernel_body (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))
-  | Just mb@(MemBlock tp return_shp return_mem _) <-
+makeSegMapCoals ::
+  (Coalesceable rep inner) =>
+  (lvl -> Bool) ->
+  lvl ->
+  TopdownEnv rep ->
+  KernelBody (Aliases rep) ->
+  (CoalsTab, InhibitTab) ->
+  (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) ->
+  (CoalsTab, InhibitTab)
+makeSegMapCoals lvlOK lvl td_env kernel_body (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))
+  | Just (MemBlock tp return_shp return_mem _) <-
       getScopeMemInfo return_name $ scope td_env <> scopeOf (kernelBodyStms kernel_body),
-    isSegThread lvl,
+    lvlOK lvl,
     MemMem pat_space <- runReader (lookupMemInfo pat_mem) $ removeScopeAliases $ scope td_env,
-    MemMem return_space <- runReader (lookupMemInfo return_mem) $ removeScopeAliases $ scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space,
+    MemMem return_space <-
+      scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space
+        & removeScopeAliases
+        & runReader (lookupMemInfo return_mem),
     pat_space == return_space =
       case M.lookup pat_mem active of
         Nothing ->
           -- We are not in a transitive case
-          if IxFun.hasOneLmad pat_ixf
-            then case ( maybe False (pat_mem `nameIn`) $ M.lookup return_mem inhb,
-                        Coalesced InPlaceCoal mb mempty
-                          & M.singleton return_name
-                          & flip (addInvAliassesVarTab td_env) return_name
-                          & fmap
-                            ( M.adjust
-                                ( \(Coalesced knd (MemBlock pt shp _ _) subst) ->
-                                    Coalesced
-                                      knd
-                                      ( MemBlock pt shp pat_mem $
-                                          IxFun.slice pat_ixf $
-                                            fullSlice (IxFun.shape pat_ixf) $
-                                              Slice $
-                                                map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $
-                                                  unSegSpace space
-                                      )
-                                      subst
-                                )
-                                return_name
-                            )
-                      ) of
-              (False, Just vtab) ->
-                (active <> M.singleton return_mem (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty), inhb)
-              _ -> (active, inhb)
-            else (active, inhb)
-        Just trans ->
-          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,
-                 Coalesced InPlaceCoal (MemBlock tp return_shp (dstmem trans) (dstind trans)) mempty
+          case ( IxFun.hasOneLmad pat_ixf
+                   && maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),
+                 Coalesced
+                   InPlaceCoal
+                   (MemBlock tp return_shp pat_mem $ resultSlice pat_ixf)
+                   mempty
                    & M.singleton return_name
                    & flip (addInvAliassesVarTab td_env) return_name
-                   & fmap
-                     ( M.adjust
-                         ( \(Coalesced knd (MemBlock pt shp mem ixf@(IxFun.IxFun _ base_shape _)) subst) ->
-                             Coalesced
-                               knd
-                               ( MemBlock pt shp mem $
-                                   IxFun.slice ixf $
-                                     fullSlice base_shape $
-                                       Slice $
-                                         map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $
-                                           unSegSpace space
-                               )
-                               subst
-                         )
-                         return_name
-                     )
                ) of
             (False, Just vtab) ->
-              let opts = if dstmem trans == pat_mem then mempty else M.insert pat_name pat_mem $ optdeps trans
+              ( active
+                  <> M.singleton
+                    return_mem
+                    (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty),
+                inhb
+              )
+            _ -> (active, inhb)
+        Just trans ->
+          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,
+                 let Coalesced _ (MemBlock _ _ trans_mem trans_ixf) _ =
+                       fromMaybe (error "Impossible") $ M.lookup pat_name $ vartab trans
+                  in Coalesced
+                       TransitiveCoal
+                       (MemBlock tp return_shp trans_mem $ resultSlice trans_ixf)
+                       mempty
+                       & M.singleton return_name
+                       & flip (addInvAliassesVarTab td_env) return_name
+               ) of
+            (False, Just vtab) ->
+              let opts =
+                    if dstmem trans == pat_mem
+                      then mempty
+                      else M.insert pat_name pat_mem $ optdeps trans
                in ( M.insert
                       return_mem
                       ( CoalsEntry
@@ -517,22 +561,28 @@
                     inhb
                   )
             _ -> (active, inhb)
-makeSegMapCoals _ td_env _ x (_, _, WriteReturns _ _ return_name _) =
+  where
+    thread_slice =
+      unSegSpace space
+        & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)
+        & Slice
+    resultSlice ixf = IxFun.slice ixf $ fullSlice (IxFun.shape ixf) thread_slice
+makeSegMapCoals _ _ td_env _ x (_, _, WriteReturns _ _ return_name _) =
   case getScopeMemInfo return_name $ scope td_env of
     Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem
     Nothing -> error "Should not happen?"
-makeSegMapCoals _ td_env _ x (_, _, result) =
+makeSegMapCoals _ _ td_env _ x (_, _, result) =
   freeIn result
     & namesToList
     & mapMaybe (flip getScopeMemInfo $ scope td_env)
-    & foldr (\(MemBlock _ _ mem _) -> flip markFailedCoal mem) x
+    & foldr (flip markFailedCoal . memName) x
 
 fullSlice :: [TPrimExp Int64 VName] -> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName)
 fullSlice shp (Slice slc) =
   Slice $ slc ++ map (\d -> DimSlice 0 d 1) (drop (length slc) shp)
 
 fixPointCoalesce ::
-  (Coalesceable rep inner) =>
+  Coalesceable rep inner =>
   LUTabFun ->
   [Param FParamMem] ->
   Body (Aliases rep) ->
@@ -593,7 +643,7 @@
 
 -- | Perform short-circuiting on 'Stms'.
 mkCoalsTabStms ::
-  (Coalesceable rep inner) =>
+  Coalesceable rep inner =>
   LUTabFun ->
   Stms (Aliases rep) ->
   TopdownEnv rep ->
@@ -636,7 +686,7 @@
 --                          then the checks should be extended to the actual
 --                          array-creation points.
 mkCoalsTabStm ::
-  (Coalesceable rep inner) =>
+  Coalesceable rep inner =>
   LUTabFun ->
   Stm (Aliases rep) ->
   TopdownEnv rep ->
@@ -1037,7 +1087,7 @@
 
 -- The case of in-place update:
 --   @let x' = x with slice <- elm@
-mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (Update safety x _ _elm))) td_env bu_env
+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (Update safety x _ _elm))) td_env bu_env
   | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =
       do
         -- (a) filter by the 3rd safety for @elm@ and @x'@
@@ -1065,14 +1115,14 @@
                         _ ->
                           markFailedCoal (actv, inhbt) m_x
 
-            -- (c) this stm is also a potential source for coalescing, so process it
-            actv'' = if safety == Unsafe then mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt' else actv'
+        -- (c) this stm is also a potential source for coalescing, so process it
+        actv'' <- if safety == Unsafe then mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'} else pure actv'
         pure $
           bu_env {activeCoals = actv'', inhibit = inhbt'}
 
 -- The case of flat in-place update:
 --   @let x' = x with flat-slice <- elm@
-mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (FlatUpdate x _ _elm))) td_env bu_env
+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (FlatUpdate x _ _elm))) td_env bu_env
   | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =
       do
         -- (a) filter by the 3rd safety for @elm@ and @x'@
@@ -1102,18 +1152,25 @@
                         _ ->
                           markFailedCoal (actv, inhbt) m_x
 
-            -- (c) this stm is also a potential source for coalescing, so process it
-            actv'' = mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt'
+        -- (c) this stm is also a potential source for coalescing, so process it
+        actv'' <- mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'}
         pure $
           bu_env {activeCoals = actv'', inhibit = inhbt'}
 --
 mkCoalsTabStm _ (Let pat _ (BasicOp Update {})) _ _ =
   error $ "In ArrayCoalescing.hs, fun mkCoalsTabStm, illegal pattern for in-place update: " ++ show pat
 -- default handling
-mkCoalsTabStm lutab (Let pat _ (Op op)) td_env bu_env = do
+mkCoalsTabStm lutab stm@(Let pat _ (Op op)) td_env bu_env = do
   -- Process body
   on_op <- asks onOp
-  on_op lutab pat op td_env bu_env
+  activeCoals' <-
+    mkCoalsHelper3PatternMatch
+      stm
+      lutab
+      td_env
+      bu_env
+  let bu_env' = bu_env {activeCoals = activeCoals'}
+  on_op lutab pat op td_env bu_env'
 mkCoalsTabStm lutab stm@(Let pat _ e) td_env bu_env = do
   --   i) Filter @activeCoals@ by the 3rd safety condition:
   --      this is now relaxed by use of LMAD eqs:
@@ -1132,8 +1189,8 @@
       ((activeCoals'', inhibit''), successCoals') =
         foldl (foldfun safe_4) ((activeCoals', inhibit'), successCoals bu_env) (getArrMemAssoc pat)
 
-      -- iii) record a potentially coalesced statement in @activeCoals@
-      activeCoals''' = mkCoalsHelper3PatternMatch pat e lutab td_env successCoals' activeCoals'' (inhibited td_env)
+  -- iii) record a potentially coalesced statement in @activeCoals@
+  activeCoals''' <- mkCoalsHelper3PatternMatch stm lutab td_env bu_env {successCoals = successCoals', activeCoals = activeCoals''}
   pure bu_env {activeCoals = activeCoals''', inhibit = inhibit'', successCoals = successCoals'}
   where
     foldfun safe_4 ((a_acc, inhb), s_acc) (b, MemBlock tp shp mb _b_indfun) =
@@ -1244,22 +1301,20 @@
 -- |   Pattern matches a potentially coalesced statement and
 --     records a new association in @activeCoals@
 mkCoalsHelper3PatternMatch ::
-  HasMemBlock (Aliases rep) =>
-  Pat (VarAliases, LetDecMem) ->
-  Exp (Aliases rep) ->
+  Coalesceable rep inner =>
+  Stm (Aliases rep) ->
   LUTabFun ->
   TopdownEnv rep ->
-  CoalsTab ->
-  CoalsTab ->
-  InhibitTab ->
-  CoalsTab
-mkCoalsHelper3PatternMatch pat e lutab td_env _ activeCoals_tab _
-  | Nothing <- genCoalStmtInfo lutab (scope td_env) pat e =
-      activeCoals_tab
-mkCoalsHelper3PatternMatch pat e lutab td_env successCoals_tab activeCoals_tab inhibit_tab
-  | Just clst <- genCoalStmtInfo lutab (scope td_env) pat e =
-      foldl processNewCoalesce activeCoals_tab clst
+  BotUpEnv ->
+  ShortCircuitM rep CoalsTab
+mkCoalsHelper3PatternMatch stm lutab td_env bu_env = do
+  clst <- genCoalStmtInfo lutab td_env (scope td_env) stm
+  case clst of
+    Nothing -> pure activeCoals_tab
+    Just clst' -> pure $ foldl processNewCoalesce activeCoals_tab clst'
   where
+    successCoals_tab = successCoals bu_env
+    activeCoals_tab = activeCoals bu_env
     processNewCoalesce acc (knd, alias_fn, x, m_x, ind_x, b, m_b, _, tp_b, shp_b) =
       -- test whether we are in a transitive coalesced case, i.e.,
       --      @let b = scratch ...@
@@ -1286,18 +1341,18 @@
                  in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0)
           success0 = IxFun.hasOneLmad ind_yx
           m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx
-       in case (success0, not m_b_aliased_m_yx, isInScope td_env m_yx) of -- nameIn m_yx (alloc td_env)
-            (True, True, True) ->
-              -- Finally update the @activeCoals@ table with a fresh
-              --   binding for @m_b@; if such one exists then overwrite.
-              -- Also, add all variables from the alias chain of @b@ to
-              --   @vartab@, for example, in the case of a sequence:
-              --   @ b0 = if cond then ... else ... @
-              --   @ b1 = alias0 b0 @
-              --   @ b  = alias1 b1 @
-              --   @ x[j] = b @
-              -- Then @b1@ and @b0@ should also be added to @vartab@ if
-              --   @alias1@ and @alias0@ are invertible, otherwise fail early!
+       in if success0 && not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)
+      -- Finally update the @activeCoals@ table with a fresh
+      --   binding for @m_b@; if such one exists then overwrite.
+      -- Also, add all variables from the alias chain of @b@ to
+      --   @vartab@, for example, in the case of a sequence:
+      --   @ b0 = if cond then ... else ... @
+      --   @ b1 = alias0 b0 @
+      --   @ b  = alias1 b1 @
+      --   @ x[j] = b @
+      -- Then @b1@ and @b0@ should also be added to @vartab@ if
+      --   @alias1@ and @alias0@ are invertible, otherwise fail early!
+            then
               let mem_info = Coalesced knd (MemBlock tp_b shp_b m_yx ind_yx) M.empty
                   opts' =
                     if m_yx == m_x
@@ -1306,7 +1361,7 @@
                   vtab = M.singleton b mem_info
                   mvtab = addInvAliassesVarTab td_env vtab b
 
-                  is_inhibited = case M.lookup m_b inhibit_tab of
+                  is_inhibited = case M.lookup m_b $ inhibited td_env of
                     Just nms -> m_yx `nameIn` nms
                     Nothing -> False
                in case (is_inhibited, mvtab) of
@@ -1323,30 +1378,139 @@
                               opts'
                               mempty
                        in M.insert m_b coal_etry acc
-            _ -> acc
-mkCoalsHelper3PatternMatch _ _ _ _ _ _ _ =
-  error "In ArrayCoalescing.hs, fun mkCoalsHelper3PatternMatch: Unreachable!!!"
+            else acc
 
-genCoalStmtInfo ::
-  HasMemBlock (Aliases rep) =>
+-- | Information about a particular short-circuit point
+type SSPointInfo =
+  ( CoalescedKind,
+    IxFun -> IxFun,
+    VName,
+    VName,
+    IxFun,
+    VName,
+    VName,
+    IxFun,
+    PrimType,
+    Shape
+  )
+
+-- | Given an op, return a list of potential short-circuit points
+type GenSSPoint rep op =
   LUTabFun ->
+  TopdownEnv rep ->
   ScopeTab rep ->
   Pat (VarAliases, LetDecMem) ->
-  Exp (Aliases rep) ->
-  Maybe [(CoalescedKind, IxFun -> IxFun, VName, VName, IxFun, VName, VName, IxFun, PrimType, Shape)]
+  op ->
+  Maybe [SSPointInfo]
+
+genSSPointInfoSeqMem ::
+  GenSSPoint SeqMem (Op (Aliases SeqMem))
+genSSPointInfoSeqMem _ _ _ _ _ =
+  Nothing
+
+-- | For 'SegOp', we currently only handle 'SegMap', and only under the following
+-- circumstances:
+--
+--  1. The 'SegMap' has only one return/pattern value.
+--
+--  2. The 'KernelBody' contains an 'Index' statement that is indexing an array using
+--  only the values from the 'SegSpace'.
+--
+--  3. The array being indexed is last-used in that statement, is free in the
+--  'SegMap', is unique or has been recently allocated (specifically, it should
+--  not be a non-unique argument to the enclosing function), has elements with
+--  the same bit-size as the pattern elements, and has the exact same 'IxFun' as
+--  the pattern of the 'SegMap' statement.
+--
+-- There can be multiple candidate arrays, but the current implementation will
+-- always just try the first one.
+--
+-- The first restriction could be relaxed by trying to match up arrays in the
+-- 'KernelBody' with patterns of the 'SegMap', but the current implementation
+-- should be enough to handle many common cases.
+--
+-- The result of the 'SegMap' is treated as the destination, while the candidate
+-- array from inside the body is treated as the source.
+genSSPointInfoSegOp ::
+  Coalesceable rep inner => GenSSPoint rep (SegOp lvl (Aliases rep))
+genSSPointInfoSegOp
+  lutab
+  td_env
+  scopetab
+  (Pat [PatElem dst (_, MemArray dst_pt _ _ (ArrayIn dst_mem dst_ixf))])
+  (SegMap _ space _ kernel_body)
+    | (src, MemBlock _ shp src_mem src_ixf) : _ <-
+        mapMaybe getPotentialMapShortCircuit $
+          stmsToList $
+            kernelBodyStms kernel_body =
+        Just [(MapCoal, id, dst, dst_mem, dst_ixf, src, src_mem, src_ixf, dst_pt, shp)]
+    where
+      iterators = map fst $ unSegSpace space
+      frees = freeIn kernel_body
+
+      getPotentialMapShortCircuit (Let (Pat [PatElem x _]) _ (BasicOp (Index src slc)))
+        | Just inds <- sliceIndices slc,
+          L.sort inds == L.sort (map Var iterators),
+          Just last_uses <- M.lookup x lutab,
+          src `nameIn` last_uses,
+          Just memblock@(MemBlock src_pt _ src_mem src_ixf) <-
+            getScopeMemInfo src scopetab,
+          src_mem `nameIn` last_uses,
+          -- The 'alloc' table contains allocated memory blocks, including
+          -- unique memory blocks from the enclosing function. It does _not_
+          -- include non-unique memory blocks from the enclosing function.
+          src_mem `M.member` alloc td_env,
+          src `nameIn` frees,
+          src_ixf == dst_ixf,
+          primBitSize src_pt == primBitSize dst_pt =
+            Just (src, memblock)
+      getPotentialMapShortCircuit _ = Nothing
+genSSPointInfoSegOp _ _ _ _ _ =
+  Nothing
+
+genSSPointInfoMemOp ::
+  GenSSPoint rep inner ->
+  GenSSPoint rep (MemOp inner)
+genSSPointInfoMemOp onOp lutab td_end scopetab pat (Inner op) =
+  onOp lutab td_end scopetab pat op
+genSSPointInfoMemOp _ _ _ _ _ _ = Nothing
+
+genSSPointInfoGPUMem ::
+  GenSSPoint GPUMem (Op (Aliases GPUMem))
+genSSPointInfoGPUMem = genSSPointInfoMemOp f
+  where
+    f lutab td_env scopetab pat (GPU.SegOp op) =
+      genSSPointInfoSegOp lutab td_env scopetab pat op
+    f _ _ _ _ _ = Nothing
+
+genSSPointInfoMCMem ::
+  GenSSPoint MCMem (Op (Aliases MCMem))
+genSSPointInfoMCMem = genSSPointInfoMemOp f
+  where
+    f lutab td_env scopetab pat (MC.ParOp Nothing op) =
+      genSSPointInfoSegOp lutab td_env scopetab pat op
+    f _ _ _ _ _ = Nothing
+
+genCoalStmtInfo ::
+  Coalesceable rep inner =>
+  LUTabFun ->
+  TopdownEnv rep ->
+  ScopeTab rep ->
+  Stm (Aliases rep) ->
+  ShortCircuitM rep (Maybe [SSPointInfo])
 -- CASE a) @let x <- copy(b^{lu})@
-genCoalStmtInfo lutab scopetab pat (BasicOp (Copy b))
+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Copy b)))
   | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =
-      case (M.lookup x lutab, getScopeMemInfo b scopetab) of
+      pure $ case (M.lookup x lutab, getScopeMemInfo b scopetab) of
         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->
           if b `notNameIn` last_uses
             then Nothing
             else Just [(CopyCoal, id, x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)]
         _ -> Nothing
 -- CASE c) @let x[i] = b^{lu}@
-genCoalStmtInfo lutab scopetab pat (BasicOp (Update _ x slice_x (Var b)))
+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Update _ x slice_x (Var b))))
   | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =
-      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of
+      pure $ case (M.lookup x' lutab, getScopeMemInfo b scopetab) of
         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->
           if b `notNameIn` last_uses
             then Nothing
@@ -1357,9 +1521,9 @@
     updateIndFunSlice ind_fun slc_x =
       let slc_x' = map (fmap pe64) $ unSlice slc_x
        in IxFun.slice ind_fun $ Slice slc_x'
-genCoalStmtInfo lutab scopetab pat (BasicOp (FlatUpdate x slice_x b))
+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (FlatUpdate x slice_x b)))
   | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =
-      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of
+      pure $ case (M.lookup x' lutab, getScopeMemInfo b scopetab) of
         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->
           if b `notNameIn` last_uses
             then Nothing
@@ -1371,9 +1535,9 @@
       IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims
 
 -- CASE b) @let x = concat(a, b^{lu})@
-genCoalStmtInfo lutab scopetab pat (BasicOp (Concat concat_dim (b0 :| bs) _))
+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Concat concat_dim (b0 :| bs) _)))
   | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =
-      case M.lookup x lutab of
+      pure $ case M.lookup x lutab of
         Nothing -> Nothing
         Just last_uses ->
           let zero = pe64 $ intConst Int64 0
@@ -1399,8 +1563,13 @@
                     _ -> (acc, offs, False)
               (res, _, _) = foldl markConcatParts ([], zero, True) (b0 : bs)
            in if null res then Nothing else Just res
--- CASE other than a), b), or c) not supported
-genCoalStmtInfo _ _ _ _ = Nothing
+-- case d) short-circuit points from ops. For instance, the result of a segmap
+-- can be considered a short-circuit point.
+genCoalStmtInfo lutab td_env scopetab (Let pat _ (Op op)) = do
+  ss_op <- asks ssPointFromOp
+  pure $ ss_op lutab td_env scopetab pat op
+-- CASE other than a), b), c), or d) not supported
+genCoalStmtInfo _ _ _ _ = pure Nothing
 
 data MemBodyResult = MemBodyResult
   { patMem :: VName,
@@ -1497,7 +1666,7 @@
     mki64subst _ = Nothing
 
 computeScalarTable ::
-  (Coalesceable rep inner) =>
+  Coalesceable rep inner =>
   ScopeTab rep ->
   Stm (Aliases rep) ->
   ScalarTableM rep (M.Map VName (PrimExp VName))
@@ -1531,18 +1700,44 @@
   on_op scope_table op
 computeScalarTable _ _ = pure mempty
 
-computeScalarTableGPUMem :: ScopeTab GPUMem -> Op (Aliases GPUMem) -> ScalarTableM GPUMem (M.Map VName (PrimExp VName))
-computeScalarTableGPUMem _ (Alloc _ _) = pure mempty
-computeScalarTableGPUMem scope_table (Inner (SegOp segop)) = do
+type ComputeScalarTable rep op =
+  ScopeTab rep -> op -> ScalarTableM rep (M.Map VName (PrimExp VName))
+
+computeScalarTableMemOp ::
+  ComputeScalarTable rep inner -> ComputeScalarTable rep (MemOp inner)
+computeScalarTableMemOp _ _ (Alloc _ _) = pure mempty
+computeScalarTableMemOp onInner scope_table (Inner op) = onInner scope_table op
+
+computeScalarTableSegOp ::
+  Coalesceable rep inner =>
+  ComputeScalarTable rep (GPU.SegOp lvl (Aliases rep))
+computeScalarTableSegOp scope_table segop = do
   concatMapM
-    (computeScalarTable $ scope_table <> scopeOf (kernelBodyStms $ segBody segop) <> scopeOfSegSpace (segSpace segop))
+    ( computeScalarTable $
+        scope_table
+          <> scopeOf (kernelBodyStms $ segBody segop)
+          <> scopeOfSegSpace (segSpace segop)
+    )
     (stmsToList $ kernelBodyStms $ segBody segop)
-computeScalarTableGPUMem _ (Inner (SizeOp _)) = pure mempty
-computeScalarTableGPUMem _ (Inner (OtherOp ())) = pure mempty
-computeScalarTableGPUMem scope_table (Inner (GPUBody _ body)) =
+
+computeScalarTableGPUMem ::
+  ComputeScalarTable GPUMem (GPU.HostOp (Aliases GPUMem) ())
+computeScalarTableGPUMem scope_table (GPU.SegOp segop) =
+  computeScalarTableSegOp scope_table segop
+computeScalarTableGPUMem _ (GPU.SizeOp _) = pure mempty
+computeScalarTableGPUMem _ (GPU.OtherOp ()) = pure mempty
+computeScalarTableGPUMem scope_table (GPU.GPUBody _ body) =
   concatMapM
     (computeScalarTable $ scope_table <> scopeOf (bodyStms body))
     (stmsToList $ bodyStms body)
+
+computeScalarTableMCMem ::
+  ComputeScalarTable MCMem (MC.MCOp (Aliases MCMem) ())
+computeScalarTableMCMem _ (MC.OtherOp ()) = pure mempty
+computeScalarTableMCMem scope_table (MC.ParOp par_op segop) =
+  (<>)
+    <$> maybe (pure mempty) (computeScalarTableSegOp scope_table) par_op
+    <*> computeScalarTableSegOp scope_table segop
 
 filterMapM1 :: (Eq k, Monad m) => (v -> m Bool) -> M.Map k v -> m (M.Map k v)
 filterMapM1 f m = fmap M.fromAscList $ filterM (f . snd) $ M.toAscList m
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
@@ -8,11 +8,7 @@
     CoalescedKind (..),
     ArrayMemBound (..),
     AllocTab,
-    AliasTab,
-    LUTabFun,
-    CreatesNewArrOp,
     HasMemBlock,
-    LUTabPrg,
     ScalarTab,
     CoalsTab,
     ScopeTab,
@@ -42,7 +38,8 @@
 import Data.Maybe
 import Data.Set qualified as S
 import Futhark.IR.Aliases
-import Futhark.IR.GPUMem
+import Futhark.IR.GPUMem as GPU
+import Futhark.IR.MCMem as MC
 import Futhark.IR.Mem.IxFun qualified as IxFun
 import Futhark.IR.SeqMem
 import Futhark.Util.Pretty hiding (line, sep, (</>))
@@ -106,6 +103,7 @@
     ConcatCoal
   | -- | transitive, i.e., other variables aliased with b.
     TransitiveCoal
+  | MapCoal
 
 -- | Information about a memory block: type, shape, name and ixfun.
 data ArrayMemBound = MemBlock
@@ -169,15 +167,6 @@
 type AllocTab = M.Map VName Space
 -- ^ the allocatted memory blocks
 
-type AliasTab = M.Map VName Names
--- ^ maps a variable or memory block to its aliases
-
-type LUTabFun = M.Map VName Names
--- ^ maps a name indentifying a stmt to the last uses in that stmt
-
-type LUTabPrg = M.Map Name LUTabFun
--- ^ maps function names to last-use tables
-
 type ScalarTab = M.Map VName (PrimExp VName)
 -- ^ maps a variable name to its PrimExp scalar expression
 
@@ -219,6 +208,7 @@
   pretty InPlaceCoal = "InPlace"
   pretty ConcatCoal = "Concat"
   pretty TransitiveCoal = "Transitive"
+  pretty MapCoal = "Map"
 
 instance Pretty ArrayMemBound where
   pretty (MemBlock ptp shp m_nm ixfn) =
@@ -324,8 +314,16 @@
       Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)
       _ -> Nothing
 
+instance HasMemBlock (Aliases MCMem) where
+  getScopeMemInfo r scope_env0 =
+    case M.lookup r scope_env0 of
+      Just (LetName (_, MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)
+      Just (FParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)
+      Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)
+      _ -> Nothing
+
 -- | @True@ if the expression returns a "fresh" array.
-createsNewArrOK :: CreatesNewArrOp (Op rep) => Exp rep -> Bool
+createsNewArrOK :: Exp rep -> Bool
 createsNewArrOK (BasicOp Replicate {}) = True
 createsNewArrOK (BasicOp Iota {}) = True
 createsNewArrOK (BasicOp Manifest {}) = True
@@ -334,28 +332,7 @@
 createsNewArrOK (BasicOp ArrayLit {}) = True
 createsNewArrOK (BasicOp Scratch {}) = True
 createsNewArrOK (BasicOp Rotate {}) = True
-createsNewArrOK (Op op) = createsNewArrOp op
 createsNewArrOK _ = False
-
-class CreatesNewArrOp rep where
-  createsNewArrOp :: rep -> Bool
-
-instance CreatesNewArrOp () where
-  createsNewArrOp () = False
-
-instance CreatesNewArrOp inner => CreatesNewArrOp (MemOp inner) where
-  createsNewArrOp (Alloc _ _) = True
-  createsNewArrOp (Inner inner) = createsNewArrOp inner
-
-instance CreatesNewArrOp inner => CreatesNewArrOp (HostOp (Aliases GPUMem) inner) where
-  createsNewArrOp (OtherOp op) = createsNewArrOp op
-  createsNewArrOp (SegOp (SegMap _ _ _ kbody)) = all isReturns $ kernelBodyResult kbody
-  createsNewArrOp (SizeOp _) = False
-  createsNewArrOp _ = undefined
-
-isReturns :: KernelResult -> Bool
-isReturns Returns {} = True
-isReturns _ = False
 
 -- | Memory-block removal from active-coalescing table
 --   should only be handled via this function, it is easy
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs b/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs
@@ -1,6 +1,5 @@
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 -- | Last use analysis for array short circuiting
 --
@@ -15,57 +14,122 @@
 -- This pass is different from "Futhark.Analysis.LastUse" in that memory blocks
 -- are used to alias arrays. For instance, an 'Update' will not result in a last
 -- use of the array being updated, because the result lives in the same memory.
-module Futhark.Optimise.ArrayShortCircuiting.LastUse (lastUseSeqMem, lastUsePrg, lastUsePrgGPU, lastUseGPUMem) where
+module Futhark.Optimise.ArrayShortCircuiting.LastUse
+  ( lastUseSeqMem,
+    lastUseGPUMem,
+    lastUseMCMem,
+    LUTabFun,
+    LUTabProg,
+  )
+where
 
 import Control.Monad.Reader
 import Control.Monad.State.Strict
 import Data.Bifunctor (bimap)
+import Data.Function ((&))
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Data.Sequence (Seq (..))
 import Futhark.IR.Aliases
 import Futhark.IR.GPUMem
+import Futhark.IR.GPUMem qualified as GPU
+import Futhark.IR.MCMem
+import Futhark.IR.MCMem qualified as MC
 import Futhark.IR.SeqMem
 import Futhark.Optimise.ArrayShortCircuiting.DataStructs
 import Futhark.Util
 
+-- | Maps a name indentifying a Stm to the last uses in that Stm.
+type LUTabFun = M.Map VName Names
+
+-- | LU-table for the constants, and for each function.
+type LUTabProg = (LUTabFun, M.Map Name LUTabFun)
+
+type LastUseOp rep = Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
+
 -- | 'LastUseReader' allows us to abstract over representations by supplying the
 -- 'onOp' function.
-newtype LastUseReader rep = LastUseReader
-  { onOp :: Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)
+data LastUseReader rep = LastUseReader
+  { onOp :: LastUseOp rep,
+    scope :: Scope (Aliases rep)
   }
 
-type LastUseM rep a = StateT AliasTab (Reader (LastUseReader rep)) a
+-- | Maps a variable or memory block to its aliases.
+type AliasTab = M.Map VName Names
 
+newtype LastUseM rep a = LastUseM (StateT AliasTab (Reader (LastUseReader rep)) a)
+  deriving
+    ( Monad,
+      Functor,
+      Applicative,
+      MonadReader (LastUseReader rep),
+      MonadState AliasTab
+    )
+
+instance
+  (RepTypes rep, CanBeAliased (Op rep)) =>
+  HasScope (Aliases rep) (LastUseM rep)
+  where
+  askScope = asks scope
+
+instance
+  (RepTypes rep, CanBeAliased (Op rep)) =>
+  LocalScope (Aliases rep) (LastUseM rep)
+  where
+  localScope sc (LastUseM m) = LastUseM $ do
+    local (\rd -> rd {scope = scope rd <> sc}) m
+
+type Constraints rep =
+  ( LocalScope (Aliases rep) (LastUseM rep),
+    ASTRep rep,
+    FreeIn (OpWithAliases (Op rep)),
+    HasMemBlock (Aliases rep),
+    CanBeAliased (Op rep)
+  )
+
+runLastUseM :: LastUseOp rep -> LastUseM rep a -> a
+runLastUseM onOp (LastUseM m) =
+  runReader (evalStateT m mempty) (LastUseReader onOp mempty)
+
 aliasLookup :: VName -> LastUseM rep Names
 aliasLookup vname =
   gets $ fromMaybe mempty . M.lookup vname
 
+lastUseProg ::
+  Constraints rep =>
+  Prog (Aliases rep) ->
+  LastUseM rep LUTabProg
+lastUseProg prog =
+  let bound_in_consts =
+        progConsts prog
+          & concatMap (patNames . stmPat)
+          & namesFromList
+      consts = progConsts prog
+      funs = progFuns prog
+   in inScopeOf consts $ do
+        (consts_lu, _) <- lastUseStms consts mempty mempty
+        lus <- mapM (lastUseFun bound_in_consts) funs
+        pure (consts_lu, M.fromList $ zip (map funDefName funs) lus)
+
+lastUseFun ::
+  Constraints rep =>
+  Names ->
+  FunDef (Aliases rep) ->
+  LastUseM rep LUTabFun
+lastUseFun bound_in_consts f =
+  inScopeOf f $ fst <$> lastUseBody (funDefBody f) (mempty, bound_in_consts)
+
 -- | Perform last-use analysis on a 'Prog' in 'SeqMem'
-lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg
-lastUsePrg prg = M.fromList $ map lastUseSeqMem $ progFuns prg
+lastUseSeqMem :: Prog (Aliases SeqMem) -> LUTabProg
+lastUseSeqMem = runLastUseM lastUseSeqOp . lastUseProg
 
 -- | Perform last-use analysis on a 'Prog' in 'GPUMem'
-lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg
-lastUsePrgGPU prg = M.fromList $ map lastUseGPUMem $ progFuns prg
-
--- | Perform last-use analysis on a 'FunDef' in 'SeqMem'
-lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, LUTabFun)
-lastUseSeqMem (FunDef _ _ fname _ _ body) =
-  let (res, _) =
-        runReader
-          (evalStateT (lastUseBody body (mempty, mempty)) mempty)
-          (LastUseReader lastUseSeqOp)
-   in (fname, res)
+lastUseGPUMem :: Prog (Aliases GPUMem) -> LUTabProg
+lastUseGPUMem = runLastUseM (lastUseMemOp lastUseGPUOp) . lastUseProg
 
--- | Perform last-use analysis on a 'FunDef' in 'GPUMem'
-lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, LUTabFun)
-lastUseGPUMem (FunDef _ _ fname _ _ body) =
-  let (res, _) =
-        runReader
-          (evalStateT (lastUseBody body (mempty, mempty)) mempty)
-          (LastUseReader lastUseGPUOp)
-   in (fname, res)
+-- | Perform last-use analysis on a 'Prog' in 'MCMem'
+lastUseMCMem :: Prog (Aliases MCMem) -> LUTabProg
+lastUseMCMem = runLastUseM (lastUseMemOp lastUseMCOp) . lastUseProg
 
 -- | Performing the last-use analysis on a body.
 --
@@ -74,7 +138,7 @@
 -- difference between the free-variables in that stmt and the set of variables
 -- known to be used after that statement.
 lastUseBody ::
-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
+  Constraints rep =>
   -- | The body of statements
   Body (Aliases rep) ->
   -- | The current last-use table, tupled with the known set of already used names
@@ -83,18 +147,19 @@
   --      (i) an updated last-use table,
   --     (ii) an updated set of used names (including the binding).
   LastUseM rep (LUTabFun, Names)
-lastUseBody bdy@(Body _ stms result) (lutab, used_nms) = do
+lastUseBody bdy@(Body _ stms result) (lutab, used_nms) =
   -- perform analysis bottom-up in bindings: results are known to be used,
   -- hence they are added to the used_nms set.
-  (lutab', _) <-
-    lastUseStms stms (lutab, used_nms) $
-      namesToList $
-        freeIn $
-          map resSubExp result
-  -- Clean up the used names by recomputing the aliasing transitive-closure
-  -- of the free names in body based on the current alias table @alstab@.
-  used_in_body <- aliasTransitiveClosure $ freeIn bdy
-  pure (lutab', used_nms <> used_in_body)
+  inScopeOf stms $ do
+    (lutab', _) <-
+      lastUseStms stms (lutab, used_nms) $
+        namesToList $
+          freeIn $
+            map resSubExp result
+    -- Clean up the used names by recomputing the aliasing transitive-closure
+    -- of the free names in body based on the current alias table @alstab@.
+    used_in_body <- aliasTransitiveClosure $ freeIn bdy
+    pure (lutab', used_nms <> used_in_body)
 
 -- | Performing the last-use analysis on a body.
 --
@@ -103,7 +168,7 @@
 -- difference between the free-variables in that stmt and the set of variables
 -- known to be used after that statement.
 lastUseKernelBody ::
-  (CanBeAliased (Op rep), ASTRep rep) =>
+  Constraints rep =>
   -- | The body of statements
   KernelBody (Aliases rep) ->
   -- | The current last-use table, tupled with the known set of already used names
@@ -112,62 +177,71 @@
   --      (i) an updated last-use table,
   --     (ii) an updated set of used names (including the binding).
   LastUseM rep (LUTabFun, Names)
-lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) = do
-  -- perform analysis bottom-up in bindings: results are known to be used,
-  -- hence they are added to the used_nms set.
-  (lutab', _) <-
-    lastUseStms stms (lutab, used_nms) $ namesToList $ freeIn result
-  -- Clean up the used names by recomputing the aliasing transitive-closure
-  -- of the free names in body based on the current alias table @alstab@.
-  used_in_body <- aliasTransitiveClosure $ freeIn bdy
-  pure (lutab', used_nms <> used_in_body)
+lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) =
+  inScopeOf stms $ do
+    -- perform analysis bottom-up in bindings: results are known to be used,
+    -- hence they are added to the used_nms set.
+    (lutab', _) <-
+      lastUseStms stms (lutab, used_nms) $ namesToList $ freeIn result
+    -- Clean up the used names by recomputing the aliasing transitive-closure
+    -- of the free names in body based on the current alias table @alstab@.
+    used_in_body <- aliasTransitiveClosure $ freeIn bdy
+    pure (lutab', used_nms <> used_in_body)
 
 lastUseStms ::
-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
+  Constraints rep =>
   Stms (Aliases rep) ->
   (LUTabFun, Names) ->
   [VName] ->
   LastUseM rep (LUTabFun, Names)
 lastUseStms Empty (lutab, nms) res_nms = do
   aliases <- concatMapM aliasLookup res_nms
-  pure (lutab, nms <> aliases)
-lastUseStms (stm@(Let pat _ e) :<| stms) (lutab, nms) res_nms = do
-  let extra_alias = case e of
-        BasicOp (Update _ old _ _) -> oneName old
-        BasicOp (FlatUpdate old _ _) -> oneName old
-        _ -> mempty
-  -- We build up aliases top-down
-  updateAliasing extra_alias pat
-  -- But compute last use bottom-up
-  (lutab', nms') <- lastUseStms stms (lutab, nms) res_nms
-  (lutab'', nms'') <- lastUseStm stm (lutab', nms')
-  pure (lutab'', nms'')
+  pure (lutab, nms <> aliases <> namesFromList res_nms)
+lastUseStms (stm@(Let pat _ e) :<| stms) (lutab, nms) res_nms =
+  inScopeOf stm $ do
+    let extra_alias = case e of
+          BasicOp (Update _ old _ _) -> oneName old
+          BasicOp (FlatUpdate old _ _) -> oneName old
+          _ -> mempty
+    -- We build up aliases top-down
+    updateAliasing extra_alias pat
+    -- But compute last use bottom-up
+    (lutab', nms') <- lastUseStms stms (lutab, nms) res_nms
+    (lutab'', nms'') <- lastUseStm stm (lutab', nms')
+    pure (lutab'', nms'')
 
 lastUseStm ::
-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
+  Constraints rep =>
   Stm (Aliases rep) ->
   (LUTabFun, Names) ->
   LastUseM rep (LUTabFun, Names)
-lastUseStm (Let pat _ e) (lutab, used_nms) =
-  do
-    -- analyse the expression and get the
-    --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)
-    -- (ii) the set of variables lastly used in the current binding.
-    -- (iii)  aliased transitive-closure of used names, and
-    (lutab', last_uses, used_nms') <- lastUseExp e used_nms
-    -- filter-out the binded names from the set of used variables,
-    -- since they go out of scope, and update the last-use table.
-    let patnms = patNames pat
-        used_nms'' = used_nms' `namesSubtract` namesFromList patnms
-        lutab'' =
-          M.union lutab' $ M.insert (head patnms) last_uses lutab
-    pure (lutab'', used_nms'')
+lastUseStm (Let pat _ e) (lutab, used_nms) = do
+  -- analyse the expression and get the
+  --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)
+  -- (ii) the set of variables lastly used in the current binding.
+  -- (iii)  aliased transitive-closure of used names, and
+  (lutab', last_uses, used_nms') <- lastUseExp e used_nms
+  sc <- asks scope
+  let lu_mems =
+        namesToList last_uses
+          & mapMaybe (`getScopeMemInfo` sc)
+          & map memName
+          & namesFromList
+          & flip namesSubtract used_nms
 
+  -- filter-out the binded names from the set of used variables,
+  -- since they go out of scope, and update the last-use table.
+  let patnms = patNames pat
+      used_nms'' = used_nms' `namesSubtract` namesFromList patnms
+      lutab'' =
+        M.union lutab' $ M.insert (head patnms) (last_uses <> lu_mems) lutab
+  pure (lutab'', used_nms'')
+
 --------------------------------
 
 -- | Last-Use Analysis for an expression.
 lastUseExp ::
-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>
+  Constraints rep =>
   -- | The expression to analyse
   Exp (Aliases rep) ->
   -- | The set of used names "after" this expression
@@ -191,7 +265,7 @@
   let used_nms' = used_cases <> body_used_nms
   (_, last_used_arrs) <- lastUsedInNames used_nms $ free_in_body <> free_in_cases
   pure (lutab_cases <> lutab', last_used_arrs, used_nms')
-lastUseExp (DoLoop var_ses _ body) used_nms0 = do
+lastUseExp (DoLoop var_ses lf body) used_nms0 = inScopeOf lf $ do
   free_in_body <- aliasTransitiveClosure $ freeIn body
   -- compute the aliasing transitive closure of initializers that are not last-uses
   var_inis <- catMaybes <$> mapM (initHelper (free_in_body <> used_nms0)) var_ses
@@ -230,56 +304,93 @@
   (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e
   pure (M.empty, lu_vars, used_nms')
 
-lastUseGPUOp :: Op (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
-lastUseGPUOp (Alloc se sp) used_nms = do
+lastUseMemOp ::
+  (inner -> Names -> LastUseM rep (LUTabFun, Names, Names)) ->
+  MemOp inner ->
+  Names ->
+  LastUseM rep (LUTabFun, Names, Names)
+lastUseMemOp _ (Alloc se sp) used_nms = do
   let free_in_e = freeIn se <> freeIn sp
   (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e
   pure (M.empty, lu_vars, used_nms')
-lastUseGPUOp (Inner (OtherOp ())) used_nms =
-  pure (mempty, mempty, used_nms)
-lastUseGPUOp (Inner (SizeOp sop)) used_nms = do
-  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn sop
-  pure (mempty, lu_vars, used_nms')
-lastUseGPUOp (Inner (SegOp (SegMap _ _ tps kbody))) used_nms = do
+lastUseMemOp onInner (Inner op) used_nms = onInner op used_nms
+
+lastUseSegOp ::
+  Constraints rep =>
+  SegOp lvl (Aliases rep) ->
+  Names ->
+  LastUseM rep (LUTabFun, Names, Names)
+lastUseSegOp (SegMap _ _ tps kbody) used_nms = do
   (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps
   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')
   pure (body_lutab, lu_vars, used_nms' <> used_nms'')
-lastUseGPUOp (Inner (SegOp (SegRed _ _ sbos tps kbody))) used_nms = do
+lastUseSegOp (SegRed _ _ sbos tps kbody) used_nms = do
   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms
   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps
   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')
   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')
-lastUseGPUOp (Inner (SegOp (SegScan _ _ sbos tps kbody))) used_nms = do
+lastUseSegOp (SegScan _ _ sbos tps kbody) used_nms = do
   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms
   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps
   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')
   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')
-lastUseGPUOp (Inner (SegOp (SegHist _ _ hos tps kbody))) used_nms = do
+lastUseSegOp (SegHist _ _ hos tps kbody) used_nms = do
   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseHistOp hos used_nms
   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps
   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')
   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')
-lastUseGPUOp (Inner (GPUBody tps body)) used_nms = do
+
+lastUseGPUOp :: HostOp (Aliases GPUMem) () -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
+lastUseGPUOp (GPU.OtherOp ()) used_nms =
+  pure (mempty, mempty, used_nms)
+lastUseGPUOp (SizeOp sop) used_nms = do
+  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn sop
+  pure (mempty, lu_vars, used_nms')
+lastUseGPUOp (GPUBody tps body) used_nms = do
   (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps
   (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')
   pure (body_lutab, lu_vars, used_nms' <> used_nms'')
+lastUseGPUOp (SegOp op) used_nms =
+  lastUseSegOp op used_nms
 
-lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
+lastUseMCOp :: MCOp (Aliases MCMem) () -> Names -> LastUseM MCMem (LUTabFun, Names, Names)
+lastUseMCOp (MC.OtherOp ()) used_nms =
+  pure (mempty, mempty, used_nms)
+lastUseMCOp (MC.ParOp par_op op) used_nms = do
+  (lutab_par_op, lu_vars_par_op, used_names_par_op) <-
+    maybe (pure mempty) (`lastUseSegOp` used_nms) par_op
+  (lutab_op, lu_vars_op, used_names_op) <-
+    lastUseSegOp op used_nms
+  pure
+    ( lutab_par_op <> lutab_op,
+      lu_vars_par_op <> lu_vars_op,
+      used_names_par_op <> used_names_op
+    )
+
+lastUseSegBinOp ::
+  Constraints rep =>
+  [SegBinOp (Aliases rep)] ->
+  Names ->
+  LastUseM rep (LUTabFun, Names, Names)
 lastUseSegBinOp sbos used_nms = do
   (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper sbos
   pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')
   where
-    helper (SegBinOp _ (Lambda _ body _) neutral shp) = do
+    helper (SegBinOp _ l@(Lambda _ body _) neutral shp) = inScopeOf l $ do
       (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn neutral <> freeIn shp
       (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')
       pure (body_lutab, lu_vars, used_nms'')
 
-lastUseHistOp :: [HistOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)
+lastUseHistOp ::
+  Constraints rep =>
+  [HistOp (Aliases rep)] ->
+  Names ->
+  LastUseM rep (LUTabFun, Names, Names)
 lastUseHistOp hos used_nms = do
   (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper hos
   pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')
   where
-    helper (HistOp shp rf dest neutral shp' (Lambda _ body _)) = do
+    helper (HistOp shp rf dest neutral shp' l@(Lambda _ body _)) = inScopeOf l $ do
       (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn shp <> freeIn rf <> freeIn dest <> freeIn neutral <> freeIn shp'
       (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')
       pure (body_lutab, lu_vars, used_nms'')
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
@@ -228,7 +228,7 @@
             if m_b `nameIn` original_mem_aliases
               then (wrt_lmads' <> lmads'', Set mempty)
               else (wrt_lmads', lmads'')
-          no_overlap = noMemOverlap td_env prev_use wrt_lmads''
+          no_overlap = noMemOverlap td_env (lmads <> prev_use) wrt_lmads''
           wrt_lmads =
             if no_overlap
               then Just wrt_lmads''
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
@@ -22,7 +22,8 @@
 import Data.Maybe
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.IR.Aliases
-import Futhark.IR.GPUMem
+import Futhark.IR.GPUMem as GPU
+import Futhark.IR.MCMem as MC
 import Futhark.IR.Mem.IxFun qualified as IxFun
 import Futhark.Optimise.ArrayShortCircuiting.DataStructs
 
@@ -131,19 +132,40 @@
 
   scopeHelper :: inner -> Scope rep
 
+instance TopDownHelper (SegOp lvl rep) where
+  innerNonNegatives _ op =
+    foldMap (oneName . fst) $ unSegSpace $ segSpace op
+
+  innerKnownLessThan op =
+    map (fmap $ primExpFromSubExp $ IntType Int64) $ unSegSpace $ segSpace op
+
+  scopeHelper op = scopeOfSegSpace $ segSpace op
+
 instance TopDownHelper (HostOp (Aliases GPUMem) ()) where
-  innerNonNegatives _ (SegOp seg_op) =
-    foldMap (oneName . fst) $ unSegSpace $ segSpace seg_op
+  innerNonNegatives vs (SegOp op) = innerNonNegatives vs op
   innerNonNegatives [vname] (SizeOp (GetSize _ _)) = oneName vname
   innerNonNegatives [vname] (SizeOp (GetSizeMax _)) = oneName vname
   innerNonNegatives _ _ = mempty
 
-  innerKnownLessThan (SegOp seg_op) =
-    map (fmap $ primExpFromSubExp $ IntType Int64) $ unSegSpace $ segSpace seg_op
+  innerKnownLessThan (SegOp op) = innerKnownLessThan op
   innerKnownLessThan _ = mempty
 
-  scopeHelper (SegOp seg_op) = scopeOfSegSpace $ segSpace seg_op
+  scopeHelper (SegOp op) = scopeHelper op
   scopeHelper _ = mempty
+
+instance TopDownHelper inner => TopDownHelper (MC.MCOp (Aliases MCMem) inner) where
+  innerNonNegatives vs (ParOp par_op op) =
+    maybe mempty (innerNonNegatives vs) par_op
+      <> innerNonNegatives vs op
+  innerNonNegatives vs (MC.OtherOp op) =
+    innerNonNegatives vs op
+  innerKnownLessThan (ParOp par_op op) =
+    maybe mempty innerKnownLessThan par_op <> innerKnownLessThan op
+  innerKnownLessThan (MC.OtherOp op) =
+    innerKnownLessThan op
+  scopeHelper (ParOp par_op op) =
+    maybe mempty scopeHelper par_op <> scopeHelper op
+  scopeHelper MC.OtherOp {} = mempty
 
 instance TopDownHelper () where
   innerNonNegatives _ () = mempty
diff --git a/src/Futhark/Pass/LiftAllocations.hs b/src/Futhark/Pass/LiftAllocations.hs
--- a/src/Futhark/Pass/LiftAllocations.hs
+++ b/src/Futhark/Pass/LiftAllocations.hs
@@ -6,11 +6,17 @@
 -- | This pass attempts to lift allocations as far towards the top in their body
 -- as possible. It does not try to hoist allocations outside across body
 -- boundaries.
-module Futhark.Pass.LiftAllocations (liftAllocationsSeqMem, liftAllocationsGPUMem) where
+module Futhark.Pass.LiftAllocations
+  ( liftAllocationsSeqMem,
+    liftAllocationsGPUMem,
+    liftAllocationsMCMem,
+  )
+where
 
 import Control.Monad.Reader
 import Data.Sequence (Seq (..))
 import Futhark.IR.GPUMem
+import Futhark.IR.MCMem
 import Futhark.IR.SeqMem
 import Futhark.Pass (Pass (..))
 
@@ -40,6 +46,19 @@
               progFuns
         }
 
+liftAllocationsMCMem :: Pass MCMem MCMem
+liftAllocationsMCMem =
+  Pass "lift allocations mc" "lift allocations mc" $ \prog@Prog {progFuns} ->
+    pure $
+      prog
+        { progFuns =
+            fmap
+              ( \f@FunDef {funDefBody} ->
+                  f {funDefBody = runReader (liftAllocationsInBody funDefBody) (Env liftAllocationsInMCOp)}
+              )
+              progFuns
+        }
+
 newtype Env inner = Env
   {onInner :: inner -> LiftM inner inner}
 
@@ -104,17 +123,28 @@
     then liftAllocationsInStms stms (stm :<| lifted) acc ((to_lift `namesSubtract` pat_names) <> freeIn stm)
     else liftAllocationsInStms stms lifted (stm :<| acc) to_lift
 
-liftAllocationsInHostOp :: HostOp GPUMem () -> LiftM (HostOp GPUMem ()) (HostOp GPUMem ())
-liftAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do
+liftAllocationsInSegOp ::
+  (Mem rep inner, LetDec rep ~ LetDecMem) =>
+  SegOp lvl rep ->
+  LiftM inner (SegOp lvl rep)
+liftAllocationsInSegOp (SegMap lvl sp tps body) = do
   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty
-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
-liftAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do
+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
+liftAllocationsInSegOp (SegRed lvl sp binops tps body) = do
   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty
-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
-liftAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do
+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
+liftAllocationsInSegOp (SegScan lvl sp binops tps body) = do
   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty
-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
-liftAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do
+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
+liftAllocationsInSegOp (SegHist lvl sp histops tps body) = do
   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty
-  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}
+  pure $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}
+
+liftAllocationsInHostOp :: HostOp GPUMem () -> LiftM (HostOp GPUMem ()) (HostOp GPUMem ())
+liftAllocationsInHostOp (SegOp op) = SegOp <$> liftAllocationsInSegOp op
 liftAllocationsInHostOp op = pure op
+
+liftAllocationsInMCOp :: MCOp MCMem () -> LiftM (MCOp MCMem ()) (MCOp MCMem ())
+liftAllocationsInMCOp (ParOp par op) =
+  ParOp <$> traverse liftAllocationsInSegOp par <*> liftAllocationsInSegOp op
+liftAllocationsInMCOp op = pure op
diff --git a/src/Futhark/Pass/LowerAllocations.hs b/src/Futhark/Pass/LowerAllocations.hs
--- a/src/Futhark/Pass/LowerAllocations.hs
+++ b/src/Futhark/Pass/LowerAllocations.hs
@@ -5,7 +5,12 @@
 
 -- | This pass attempts to lower allocations as far towards the bottom of their
 -- body as possible.
-module Futhark.Pass.LowerAllocations (lowerAllocationsSeqMem, lowerAllocationsGPUMem) where
+module Futhark.Pass.LowerAllocations
+  ( lowerAllocationsSeqMem,
+    lowerAllocationsGPUMem,
+    lowerAllocationsMCMem,
+  )
+where
 
 import Control.Monad.Reader
 import Data.Function ((&))
@@ -13,6 +18,7 @@
 import Data.Sequence (Seq (..))
 import Data.Sequence qualified as Seq
 import Futhark.IR.GPUMem
+import Futhark.IR.MCMem
 import Futhark.IR.SeqMem
 import Futhark.Pass (Pass (..))
 
@@ -42,6 +48,19 @@
               progFuns
         }
 
+lowerAllocationsMCMem :: Pass MCMem MCMem
+lowerAllocationsMCMem =
+  Pass "lower allocations mc" "lower allocations mc" $ \prog@Prog {progFuns} ->
+    pure $
+      prog
+        { progFuns =
+            fmap
+              ( \f@FunDef {funDefBody} ->
+                  f {funDefBody = runReader (lowerAllocationsInBody funDefBody) (Env lowerAllocationsInMCOp)}
+              )
+              progFuns
+        }
+
 newtype Env inner = Env
   {onInner :: inner -> LowerM inner inner}
 
@@ -98,17 +117,28 @@
       )
       (alloc, acc)
 
-lowerAllocationsInHostOp :: HostOp GPUMem () -> LowerM (HostOp GPUMem ()) (HostOp GPUMem ())
-lowerAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do
+lowerAllocationsInSegOp ::
+  (Mem rep inner, LetDec rep ~ LetDecMem) =>
+  SegOp lvl rep ->
+  LowerM inner (SegOp lvl rep)
+lowerAllocationsInSegOp (SegMap lvl sp tps body) = do
   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty
-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
-lowerAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do
+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}
+lowerAllocationsInSegOp (SegRed lvl sp binops tps body) = do
   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty
-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
-lowerAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do
+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}
+lowerAllocationsInSegOp (SegScan lvl sp binops tps body) = do
   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty
-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
-lowerAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do
+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}
+lowerAllocationsInSegOp (SegHist lvl sp histops tps body) = do
   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty
-  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}
+  pure $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}
+
+lowerAllocationsInHostOp :: HostOp GPUMem () -> LowerM (HostOp GPUMem ()) (HostOp GPUMem ())
+lowerAllocationsInHostOp (SegOp op) = SegOp <$> lowerAllocationsInSegOp op
 lowerAllocationsInHostOp op = pure op
+
+lowerAllocationsInMCOp :: MCOp MCMem () -> LowerM (MCOp MCMem ()) (MCOp MCMem ())
+lowerAllocationsInMCOp (ParOp par op) =
+  ParOp <$> traverse lowerAllocationsInSegOp par <*> lowerAllocationsInSegOp op
+lowerAllocationsInMCOp op = pure op
diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs
--- a/src/Futhark/Passes.hs
+++ b/src/Futhark/Passes.hs
@@ -195,5 +195,15 @@
         simplifyMCMem,
         entryPointMemMC,
         doubleBufferMC,
+        simplifyMCMem,
+        performCSE False,
+        LiftAllocations.liftAllocationsMCMem,
+        simplifyMCMem,
+        ArrayShortCircuiting.optimiseMCMem,
+        simplifyMCMem,
+        performCSE False,
+        simplifyMCMem,
+        LowerAllocations.lowerAllocationsMCMem,
+        performCSE False,
         simplifyMCMem
       ]
diff --git a/src/Futhark/Test/Spec.hs b/src/Futhark/Test/Spec.hs
--- a/src/Futhark/Test/Spec.hs
+++ b/src/Futhark/Test/Spec.hs
@@ -86,9 +86,11 @@
 -- | How a program can be transformed.
 data StructurePipeline
   = GpuPipeline
+  | MCPipeline
   | SOACSPipeline
   | SeqMemPipeline
   | GpuMemPipeline
+  | MCMemPipeline
   | NoPipeline
   deriving (Show)
 
@@ -342,6 +344,8 @@
   choice
     [ lexeme sep "gpu-mem" $> GpuMemPipeline,
       lexeme sep "gpu" $> GpuPipeline,
+      lexeme sep "mc-mem" $> MCMemPipeline,
+      lexeme sep "mc" $> MCPipeline,
       lexeme sep "seq-mem" $> SeqMemPipeline,
       lexeme sep "internalised" $> NoPipeline,
       pure SOACSPipeline
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
@@ -155,10 +155,10 @@
 prettyValue :: Value m -> Doc a
 prettyValue = prettyValueWith pprPrim
   where
-    pprPrim (UnsignedValue (Int8Value v)) = pretty v
-    pprPrim (UnsignedValue (Int16Value v)) = pretty v
-    pprPrim (UnsignedValue (Int32Value v)) = pretty v
-    pprPrim (UnsignedValue (Int64Value v)) = pretty v
+    pprPrim (UnsignedValue (Int8Value v)) = pretty (fromIntegral v :: Word8)
+    pprPrim (UnsignedValue (Int16Value v)) = pretty (fromIntegral v :: Word16)
+    pprPrim (UnsignedValue (Int32Value v)) = pretty (fromIntegral v :: Word32)
+    pprPrim (UnsignedValue (Int64Value v)) = pretty (fromIntegral v :: Word64)
     pprPrim (SignedValue (Int8Value v)) = pretty v
     pprPrim (SignedValue (Int16Value v)) = pretty v
     pprPrim (SignedValue (Int32Value v)) = pretty v
diff --git a/src/Language/Futhark/Primitive.hs b/src/Language/Futhark/Primitive.hs
--- a/src/Language/Futhark/Primitive.hs
+++ b/src/Language/Futhark/Primitive.hs
@@ -31,6 +31,7 @@
     PrimValue (..),
     primValueType,
     blankPrimValue,
+    onePrimValue,
 
     -- * Operations
     Overflow (..),
@@ -357,6 +358,19 @@
 blankPrimValue (FloatType Float64) = FloatValue $ Float64Value 0.0
 blankPrimValue Bool = BoolValue False
 blankPrimValue Unit = UnitValue
+
+-- | A one value of the given primitive type - this is one
+-- whatever is close to it.
+onePrimValue :: PrimType -> PrimValue
+onePrimValue (IntType Int8) = IntValue $ Int8Value 1
+onePrimValue (IntType Int16) = IntValue $ Int16Value 1
+onePrimValue (IntType Int32) = IntValue $ Int32Value 1
+onePrimValue (IntType Int64) = IntValue $ Int64Value 1
+onePrimValue (FloatType Float16) = FloatValue $ Float16Value 1.0
+onePrimValue (FloatType Float32) = FloatValue $ Float32Value 1.0
+onePrimValue (FloatType Float64) = FloatValue $ Float64Value 1.0
+onePrimValue Bool = BoolValue True
+onePrimValue Unit = UnitValue
 
 -- | Various unary operators.  It is a bit ad-hoc what is a unary
 -- operator and what is a built-in function.  Perhaps these should all
diff --git a/src/Language/Futhark/Primitive/Parse.hs b/src/Language/Futhark/Primitive/Parse.hs
--- a/src/Language/Futhark/Primitive/Parse.hs
+++ b/src/Language/Futhark/Primitive/Parse.hs
@@ -25,7 +25,7 @@
 
 -- | Is this character a valid member of an identifier?
 constituent :: Char -> Bool
-constituent c = isAlphaNum c || (c `elem` ("_/'+-=!&^.<>*|" :: String))
+constituent c = isAlphaNum c || (c `elem` ("_/'+-=!&^.<>*|%" :: String))
 
 -- | Consume whitespace (including skipping line comments).
 whitespace :: Parsec Void T.Text ()
diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs
--- a/src/Language/Futhark/Prop.hs
+++ b/src/Language/Futhark/Prop.hs
@@ -21,6 +21,7 @@
     prettyStacktrace,
     progHoles,
     defaultEntryPoint,
+    paramName,
 
     -- * Queries on expressions
     typeOf,
@@ -347,6 +348,11 @@
     (combineTypeShapes (Scalar et1) (Scalar et2) `setAliases` mempty)
 combineTypeShapes _ new_tp = new_tp
 
+-- | The name, if any.
+paramName :: PName -> Maybe VName
+paramName (Named v) = Just v
+paramName Unnamed = Nothing
+
 -- | Match the dimensions of otherwise assumed-equal types.  The
 -- combining function is also passed the names bound within the type
 -- (from named parameters or return types).
@@ -380,7 +386,7 @@
         ( Scalar (Arrow als1 p1 a1 (RetType dims1 b1)),
           Scalar (Arrow als2 p2 a2 (RetType dims2 b2))
           ) ->
-            let bound' = mapMaybe maybePName [p1, p2] <> dims1 <> dims2 <> bound
+            let bound' = mapMaybe paramName [p1, p2] <> dims1 <> dims2 <> bound
              in Scalar
                   <$> ( Arrow (als1 <> als2) p1
                           <$> matchDims' bound' a1 a2
@@ -397,9 +403,6 @@
     matchTypeArg bound (TypeArgDim x loc) (TypeArgDim y _) =
       TypeArgDim <$> onDims bound x y <*> pure loc
     matchTypeArg _ a _ = pure a
-
-    maybePName (Named v) = Just v
-    maybePName Unnamed = Nothing
 
     onShapes bound shape1 shape2 =
       Shape <$> zipWithM (onDims bound) (shapeDims shape1) (shapeDims shape2)
diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs
--- a/src/Language/Futhark/TypeChecker.hs
+++ b/src/Language/Futhark/TypeChecker.hs
@@ -642,17 +642,26 @@
   case entry' of
     Just _
       | any isTypeParam tparams' ->
-          typeError loc mempty "Entry point functions may not be polymorphic."
+          typeError loc mempty $
+            withIndexLink
+              "polymorphic-entry"
+              "Entry point functions may not be polymorphic."
       | not (all patternOrderZero params')
           || not (all orderZero rettype_params)
           || not (orderZero rettype') ->
-          typeError loc mempty "Entry point functions may not be higher-order."
+          typeError loc mempty $
+            withIndexLink
+              "higher-order-entry"
+              "Entry point functions may not be higher-order."
       | sizes_only_in_ret <-
           S.fromList (map typeParamName tparams')
             `S.intersection` freeInType rettype'
             `S.difference` foldMap freeInType (map patternStructType params' ++ rettype_params),
         not $ S.null sizes_only_in_ret ->
-          typeError loc mempty "Entry point functions must not be size-polymorphic in their return type."
+          typeError loc mempty $
+            withIndexLink
+              "size-polymorphic-entry"
+              "Entry point functions must not be size-polymorphic in their return type."
       | p : _ <- filter nastyParameter params' ->
           warn loc $
             "Entry point parameter\n"
diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs
--- a/src/Language/Futhark/TypeChecker/Terms.hs
+++ b/src/Language/Futhark/TypeChecker/Terms.hs
@@ -37,7 +37,7 @@
 overloadedTypeVars :: Constraints -> Names
 overloadedTypeVars = mconcat . map f . M.elems
   where
-    f (_, HasFields fs _) = mconcat $ map typeVars $ M.elems fs
+    f (_, HasFields _ fs _) = mconcat $ map typeVars $ M.elems fs
     f _ = mempty
 
 --- Basic checking
@@ -375,7 +375,7 @@
     zeroOrderType
       (mkUsage loc "returning value of this type from 'if' expression")
       "type returned from branch"
-      t'
+      (toStruct t')
 
     pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes t' retext)
   where
@@ -445,8 +445,10 @@
     t <- expType e'
     case anyConsumption e_occs of
       Just c ->
-        let msg = "type computed with consumption at " <> locText (location c)
-         in zeroOrderType (mkUsage loc "consumption in right-hand side of 'let'-binding") msg t
+        zeroOrderType
+          (mkUsage loc "consumption in right-hand side of 'let'-binding")
+          ("type computed with consumption at " <> locText (location c))
+          (toStruct t)
       _ -> pure ()
 
     incLevel . bindingSizes sizes $ \sizes' ->
@@ -499,8 +501,7 @@
             )
             (Info $ AppRes body_t ext)
 checkExp (AppExp (LetWith dest src slice ve body loc) _) =
-  sequentially (checkIdent src) $ \src' _ -> do
-    slice' <- checkSlice slice
+  sequentially ((,) <$> checkIdent src <*> checkSlice slice) $ \(src', slice') _ -> do
     (t, _) <- newArrayType (srclocOf src) "src" $ sliceDims slice'
     unify (mkUsage loc "type of target array") t $ toStruct $ unInfo $ identType src'
 
@@ -719,7 +720,7 @@
     zeroOrderType
       (mkUsage loc "being returned 'match'")
       "type returned from pattern match"
-      t
+      (toStruct t)
     pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext)
 checkExp (Attr info e loc) =
   Attr <$> checkAttr info <*> checkExp e <*> pure loc
@@ -1258,14 +1259,14 @@
       typeError usage mempty . withIndexLink "ambiguous-type" $
         "Type is ambiguous (must be equality type)."
           </> "Add a type annotation to disambiguate the type."
-    fixOverloaded (_, HasFields fs usage) =
+    fixOverloaded (_, HasFields _ fs usage) =
       typeError usage mempty . withIndexLink "ambiguous-type" $
         "Type is ambiguous.  Must be record with fields:"
           </> indent 2 (stack $ map field $ M.toList fs)
           </> "Add a type annotation to disambiguate the type."
       where
         field (l, t) = pretty l <> colon <+> align (pretty t)
-    fixOverloaded (_, HasConstrs cs usage) =
+    fixOverloaded (_, HasConstrs _ cs usage) =
       typeError usage mempty . withIndexLink "ambiguous-type" $
         "Type is ambiguous (must be a sum type with constructors:"
           <+> pretty (Sum cs) <> ")."
diff --git a/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs b/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs
--- a/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/DoLoop.hs
@@ -206,6 +206,7 @@
     zeroOrderType
       (mkUsage (srclocOf mergeexp) "use as loop variable")
       "type used as loop variable"
+      . toStruct
       =<< expTypeFully mergeexp'
 
     -- The handling of dimension sizes is a bit intricate, but very
diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
--- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
@@ -235,7 +235,6 @@
   (Pat -> TermTypeM a) ->
   TermTypeM a
 bindingPat sizes p t m = do
-  checkForDuplicateNames (map sizeBinderToParam sizes) [p]
   checkPat sizes p t $ \p' -> binding True (S.toList $ patIdents p') $ do
     -- Perform an observation of every declared dimension.  This
     -- prevents unused-name warnings for otherwise unused dimensions.
@@ -354,6 +353,13 @@
   pure $ PatLit l (Info (fromStruct t')) loc
 checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs)))
   | Just ts <- M.lookup n cs = do
+      when (length ps /= length ts) $
+        typeError loc mempty $
+          "Pattern #" <> pretty n <> " expects"
+            <+> pretty (length ps)
+            <+> "constructor arguments, but type provides"
+            <+> pretty (length ts)
+            <+> "arguments."
       ps' <- zipWithM (checkPat' sizes) ps $ map Ascribed ts
       pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc
 checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed t) = do
diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs
--- a/src/Language/Futhark/TypeChecker/Unify.hs
+++ b/src/Language/Futhark/TypeChecker/Unify.hs
@@ -122,9 +122,9 @@
   | ParamType Liftedness SrcLoc
   | Constraint StructRetType Usage
   | Overloaded [PrimType] Usage
-  | HasFields (M.Map Name StructType) Usage
+  | HasFields Liftedness (M.Map Name StructType) Usage
   | Equality Usage
-  | HasConstrs (M.Map Name [StructType]) Usage
+  | HasConstrs Liftedness (M.Map Name [StructType]) Usage
   | ParamSize SrcLoc
   | -- | Is not actually a type, but a term-level size,
     -- possibly already set to something specific.
@@ -141,9 +141,9 @@
   locOf (ParamType _ usage) = locOf usage
   locOf (Constraint _ usage) = locOf usage
   locOf (Overloaded _ usage) = locOf usage
-  locOf (HasFields _ usage) = locOf usage
+  locOf (HasFields _ _ usage) = locOf usage
   locOf (Equality usage) = locOf usage
-  locOf (HasConstrs _ usage) = locOf usage
+  locOf (HasConstrs _ _ usage) = locOf usage
   locOf (ParamSize loc) = locOf loc
   locOf (Size _ usage) = locOf usage
   locOf (UnknowableSize loc _) = locOf loc
@@ -270,7 +270,7 @@
 typeVarNotes :: MonadUnify m => VName -> m Notes
 typeVarNotes v = maybe mempty (note . snd) . M.lookup v <$> getConstraints
   where
-    note (HasConstrs cs _) =
+    note (HasConstrs _ cs _) =
       aNote $
         prettyName v
           <+> "="
@@ -278,7 +278,7 @@
           <+> "..."
     note (Overloaded ts _) =
       aNote $ prettyName v <+> "must be one of" <+> mconcat (punctuate ", " (map pretty ts))
-    note (HasFields fs _) =
+    note (HasFields _ fs _) =
       aNote $
         prettyName v
           <+> "="
@@ -678,21 +678,21 @@
                 <> commasep (map (dquotes . prettyName) problems)
                 <> " used as size(s) would go out of scope."
 
+  let unliftedBcs unlifted_usage =
+        breadCrumb
+          ( Matching $
+              "When verifying that"
+                <+> dquotes (prettyName vn)
+                <+> textwrap "is not instantiated with a function type, due to"
+                <+> pretty unlifted_usage
+          )
+          bcs
+
   case snd <$> M.lookup vn constraints of
     Just (NoConstraint Unlifted unlift_usage) -> do
-      let bcs' =
-            breadCrumb
-              ( Matching $
-                  "When verifying that"
-                    <+> dquotes (prettyName vn)
-                    <+> textwrap "is not instantiated with a function type, due to"
-                    <+> pretty unlift_usage
-              )
-              bcs
-
       link
 
-      arrayElemTypeWith usage bcs' tp
+      arrayElemTypeWith usage (unliftedBcs unlift_usage) tp
       when (any (`elem` bound) (freeInType tp)) $
         unifyError usage mempty bcs $
           "Type variable"
@@ -722,7 +722,8 @@
                   <+> commasep (map pretty ts)
                   </> "due to"
                   <+> pretty old_usage <> "."
-    Just (HasFields required_fields old_usage) -> do
+    Just (HasFields l required_fields old_usage) -> do
+      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp
       case tp of
         Scalar (Record tp_fields)
           | all (`M.member` tp_fields) $ M.keys required_fields -> do
@@ -734,7 +735,7 @@
               unifySharedFields onDims usage bound bcs required_fields' tp_fields
         Scalar (TypeVar _ _ (QualName [] v) []) -> do
           case M.lookup v constraints of
-            Just (_, HasFields tp_fields _) ->
+            Just (_, HasFields _ tp_fields _) ->
               unifySharedFields onDims usage bound bcs required_fields tp_fields
             Just (_, NoConstraint {}) -> pure ()
             Just (_, Equality {}) -> pure ()
@@ -746,10 +747,10 @@
             M.insertWith
               combineFields
               v
-              (lvl, HasFields required_fields old_usage)
+              (lvl, HasFields l required_fields old_usage)
           where
-            combineFields (_, HasFields fs1 usage1) (_, HasFields fs2 _) =
-              (lvl, HasFields (M.union fs1 fs2) usage1)
+            combineFields (_, HasFields l1 fs1 usage1) (_, HasFields l2 fs2 _) =
+              (lvl, HasFields (l1 `min` l2) (M.union fs1 fs2) usage1)
             combineFields hasfs _ = hasfs
         _ ->
           unifyError usage mempty bcs $
@@ -764,7 +765,8 @@
               </> "due to"
               <+> pretty old_usage <> "."
     -- See Note [Linking variables to sum types]
-    Just (HasConstrs required_cs old_usage) ->
+    Just (HasConstrs l required_cs old_usage) -> do
+      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp
       case tp of
         Scalar (Sum ts)
           | all (`M.member` ts) $ M.keys required_cs -> do
@@ -775,7 +777,7 @@
               unifySharedConstructors onDims usage bound bcs required_cs ts
         Scalar (TypeVar _ _ (QualName [] v) []) -> do
           case M.lookup v constraints of
-            Just (_, HasConstrs v_cs _) ->
+            Just (_, HasConstrs _ v_cs _) ->
               unifySharedConstructors onDims usage bound bcs required_cs v_cs
             Just (_, NoConstraint {}) -> pure ()
             Just (_, Equality {}) -> pure ()
@@ -787,10 +789,10 @@
             M.insertWith
               combineConstrs
               v
-              (lvl, HasConstrs required_cs old_usage)
+              (lvl, HasConstrs l required_cs old_usage)
           where
-            combineConstrs (_, HasConstrs cs1 usage1) (_, HasConstrs cs2 _) =
-              (lvl, HasConstrs (M.union cs1 cs2) usage1)
+            combineConstrs (_, HasConstrs l1 cs1 usage1) (_, HasConstrs l2 cs2 _) =
+              (lvl, HasConstrs (l1 `min` l2) (M.union cs1 cs2) usage1)
             combineConstrs hasCs _ = hasCs
         _ -> noSumType mempty
     _ -> link
@@ -875,14 +877,14 @@
               <+> "due to"
               <+> pretty vn_usage <> "."
         ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage)
-    Just (_, HasConstrs _ vn_usage) ->
+    Just (_, HasConstrs _ _ vn_usage) ->
       unifyError usage mempty noBreadCrumbs $
         "Type constrained to one of"
           <+> commasep (map pretty ts)
             <> ", but also inferred to be sum type due to"
           <+> pretty vn_usage
             <> "."
-    Just (_, HasFields _ vn_usage) ->
+    Just (_, HasFields _ _ vn_usage) ->
       unifyError usage mempty noBreadCrumbs $
         "Type constrained to one of"
           <+> commasep (map pretty ts)
@@ -927,29 +929,33 @@
           pure () -- All primtypes support equality.
         Just (_, Equality {}) ->
           pure ()
-        Just (_, HasConstrs cs _) ->
+        Just (_, HasConstrs _ cs _) ->
           mapM_ (equalityType usage) $ concat $ M.elems cs
         _ ->
           unifyError usage mempty noBreadCrumbs $
             "Type" <+> prettyName vn <+> "does not support equality."
 
 zeroOrderTypeWith ::
-  (MonadUnify m, Pretty (Shape dim), Monoid as) =>
+  MonadUnify m =>
   Usage ->
   BreadCrumbs ->
-  TypeBase dim as ->
+  StructType ->
   m ()
 zeroOrderTypeWith usage bcs t = do
   unless (orderZero t) $
     unifyError usage mempty bcs $
       "Type" </> indent 2 (pretty t) </> "found to be functional."
-  mapM_ mustBeZeroOrder . S.toList . typeVars $ t
+  mapM_ mustBeZeroOrder . S.toList . typeVars =<< normType t
   where
     mustBeZeroOrder vn = do
       constraints <- getConstraints
       case M.lookup vn constraints of
         Just (lvl, NoConstraint _ _) ->
           modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage)
+        Just (lvl, HasFields _ fs _) ->
+          modifyConstraints $ M.insert vn (lvl, HasFields Unlifted fs usage)
+        Just (lvl, HasConstrs _ cs _) ->
+          modifyConstraints $ M.insert vn (lvl, HasConstrs Unlifted cs usage)
         Just (_, ParamType Lifted ploc) ->
           unifyError usage mempty bcs $
             "Type parameter"
@@ -961,11 +967,7 @@
 
 -- | Assert that this type must be zero-order.
 zeroOrderType ::
-  (MonadUnify m, Pretty (Shape dim), Monoid as) =>
-  Usage ->
-  T.Text ->
-  TypeBase dim as ->
-  m ()
+  MonadUnify m => Usage -> T.Text -> StructType -> m ()
 zeroOrderType usage desc =
   zeroOrderTypeWith usage $ breadCrumb bc noBreadCrumbs
   where
@@ -1057,12 +1059,14 @@
   constraints <- getConstraints
   case t of
     Scalar (TypeVar _ _ (QualName _ tn) [])
-      | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do
+      | Just (lvl, NoConstraint l _) <- M.lookup tn constraints -> do
           mapM_ (scopeCheck usage noBreadCrumbs tn lvl) fs
-          modifyConstraints $ M.insert tn (lvl, HasConstrs (M.singleton c fs) usage)
-      | Just (lvl, HasConstrs cs _) <- M.lookup tn constraints ->
+          modifyConstraints $ M.insert tn (lvl, HasConstrs l (M.singleton c fs) usage)
+      | Just (lvl, HasConstrs l cs _) <- M.lookup tn constraints ->
           case M.lookup c cs of
-            Nothing -> modifyConstraints $ M.insert tn (lvl, HasConstrs (M.insert c fs cs) usage)
+            Nothing ->
+              modifyConstraints $
+                M.insert tn (lvl, HasConstrs l (M.insert c fs cs) usage)
             Just fs'
               | length fs == length fs' -> zipWithM_ (unify usage) fs fs'
               | otherwise ->
@@ -1098,16 +1102,16 @@
     Scalar (TypeVar _ _ (QualName _ tn) [])
       | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do
           scopeCheck usage bcs tn lvl l_type
-          modifyConstraints $ M.insert tn (lvl, HasFields (M.singleton l l_type) usage)
+          modifyConstraints $ M.insert tn (lvl, HasFields Lifted (M.singleton l l_type) usage)
           pure l_type'
-      | Just (lvl, HasFields fields _) <- M.lookup tn constraints -> do
+      | Just (lvl, HasFields lifted fields _) <- M.lookup tn constraints -> do
           case M.lookup l fields of
             Just t' -> unifyWith onDims usage bound bcs l_type t'
             Nothing ->
               modifyConstraints $
                 M.insert
                   tn
-                  (lvl, HasFields (M.insert l l_type fields) usage)
+                  (lvl, HasFields lifted (M.insert l l_type fields) usage)
           pure l_type'
     Scalar (Record fields)
       | Just t' <- M.lookup l fields -> do