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.20.5
+version:        0.20.6
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -312,7 +312,7 @@
       alex:alex
     , happy:happy
   build-depends:
-      aeson < 2
+      aeson >=2.0.0.0
     , ansi-terminal >=0.6.3.1
     , array >=0.4
     , base >=4.13 && <5
@@ -331,7 +331,7 @@
     , filepath >=1.4.1.1
     , free >=4.12.4
     , futhark-data >= 1.0.2.0
-    , futhark-server >= 1.1.1.0
+    , futhark-server >= 1.1.2.0
     , githash >=0.1.6.1
     , half >= 0.3
     , haskeline
@@ -353,7 +353,6 @@
     , text >=1.2.2.2
     , time >=1.6.0.1
     , transformers >=0.3
-    , unordered-containers >=0.2.7
     , vector >=0.12
     , vector-binary-instances >=0.2.2.0
     , versions >=5.0.0
diff --git a/rts/c/scheduler.h b/rts/c/scheduler.h
--- a/rts/c/scheduler.h
+++ b/rts/c/scheduler.h
@@ -581,7 +581,7 @@
 #endif
   worker->nested++;
   int err = subtask->fn(subtask->args, subtask->start, subtask->end,
-                        subtask->chunkable ? worker->tid : subtask->id,
+                        subtask->id,
                         worker->tid);
   worker->nested--;
   // Some error occured during some other subtask
@@ -1131,6 +1131,9 @@
     struct worker *cur_worker = &scheduler->workers[i];
     CHECK_ERR(pthread_join(scheduler->workers[i].thread, NULL), "pthread_join");
   }
+
+  // And then destroy our local queue.
+  subtask_queue_destroy(&worker_local->q);
 
   free(scheduler->workers);
 
diff --git a/rts/c/server.h b/rts/c/server.h
--- a/rts/c/server.h
+++ b/rts/c/server.h
@@ -170,7 +170,7 @@
 
 struct server_state {
   struct futhark_prog prog;
-  struct futhark_context *cfg;
+  struct futhark_context_config *cfg;
   struct futhark_context *ctx;
   int variables_capacity;
   struct variable *variables;
@@ -661,6 +661,7 @@
     ok();
   }
 
+  free(s.variables);
   free(line);
 }
 
diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs
--- a/src/Futhark/Analysis/HORep/SOAC.hs
+++ b/src/Futhark/Analysis/HORep/SOAC.hs
@@ -397,8 +397,7 @@
 
 instance PrettyRep rep => PP.Pretty (SOAC rep) where
   ppr (Screma w form arrs) = Futhark.ppScrema w arrs form
-  ppr (Hist len ops bucket_fun imgs) =
-    Futhark.ppHist len ops bucket_fun imgs
+  ppr (Hist len ops bucket_fun imgs) = Futhark.ppHist len imgs ops bucket_fun
   ppr soac = text $ show soac
 
 -- | Returns the inputs used in a SOAC.
@@ -478,19 +477,15 @@
 toExp soac = Op <$> toSOAC soac
 
 -- | Convert a SOAC to a Futhark-level SOAC.
-toSOAC ::
-  MonadBuilder m =>
-  SOAC (Rep m) ->
-  m (Futhark.SOAC (Rep m))
+toSOAC :: MonadBuilder m => SOAC (Rep m) -> m (Futhark.SOAC (Rep m))
 toSOAC (Stream w form lam nes inps) =
   Futhark.Stream w <$> inputsToSubExps inps <*> pure form <*> pure nes <*> pure lam
-toSOAC (Scatter len lam ivs dests) = do
-  ivs' <- inputsToSubExps ivs
-  return $ Futhark.Scatter len lam ivs' dests
+toSOAC (Scatter w lam ivs dests) =
+  Futhark.Scatter w <$> inputsToSubExps ivs <*> pure lam <*> pure dests
 toSOAC (Screma w form arrs) =
   Futhark.Screma w <$> inputsToSubExps arrs <*> pure form
-toSOAC (Hist w ops lam inps) =
-  Futhark.Hist w ops lam <$> inputsToSubExps inps
+toSOAC (Hist w ops lam arrs) =
+  Futhark.Hist w <$> inputsToSubExps arrs <*> pure ops <*> pure lam
 
 -- | The reason why some expression cannot be converted to a 'SOAC'
 -- value.
@@ -508,11 +503,11 @@
   m (Either NotSOAC (SOAC rep))
 fromExp (Op (Futhark.Stream w as form nes lam)) =
   Right . Stream w form lam nes <$> traverse varInput as
-fromExp (Op (Futhark.Scatter len lam ivs as)) =
-  Right <$> (Scatter len lam <$> traverse varInput ivs <*> pure as)
+fromExp (Op (Futhark.Scatter w ivs lam as)) =
+  Right <$> (Scatter w lam <$> traverse varInput ivs <*> pure as)
 fromExp (Op (Futhark.Screma w arrs form)) =
   Right . Screma w form <$> traverse varInput arrs
-fromExp (Op (Futhark.Hist w ops lam arrs)) =
+fromExp (Op (Futhark.Hist w arrs ops lam)) =
   Right . Hist w ops lam <$> traverse varInput arrs
 fromExp _ = pure $ Left NotSOAC
 
diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs
--- a/src/Futhark/Analysis/Metrics.hs
+++ b/src/Futhark/Analysis/Metrics.hs
@@ -112,7 +112,7 @@
 expMetrics Apply {} =
   seen "Apply"
 expMetrics (WithAcc _ lam) =
-  inside "MkAcc" $ lambdaMetrics lam
+  inside "WithAcc" $ lambdaMetrics lam
 expMetrics (Op op) =
   opMetrics op
 
diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs
--- a/src/Futhark/Analysis/SymbolTable.hs
+++ b/src/Futhark/Analysis/SymbolTable.hs
@@ -11,11 +11,13 @@
     -- * Entries
     Entry,
     deepen,
+    entryAccInput,
     entryDepth,
     entryLetBoundDec,
     entryIsSize,
     entryStm,
     entryFParam,
+    entryLParam,
 
     -- * Lookup
     elem,
@@ -46,6 +48,7 @@
 
     -- * Misc
     hideCertified,
+    noteAccTokens,
   )
 where
 
@@ -132,6 +135,9 @@
     -- | True if this name has been used as an array size,
     -- implying that it is non-negative.
     entryIsSize :: Bool,
+    -- | For names that are tokens of an accumulator, this is the
+    -- corresponding combining function and neutral element.
+    entryAccInput :: Maybe (WithAccInput rep),
     entryType :: EntryType rep
   }
 
@@ -165,11 +171,13 @@
 
 data LParamEntry rep = LParamEntry
   { lparamDec :: LParamInfo rep,
+    lparamAliases :: Names,
     lparamIndex :: IndexArray
   }
 
 data FreeVarEntry rep = FreeVarEntry
   { freeVarDec :: NameInfo rep,
+    freeVarAliases :: Names,
     -- | Index a delayed array, if possible.
     freeVarIndex :: VName -> IndexArray
   }
@@ -198,9 +206,22 @@
   FParam e' -> Just $ fparamDec e'
   _ -> Nothing
 
+entryLParam :: Entry rep -> Maybe (LParamInfo rep)
+entryLParam e = case entryType e of
+  LParam e' -> Just $ lparamDec e'
+  _ -> Nothing
+
 entryLetBoundDec :: Entry rep -> Maybe (LetDec rep)
 entryLetBoundDec = fmap letBoundDec . isLetBound
 
+entryAliases :: EntryType rep -> Names
+entryAliases (LetBound e) = letBoundAliases e
+entryAliases (FParam e) = fparamAliases e
+entryAliases (LParam e) = lparamAliases e
+entryAliases (FreeVar e) = freeVarAliases e
+entryAliases (LoopVar _) = mempty -- Integers have no aliases.
+
+-- | You almost always want 'available' instead of this one.
 elem :: VName -> SymbolTable rep -> Bool
 elem name = isJust . lookup name
 
@@ -234,10 +255,7 @@
 
 lookupAliases :: VName -> SymbolTable rep -> Names
 lookupAliases name vtable =
-  case entryType <$> M.lookup name (bindings vtable) of
-    Just (LetBound e) -> letBoundAliases e
-    Just (FParam e) -> fparamAliases e
-    _ -> mempty
+  maybe mempty (entryAliases . entryType) $ M.lookup name (bindings vtable)
 
 -- | If the given variable name is the name of a 'ForLoop' parameter,
 -- then return the bound of that loop.
@@ -321,7 +339,8 @@
   | length ds == length is,
     Just (Prim t) <- lookupSubExpType v table =
     Just $ Indexed mempty $ primExpFromSubExp t v
-indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_ : is) =
+indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_ : is) = do
+  guard $ v `available` table
   index' v is table
 indexExp table (BasicOp (Reshape newshape v)) _ is
   | Just oldshape <- arrayDims <$> lookupType v table =
@@ -331,7 +350,8 @@
             (map pe64 $ newDims newshape)
             is
      in index' v is' table
-indexExp table (BasicOp (Index v slice)) _ is =
+indexExp table (BasicOp (Index v slice)) _ is = do
+  guard $ v `available` table
   index' v (adjust (unSlice slice) is) table
   where
     adjust (DimFix j : js') is' =
@@ -384,6 +404,7 @@
           { entryConsumed = False,
             entryDepth = loopDepth vtable,
             entryIsSize = False,
+            entryAccInput = Nothing,
             entryType = entry
           }
       dims = mapMaybe subExpVar $ arrayDims $ typeOf entry'
@@ -431,6 +452,16 @@
             entry
               { fparamAliases = oneName (patElemName pe) <> fparamAliases entry
               }
+        update' (LParam entry) =
+          LParam
+            entry
+              { lparamAliases = oneName (patElemName pe) <> lparamAliases entry
+              }
+        update' (FreeVar entry) =
+          FreeVar
+            entry
+              { freeVarAliases = oneName (patElemName pe) <> freeVarAliases entry
+              }
         update' e = e
 
 insertStms ::
@@ -476,6 +507,7 @@
       LParam
         LParamEntry
           { lparamDec = AST.paramDec param,
+            lparamAliases = mempty,
             lparamIndex = const Nothing
           }
     name = AST.paramName param
@@ -520,7 +552,8 @@
       FreeVar
         FreeVarEntry
           { freeVarDec = dec,
-            freeVarIndex = \_ _ -> Nothing
+            freeVarIndex = \_ _ -> Nothing,
+            freeVarAliases = mempty
           }
 
 consume :: VName -> SymbolTable rep -> SymbolTable rep
@@ -544,7 +577,8 @@
               FreeVar
                 FreeVarEntry
                   { freeVarDec = entryInfo entry,
-                    freeVarIndex = \_ _ -> Nothing
+                    freeVarIndex = \_ _ -> Nothing,
+                    freeVarAliases = entryAliases $ entryType entry
                   }
           }
       | otherwise = entry
@@ -555,3 +589,21 @@
 hideCertified to_hide = hideIf $ maybe False hide . entryStm
   where
     hide = any (`nameIn` to_hide) . unCerts . stmCerts
+
+-- | Note that these names are tokens for the corresponding
+-- accumulators.  The names must already be present in the symbol
+-- table.
+noteAccTokens ::
+  [(VName, WithAccInput rep)] ->
+  SymbolTable rep ->
+  SymbolTable rep
+noteAccTokens = flip (foldl' f)
+  where
+    f vtable (v, accum) =
+      case M.lookup v $ bindings vtable of
+        Nothing -> vtable
+        Just e ->
+          vtable
+            { bindings =
+                M.insert v (e {entryAccInput = Just accum}) $ bindings vtable
+            }
diff --git a/src/Futhark/Bench.hs b/src/Futhark/Bench.hs
--- a/src/Futhark/Bench.hs
+++ b/src/Futhark/Bench.hs
@@ -21,9 +21,10 @@
 import Control.Applicative
 import Control.Monad.Except
 import qualified Data.Aeson as JSON
+import qualified Data.Aeson.Key as JSON
+import qualified Data.Aeson.KeyMap as JSON
 import qualified Data.ByteString.Char8 as SBS
 import qualified Data.ByteString.Lazy.Char8 as LBS
-import qualified Data.HashMap.Strict as HM
 import qualified Data.Map as M
 import Data.Maybe
 import qualified Data.Text as T
@@ -81,19 +82,19 @@
 
 instance JSON.FromJSON DataResults where
   parseJSON = JSON.withObject "datasets" $ \o ->
-    DataResults <$> mapM datasetResult (HM.toList o)
+    DataResults <$> mapM datasetResult (JSON.toList o)
     where
       datasetResult (k, v) =
-        DataResult (T.unpack k)
+        DataResult (JSON.toString k)
           <$> ((Right <$> success v) <|> (Left <$> JSON.parseJSON v))
       success = JSON.withObject "result" $ \o ->
         Result <$> o JSON..: "runtimes" <*> o JSON..: "bytes" <*> o JSON..: "stderr"
 
-dataResultJSON :: DataResult -> (T.Text, JSON.Value)
+dataResultJSON :: DataResult -> (JSON.Key, JSON.Value)
 dataResultJSON (DataResult desc (Left err)) =
-  (T.pack desc, JSON.toJSON err)
+  (JSON.fromString desc, JSON.toJSON err)
 dataResultJSON (DataResult desc (Right (Result runtimes bytes progerr))) =
-  ( T.pack desc,
+  ( JSON.fromString desc,
     JSON.object
       [ ("runtimes", JSON.toJSON $ map runMicroseconds runtimes),
         ("bytes", JSON.toJSON bytes),
@@ -101,22 +102,22 @@
       ]
   )
 
-benchResultJSON :: BenchResult -> (T.Text, JSON.Value)
+benchResultJSON :: BenchResult -> (JSON.Key, JSON.Value)
 benchResultJSON (BenchResult prog r) =
-  ( T.pack prog,
-    JSON.Object $ HM.singleton "datasets" (JSON.toJSON $ DataResults r)
+  ( JSON.fromString prog,
+    JSON.object [("datasets", JSON.toJSON $ DataResults r)]
   )
 
 instance JSON.ToJSON BenchResults where
   toJSON (BenchResults rs) =
-    JSON.Object $ HM.fromList $ map benchResultJSON rs
+    JSON.object $ map benchResultJSON rs
 
 instance JSON.FromJSON BenchResults where
   parseJSON = JSON.withObject "benchmarks" $ \o ->
-    BenchResults <$> mapM onBenchmark (HM.toList o)
+    BenchResults <$> mapM onBenchmark (JSON.toList o)
     where
       onBenchmark (k, v) =
-        BenchResult (T.unpack k)
+        BenchResult (JSON.toString k)
           <$> JSON.withObject "benchmark" onBenchmark' v
       onBenchmark' o =
         fmap unDataResults . JSON.parseJSON =<< o JSON..: "datasets"
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
@@ -125,11 +125,11 @@
   case pipeline of
     SOACSPipeline ->
       check ["-s"]
-    KernelsPipeline ->
+    GpuPipeline ->
       check ["--gpu"]
-    SequentialCpuPipeline ->
+    SeqMemPipeline ->
       check ["--seq-mem"]
-    GpuPipeline ->
+    GpuMemPipeline ->
       check ["--gpu-mem"]
     NoPipeline ->
       check []
@@ -154,9 +154,9 @@
   where
     maybePipeline :: StructurePipeline -> T.Text
     maybePipeline SOACSPipeline = "(soacs) "
-    maybePipeline KernelsPipeline = "(kernels) "
-    maybePipeline SequentialCpuPipeline = "(seq-mem) "
-    maybePipeline GpuPipeline = "(gpu-mem) "
+    maybePipeline GpuPipeline = "(kernels) "
+    maybePipeline SeqMemPipeline = "(seq-mem) "
+    maybePipeline GpuMemPipeline = "(gpu-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
@@ -18,6 +18,7 @@
 import Futhark.CodeGen.Backends.COpenCL.Boilerplate (commonOptions, sizeLoggingCode)
 import qualified Futhark.CodeGen.Backends.GenericC as GC
 import Futhark.CodeGen.Backends.GenericC.Options
+import Futhark.CodeGen.Backends.SimpleRep (primStorageType, toStorage)
 import Futhark.CodeGen.ImpCode.OpenCL
 import qualified Futhark.CodeGen.ImpGen.CUDA as ImpGen
 import Futhark.IR.GPUMem hiding
@@ -359,6 +360,11 @@
     expNotZero e = [C.cexp|$exp:e != 0|]
     expAnd a b = [C.cexp|$exp:a && $exp:b|]
     expsNotZero = foldl expAnd [C.cexp|1|] . map expNotZero
+    mkArgs (ValueKArg e t@(FloatType Float16)) = do
+      arg <- newVName "kernel_arg"
+      e' <- GC.compileExp e
+      GC.item [C.citem|$ty:(primStorageType t) $id:arg = $exp:(toStorage t e');|]
+      pure (arg, Nothing)
     mkArgs (ValueKArg e t) =
       (,Nothing) <$> GC.compileExpToName "kernel_arg" t e
     mkArgs (MemKArg v) = do
diff --git a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
@@ -297,7 +297,7 @@
   GC.CompilerM OpenCL () ()
 generateContextFuns cfg cost_centres kernels sizes failures = do
   final_inits <- GC.contextFinalInits
-  (fields, init_fields) <- GC.contextContents
+  (fields, init_fields, free_fields) <- GC.contextContents
   let forCostCentre name =
         [ ( [C.csdecl|typename int64_t $id:(kernelRuntime name);|],
             [C.cstm|ctx->$id:(kernelRuntime name) = 0;|]
@@ -411,6 +411,7 @@
   GC.publicDef_ "context_free" GC.InitDecl $ \s ->
     ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|],
       [C.cedecl|void $id:s(struct $id:ctx* ctx) {
+                                 $stms:free_fields
                                  free_constants(ctx);
                                  cuda_cleanup(&ctx->cuda);
                                  free_lock(&ctx->lock);
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
@@ -18,6 +18,7 @@
 import Futhark.CodeGen.Backends.COpenCL.Boilerplate
 import qualified Futhark.CodeGen.Backends.GenericC as GC
 import Futhark.CodeGen.Backends.GenericC.Options
+import Futhark.CodeGen.Backends.SimpleRep (primStorageType, toStorage)
 import Futhark.CodeGen.ImpCode.OpenCL
 import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen
 import Futhark.IR.GPUMem hiding
@@ -344,8 +345,17 @@
   when (safety >= SafetyFull) $
     GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]
   where
-    setKernelArg i (ValueKArg e bt) = do
-      v <- GC.compileExpToName "kernel_arg" bt e
+    setKernelArg i (ValueKArg e pt) = do
+      v <- case pt of
+        -- We always transfer f16 values to the kernel as 16 bits, but
+        -- the actual host type may be typedef'd to a 32-bit float.
+        -- This requires some care.
+        FloatType Float16 -> do
+          v <- newVName "kernel_arg"
+          e' <- toStorage pt <$> GC.compileExp e
+          GC.decl [C.cdecl|$ty:(primStorageType pt) $id:v = $e';|]
+          pure v
+        _ -> GC.compileExpToName "kernel_arg" pt e
       GC.stm
         [C.cstm|
             OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, sizeof($id:v), &$id:v));
diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
@@ -303,7 +303,7 @@
                        }|]
     )
 
-  (fields, init_fields) <- GC.contextContents
+  (fields, init_fields, free_fields) <- GC.contextContents
   ctx <- GC.publicDef "context" GC.InitDecl $ \s ->
     ( [C.cedecl|struct $id:s;|],
       [C.cedecl|struct $id:s {
@@ -436,6 +436,7 @@
   GC.publicDef_ "context_free" GC.InitDecl $ \s ->
     ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|],
       [C.cedecl|void $id:s(struct $id:ctx* ctx) {
+                                 $stms:free_fields
                                  free_constants(ctx);
                                  free_lock(&ctx->lock);
                                  $stms:(map releaseKernel (M.toList kernels))
diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs
--- a/src/Futhark/CodeGen/Backends/GenericC.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC.hs
@@ -65,12 +65,14 @@
     publicName,
     contextType,
     contextField,
+    contextFieldDyn,
     memToCType,
     cacheMem,
     fatMemory,
     rawMemCType,
     cproduct,
     fatMemType,
+    freeAllocatedMem,
 
     -- * Building Blocks
     primTypeToCType,
@@ -84,6 +86,7 @@
 import Control.Monad.State
 import Data.Bifunctor (first)
 import qualified Data.DList as DL
+import Data.List (unzip4)
 import Data.Loc
 import qualified Data.Map.Strict as M
 import Data.Maybe
@@ -119,7 +122,7 @@
     compUserState :: s,
     compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition),
     compLibDecls :: DL.DList C.Definition,
-    compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp),
+    compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp, Maybe C.Stm),
     compProfileItems :: DL.DList C.BlockItem,
     compClearItems :: DL.DList C.BlockItem,
     compDeclaredMem :: [(VName, Space)],
@@ -244,15 +247,19 @@
   (formatstrs, formatargs) <- unzip <$> mapM onPart parts
   pure (mconcat formatstrs, formatargs)
 
+freeAllocatedMem :: CompilerM op s [C.BlockItem]
+freeAllocatedMem = collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem
+
 defError :: ErrorCompiler op s
 defError msg stacktrace = do
-  free_all_mem <- collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem
+  free_all_mem <- freeAllocatedMem
   (formatstr, formatargs) <- errorMsgString msg
   let formatstr' = "Error: " <> formatstr <> "\n\nBacktrace:\n%s"
   items
     [C.citems|ctx->error = msgprintf($string:formatstr', $args:formatargs, $string:stacktrace);
-                  $items:free_all_mem
-                  return 1;|]
+              $items:free_all_mem
+              err = 1;
+              goto cleanup;|]
 
 defCall :: CallCompiler op s
 defCall dests fname args = do
@@ -366,9 +373,10 @@
 entryDecls = declsCode (== EntryDecl)
 miscDecls = declsCode (== MiscDecl)
 
-contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm])
+contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm], [C.Stm])
 contextContents = do
-  (field_names, field_types, field_values) <- gets $ unzip3 . DL.toList . compCtxFields
+  (field_names, field_types, field_values, field_frees) <-
+    gets $ unzip4 . DL.toList . compCtxFields
   let fields =
         [ [C.csdecl|$ty:ty $id:name;|]
           | (name, ty) <- zip field_names field_types
@@ -377,7 +385,7 @@
         [ [C.cstm|ctx->$id:name = $exp:e;|]
           | (name, Just e) <- zip field_names field_values
         ]
-  return (fields, init_fields)
+  return (fields, init_fields, catMaybes field_frees)
 
 contextFinalInits :: CompilerM op s [C.Stm]
 contextFinalInits = gets compInit
@@ -503,7 +511,11 @@
 
 contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s ()
 contextField name ty initial = modify $ \s ->
-  s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial)}
+  s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Nothing)}
+
+contextFieldDyn :: C.Id -> C.Type -> C.Exp -> C.Stm -> CompilerM op s ()
+contextFieldDyn name ty initial free = modify $ \s ->
+  s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, Just initial, Just free)}
 
 profileReport :: C.BlockItem -> CompilerM op s ()
 profileReport x = modify $ \s ->
diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
@@ -19,7 +19,6 @@
 import Futhark.CodeGen.Backends.GenericC.Options
 import Futhark.CodeGen.Backends.SimpleRep
   ( cproduct,
-    fromStorage,
     primAPIType,
     primStorageType,
     scalarToPrim,
@@ -176,7 +175,7 @@
             [C.cstm|;|],
             [C.cstm|;|],
             [C.cstm|;|],
-            fromStorage t [C.cexp|$id:dest|]
+            [C.cexp|$id:dest|]
           )
     Just (TypeOpaque desc _) ->
       ( [C.citems|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:(T.unpack desc));|],
diff --git a/src/Futhark/CodeGen/Backends/GenericC/Manifest.hs b/src/Futhark/CodeGen/Backends/GenericC/Manifest.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/Manifest.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/Manifest.hs
@@ -19,7 +19,9 @@
 
 import Data.Aeson (ToJSON (..), object)
 import qualified Data.Aeson as JSON
+import qualified Data.Aeson.Key as JSON
 import Data.Aeson.Text (encodeToLazyText)
+import Data.Bifunctor (bimap)
 import qualified Data.Map as M
 import qualified Data.Text as T
 import Data.Text.Lazy (toStrict)
@@ -112,10 +114,10 @@
     object
       [ ("backend", toJSON backend),
         ( "entry_points",
-          object $ M.toList $ fmap onEntryPoint entry_points
+          object $ map (bimap JSON.fromText onEntryPoint) $ M.toList entry_points
         ),
         ( "types",
-          object $ M.toList $ fmap onType types
+          object $ map (bimap JSON.fromText onType) $ M.toList types
         )
       ]
     where
diff --git a/src/Futhark/CodeGen/Backends/MulticoreC.hs b/src/Futhark/CodeGen/Backends/MulticoreC.hs
--- a/src/Futhark/CodeGen/Backends/MulticoreC.hs
+++ b/src/Futhark/CodeGen/Backends/MulticoreC.hs
@@ -110,7 +110,7 @@
                            }|]
     )
 
-  (fields, init_fields) <- GC.contextContents
+  (fields, init_fields, free_fields) <- GC.contextContents
 
   ctx <- GC.publicDef "context" GC.InitDecl $ \s ->
     ( [C.cedecl|struct $id:s;|],
@@ -180,6 +180,7 @@
   GC.publicDef_ "context_free" GC.InitDecl $ \s ->
     ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|],
       [C.cedecl|void $id:s(struct $id:ctx* ctx) {
+             $stms:free_fields
              free_constants(ctx);
              (void)scheduler_destroy(&ctx->scheduler);
              free_lock(&ctx->lock);
@@ -408,9 +409,21 @@
 
 addBenchmarkFields :: Name -> Maybe VName -> GC.CompilerM op s ()
 addBenchmarkFields name (Just _) = do
-  GC.contextField (functionRuntime name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(typename int64_t), ctx->scheduler.num_threads)|]
-  GC.contextField (functionRuns name) [C.cty|int*|] $ Just [C.cexp|calloc(sizeof(int), ctx->scheduler.num_threads)|]
-  GC.contextField (functionIter name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(sizeof(typename int64_t)), ctx->scheduler.num_threads)|]
+  GC.contextFieldDyn
+    (functionRuntime name)
+    [C.cty|typename int64_t*|]
+    [C.cexp|calloc(sizeof(typename int64_t), ctx->scheduler.num_threads)|]
+    [C.cstm|free(ctx->$id:(functionRuntime name));|]
+  GC.contextFieldDyn
+    (functionRuns name)
+    [C.cty|int*|]
+    [C.cexp|calloc(sizeof(int), ctx->scheduler.num_threads)|]
+    [C.cstm|free(ctx->$id:(functionRuns name));|]
+  GC.contextFieldDyn
+    (functionIter name)
+    [C.cty|typename int64_t*|]
+    [C.cexp|calloc(sizeof(sizeof(typename int64_t)), ctx->scheduler.num_threads)|]
+    [C.cstm|free(ctx->$id:(functionIter name));|]
 addBenchmarkFields name Nothing = do
   GC.contextField (functionRuntime name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]
   GC.contextField (functionRuns name) [C.cty|int|] $ Just [C.cexp|0|]
@@ -489,7 +502,7 @@
           mapM_ GC.item decl_cached
           code' <- GC.blockScope $ GC.compileCode code
           mapM_ GC.item [C.citems|$items:code'|]
-          mapM_ GC.stm free_cached
+          GC.stm [C.cstm|cleanup: {$stms:free_cached}|]
     return
       [C.cedecl|int $id:s(void *args, typename int64_t iterations, int tid, struct scheduler_info info) {
                            int err = 0;
@@ -498,8 +511,9 @@
                            struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args;
                            struct futhark_context *ctx = $id:fstruct->ctx;
                            $items:fbody
-                           $stms:(compileWriteBackResVals fstruct retval_args retval_ctypes)
-                           cleanup: {}
+                           if (err == 0) {
+                             $stms:(compileWriteBackResVals fstruct retval_args retval_ctypes)
+                           }
                            return err;
                       }|]
 
@@ -570,11 +584,14 @@
       GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]
       return mempty
 
+  free_all_mem <- GC.freeAllocatedMem
   let ftask_err = fpar_task <> "_err"
-  let code =
+      code =
         [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);
                   if ($id:ftask_err != 0) {
-                    err = 1; goto cleanup;
+                    $items:free_all_mem;
+                    err = 1;
+                    goto cleanup;
                   }|]
 
   mapM_ GC.item code
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
@@ -255,9 +255,8 @@
   finishIfSynchronous
   where
     processKernelArg :: Imp.KernelArg -> Py.CompilerM op s PyExp
-    processKernelArg (Imp.ValueKArg e bt) = do
-      e' <- Py.compileExp e
-      return $ Py.simpleCall (Py.compilePrimToNp bt) [e']
+    processKernelArg (Imp.ValueKArg e bt) =
+      Py.toStorage bt <$> Py.compileExp e
     processKernelArg (Imp.MemKArg v) = Py.compileVar v
     processKernelArg (Imp.SharedMemoryKArg (Imp.Count num_bytes)) = do
       num_bytes' <- Py.compileExp num_bytes
diff --git a/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs
@@ -53,7 +53,7 @@
                                }|]
     )
 
-  (fields, init_fields) <- GC.contextContents
+  (fields, init_fields, free_fields) <- GC.contextContents
 
   ctx <- GC.publicDef "context" GC.InitDecl $ \s ->
     ( [C.cedecl|struct $id:s;|],
@@ -93,6 +93,7 @@
   GC.publicDef_ "context_free" GC.InitDecl $ \s ->
     ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|],
       [C.cedecl|void $id:s(struct $id:ctx* ctx) {
+                                 $stms:free_fields
                                  free_constants(ctx);
                                  free_lock(&ctx->lock);
                                  free(ctx);
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
@@ -753,6 +753,10 @@
   -- While-loop: Try to insert your value
   let (toBits, fromBits) =
         case t of
+          FloatType Float16 ->
+            ( \v -> Imp.FunExp "to_bits16" [v] int16,
+              \v -> Imp.FunExp "from_bits16" [v] t
+            )
           FloatType Float32 ->
             ( \v -> Imp.FunExp "to_bits32" [v] int32,
               \v -> Imp.FunExp "from_bits32" [v] t
@@ -764,6 +768,7 @@
           _ -> (id, id)
 
       int
+        | primBitSize t == 16 = int16
         | primBitSize t == 32 = int32
         | otherwise = int64
 
@@ -774,15 +779,14 @@
     old_bits_v <- newVName "old_bits"
     dPrim_ old_bits_v int
     let old_bits = Imp.var old_bits_v int
-    sOp $
-      Imp.Atomic space $
-        Imp.AtomicCmpXchg
-          int
-          old_bits_v
-          arr'
-          bucket_offset
-          (toBits (Imp.var assumed t))
-          (toBits (Imp.var x t))
+    sOp . Imp.Atomic space $
+      Imp.AtomicCmpXchg
+        int
+        old_bits_v
+        arr'
+        bucket_offset
+        (toBits (Imp.var assumed t))
+        (toBits (Imp.var x t))
     old <~~ fromBits old_bits
     let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits
     sWhen (isBool won) (run_loop <-- false)
@@ -1141,35 +1145,35 @@
 
     barrier
 
-  let read_carry_in = do
+  no_carry_in <- dPrimVE "no_carry_in" $ is_first_block .||. bNot ltid_in_bounds
+
+  let read_carry_in = sUnless no_carry_in $ do
         forM_ (zip x_params y_params) $ \(x, y) ->
           copyDWIM (paramName y) [] (Var (paramName x)) []
         zipWithM_ readPrevBlockResult x_params arrs
 
-      y_to_x = forM_ (zip x_params y_params) $ \(x, y) ->
-        when (primType (paramType x)) $
-          copyDWIM (paramName x) [] (Var (paramName y)) []
-
       op_to_x
         | Nothing <- seg_flag =
-          compileBody' x_params $ lambdaBody lam
+          sUnless no_carry_in $ compileBody' x_params $ lambdaBody lam
         | Just flag_true <- seg_flag = do
           inactive <-
             dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid32
-          sWhen inactive y_to_x
+          sUnless no_carry_in . sWhen inactive . forM_ (zip x_params y_params) $ \(x, y) ->
+            copyDWIM (paramName x) [] (Var (paramName y)) []
+          -- The convoluted control flow is to ensure all threads
+          -- hit this barrier (if applicable).
           when array_scan barrier
-          sUnless inactive $ compileBody' x_params $ lambdaBody lam
+          sUnless no_carry_in $ sUnless inactive $ compileBody' x_params $ lambdaBody lam
 
       write_final_result =
         forM_ (zip x_params arrs) $ \(p, arr) ->
           when (primType $ paramType p) $
             copyDWIM arr [DimFix ltid] (Var $ paramName p) []
 
-  sComment "carry-in for every block except the first" $
-    sUnless (is_first_block .||. bNot ltid_in_bounds) $ do
-      sComment "read operands" read_carry_in
-      sComment "perform operation" op_to_x
-      sComment "write final result" write_final_result
+  sComment "carry-in for every block except the first" $ do
+    sComment "read operands" read_carry_in
+    sComment "perform operation" op_to_x
+    sComment "write final result" $ sUnless no_carry_in write_final_result
 
   barrier
 
@@ -1195,9 +1199,7 @@
   InKernelGen ()
 inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do
   skip_threads <- dPrim "skip_threads" int32
-  let in_block_thread_active =
-        tvExp skip_threads .<=. in_block_id
-      actual_params = lambdaParams scan_lam
+  let actual_params = lambdaParams scan_lam
       (x_params, y_params) =
         splitAt (length actual_params `div` 2) actual_params
       y_to_x =
@@ -1215,16 +1217,21 @@
 
   when array_scan barrier
 
-  let op_to_x
+  let op_to_x in_block_thread_active
         | Nothing <- seg_flag =
-          compileBody' x_params $ lambdaBody scan_lam
+          sWhen in_block_thread_active $
+            compileBody' x_params $ lambdaBody scan_lam
         | Just flag_true <- seg_flag = do
           inactive <-
-            dPrimVE "inactive" $
-              flag_true (ltid32 - tvExp skip_threads) ltid32
-          sWhen inactive y_to_x
+            dPrimVE "inactive" $ flag_true (ltid32 - tvExp skip_threads) ltid32
+          sWhen (in_block_thread_active .&&. inactive) $
+            forM_ (zip x_params y_params) $ \(x, y) ->
+              copyDWIM (paramName x) [] (Var (paramName y)) []
+          -- The convoluted control flow is to ensure all threads
+          -- hit this barrier (if applicable).
           when array_scan barrier
-          sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam
+          sWhen in_block_thread_active . sUnless inactive $
+            compileBody' x_params $ lambdaBody scan_lam
 
       maybeBarrier =
         sWhen
@@ -1234,16 +1241,17 @@
   sComment "in-block scan (hopefully no barriers needed)" $ do
     skip_threads <-- 1
     sWhile (tvExp skip_threads .<. block_size) $ do
-      sWhen (in_block_thread_active .&&. active) $ do
-        sComment "read operands" $
-          zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs
-        sComment "perform operation" op_to_x
+      thread_active <-
+        dPrimVE "thread_active" $ tvExp skip_threads .<=. in_block_id .&&. active
 
+      sWhen thread_active . sComment "read operands" $
+        zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs
+      sComment "perform operation" $ op_to_x thread_active
+
       maybeBarrier
 
-      sWhen (in_block_thread_active .&&. active) $
-        sComment "write result" $
-          sequence_ $ zipWith3 writeResult x_params y_params arrs
+      sWhen thread_active . sComment "write result" $
+        sequence_ $ zipWith3 writeResult x_params y_params arrs
 
       maybeBarrier
 
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
@@ -239,12 +239,9 @@
           GC.blockScope $ GC.compileCode $ kernelBody kernel
       kstate = GC.compUserState cstate
 
-      use_params = mapMaybe useAsParam $ kernelUses kernel
-
       (local_memory_args, local_memory_params, local_memory_init) =
-        unzip3 $
-          flip evalState (blankNameSource :: VNameSource) $
-            mapM (prepareLocalMemory target) $ kernelLocalMemory kstate
+        unzip3 . flip evalState (blankNameSource :: VNameSource) $
+          mapM (prepareLocalMemory target) $ kernelLocalMemory kstate
 
       -- CUDA has very strict restrictions on the number of blocks
       -- permitted along the 'y' and 'z' dimensions of the grid
@@ -275,6 +272,9 @@
 
       (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel
 
+  let (use_params, unpack_params) =
+        unzip $ mapMaybe useAsParam $ kernelUses kernel
+
   let (safety, error_init)
         -- We conservatively assume that any called function can fail.
         | not $ null called =
@@ -324,6 +324,7 @@
 
       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
@@ -368,16 +369,24 @@
           [C.citem|volatile $ty:defaultMemBlockType $id:mem = &shared_mem[$id:param];|]
         )
 
-useAsParam :: KernelUse -> Maybe C.Param
-useAsParam (ScalarUse name bt) =
-  let ctp = case bt of
+useAsParam :: KernelUse -> Maybe (C.Param, [C.BlockItem])
+useAsParam (ScalarUse name pt) = do
+  let name_bits = zEncodeString (pretty name) <> "_bits"
+      ctp = case pt of
         -- OpenCL does not permit bool as a kernel parameter type.
         Bool -> [C.cty|unsigned char|]
         Unit -> [C.cty|unsigned char|]
-        _ -> GC.primTypeToCType bt
-   in Just [C.cparam|$ty:ctp $id:name|]
+        _ -> primStorageType pt
+  if ctp == primTypeToCType pt
+    then Just ([C.cparam|$ty:ctp $id:name|], [])
+    else
+      let name_bits_e = [C.cexp|$id:name_bits|]
+       in Just
+            ( [C.cparam|$ty:ctp $id:name_bits|],
+              [[C.citem|$ty:(primTypeToCType pt) $id:name = $exp:(fromStorage pt name_bits_e);|]]
+            )
 useAsParam (MemoryUse name) =
-  Just [C.cparam|__global $ty:defaultMemBlockType $id:name|]
+  Just ([C.cparam|__global $ty:defaultMemBlockType $id:name|], [])
 useAsParam ConstUse {} =
   Nothing
 
diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
--- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
+++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
@@ -94,10 +94,15 @@
 -- When the SegRed's return value is a scalar
 -- we perform a call by value-result in the segop function
 getReturnParams :: Pat MCMem -> SegOp () MCMem -> MulticoreGen [Imp.Param]
-getReturnParams pat SegRed {} = do
-  let retvals = map patElemName $ patElems pat
-  retvals_ts <- mapM lookupType retvals
-  concat <$> zipWithM toParam retvals retvals_ts
+getReturnParams pat SegRed {} =
+  -- It's a good idea to make sure any prim values are initialised, as
+  -- we will load them (redundantly) in the task code, and
+  -- uninitialised values are UB.
+  fmap concat . forM (patElems pat) $ \pe -> do
+    case patElemType pe of
+      Prim pt -> patElemName pe <~~ ValueExp (blankPrimValue pt)
+      _ -> pure ()
+    toParam (patElemName pe) (patElemType pe)
 getReturnParams _ _ = return mempty
 
 renameSegBinOp :: [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem]
@@ -371,50 +376,50 @@
   MulticoreGen () ->
   MulticoreGen ()
 atomicUpdateCAS t arr old bucket x do_op = do
-  -- Code generation target:
-  --
-  -- old = d_his[idx];
-  -- do {
-  --   assumed = old;
-  --   x = do_op(assumed, y);
-  --   old = atomicCAS(&d_his[idx], assumed, tmp);
-  -- } while(assumed != old);
   run_loop <- dPrimV "run_loop" (0 :: Imp.TExp Int32)
-  everythingVolatile $ copyDWIMFix old [] (Var arr) bucket
   (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket
 
   bytes <- toIntegral $ primBitSize t
-  (to, from) <- getBitConvertFunc $ primBitSize t
-  -- While-loop: Try to insert your value
-  let (toBits, _fromBits) =
+  let (toBits, fromBits) =
         case t of
-          FloatType _ ->
-            ( \v -> Imp.FunExp to [v] bytes,
-              \v -> Imp.FunExp from [v] t
+          FloatType Float16 ->
+            ( \v -> Imp.FunExp "to_bits16" [v] int16,
+              \v -> Imp.FunExp "from_bits16" [v] t
             )
+          FloatType Float32 ->
+            ( \v -> Imp.FunExp "to_bits32" [v] int32,
+              \v -> Imp.FunExp "from_bits32" [v] t
+            )
+          FloatType Float64 ->
+            ( \v -> Imp.FunExp "to_bits64" [v] int64,
+              \v -> Imp.FunExp "from_bits64" [v] t
+            )
           _ -> (id, id)
 
+      int
+        | primBitSize t == 16 = int16
+        | primBitSize t == 32 = int32
+        | otherwise = int64
+
+  everythingVolatile $ copyDWIMFix old [] (Var arr) bucket
+
+  old_bits_v <- tvVar <$> dPrim "old_bits" int
+  old_bits_v <~~ toBits (Imp.var old t)
+  let old_bits = Imp.var old_bits_v int
+
+  -- While-loop: Try to insert your value
   sWhile (tvExp run_loop .==. 0) $ do
     x <~~ Imp.var old t
     do_op -- Writes result into x
-    sOp $
-      Imp.Atomic $
-        Imp.AtomicCmpXchg
-          bytes
-          old
-          arr'
-          (sExt32 <$> bucket_offset)
-          (tvVar run_loop)
-          (toBits (Imp.var x t))
-
--- TODO for supporting 8 and 16 bits (and 128)
--- we need a functions for converting to and from bits
-getBitConvertFunc :: Int -> MulticoreGen (String, String)
--- getBitConvertFunc 8 = return $ ("to_bits8, from_bits8")
--- getBitConvertFunc 16 = return $ ("to_bits8, from_bits8")
-getBitConvertFunc 32 = return ("to_bits32", "from_bits32")
-getBitConvertFunc 64 = return ("to_bits64", "from_bits64")
-getBitConvertFunc b = error $ "number of bytes is not supported " ++ pretty b
+    sOp . Imp.Atomic $
+      Imp.AtomicCmpXchg
+        bytes
+        old_bits_v
+        arr'
+        (sExt32 <$> bucket_offset)
+        (tvVar run_loop)
+        (toBits (Imp.var x t))
+    old <~~ fromBits old_bits
 
 supportedPrims :: Int -> Bool
 supportedPrims 8 = True
diff --git a/src/Futhark/IR/GPU/Op.hs b/src/Futhark/IR/GPU/Op.hs
--- a/src/Futhark/IR/GPU/Op.hs
+++ b/src/Futhark/IR/GPU/Op.hs
@@ -13,6 +13,7 @@
 
     -- * Host operations
     HostOp (..),
+    traverseHostOpStms,
     typeCheckHostOp,
 
     -- * SegOp refinements
@@ -240,6 +241,15 @@
   | SizeOp SizeOp
   | OtherOp op
   deriving (Eq, Ord, Show)
+
+-- | A helper for defining 'TraverseOpStms'.
+traverseHostOpStms ::
+  Monad m =>
+  OpStmsTraverser m op rep ->
+  OpStmsTraverser m (HostOp rep op) rep
+traverseHostOpStms _ f (SegOp segop) = SegOp <$> traverseSegOpStms f segop
+traverseHostOpStms _ _ (SizeOp sizeop) = pure $ SizeOp sizeop
+traverseHostOpStms onOtherOp f (OtherOp other) = OtherOp <$> onOtherOp f other
 
 instance (ASTRep rep, Substitute op) => Substitute (HostOp rep op) where
   substituteNames substs (SegOp op) =
diff --git a/src/Futhark/IR/GPU/Simplify.hs b/src/Futhark/IR/GPU/Simplify.hs
--- a/src/Futhark/IR/GPU/Simplify.hs
+++ b/src/Futhark/IR/GPU/Simplify.hs
@@ -75,6 +75,9 @@
   w' <- Engine.simplify w
   return (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty)
 
+instance TraverseOpStms (Wise GPU) where
+  traverseOpStms = traverseHostOpStms traverseSOACStms
+
 instance BuilderOps (Wise GPU)
 
 instance HasSegOp (Wise GPU) where
diff --git a/src/Futhark/IR/GPUMem.hs b/src/Futhark/IR/GPUMem.hs
--- a/src/Futhark/IR/GPUMem.hs
+++ b/src/Futhark/IR/GPUMem.hs
@@ -82,6 +82,9 @@
   mkBodyB stms res = return $ Engine.mkWiseBody () stms res
   mkLetNamesB = mkLetNamesB''
 
+instance TraverseOpStms (Engine.Wise GPUMem) where
+  traverseOpStms = traverseMemOpStms (traverseHostOpStms (const pure))
+
 simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem)
 simplifyProg = simplifyProgGeneric simpleGPUMem
 
diff --git a/src/Futhark/IR/MC.hs b/src/Futhark/IR/MC.hs
--- a/src/Futhark/IR/MC.hs
+++ b/src/Futhark/IR/MC.hs
@@ -61,6 +61,9 @@
 
 instance PrettyRep MC
 
+instance TraverseOpStms (Engine.Wise MC) where
+  traverseOpStms = traverseMCOpStms traverseSOACStms
+
 simpleMC :: Simplify.SimpleOps MC
 simpleMC = Simplify.bindableSimpleOps $ simplifyMCOp SOAC.simplifySOAC
 
diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs
--- a/src/Futhark/IR/MC/Op.hs
+++ b/src/Futhark/IR/MC/Op.hs
@@ -9,6 +9,7 @@
 -- also re-exported from here.
 module Futhark.IR.MC.Op
   ( MCOp (..),
+    traverseMCOpStms,
     typeCheckMCOp,
     simplifyMCOp,
     module Futhark.IR.SegOp,
@@ -50,6 +51,11 @@
   | -- | Something else (in practice often a SOAC).
     OtherOp op
   deriving (Eq, Ord, Show)
+
+traverseMCOpStms :: Monad m => OpStmsTraverser m op rep -> OpStmsTraverser m (MCOp rep op) rep
+traverseMCOpStms _ f (ParOp par_op op) =
+  ParOp <$> traverse (traverseSegOpStms f) par_op <*> traverseSegOpStms f op
+traverseMCOpStms onInner f (OtherOp op) = OtherOp <$> onInner f op
 
 instance (ASTRep rep, Substitute op) => Substitute (MCOp rep op) where
   substituteNames substs (ParOp par_op op) =
diff --git a/src/Futhark/IR/MCMem.hs b/src/Futhark/IR/MCMem.hs
--- a/src/Futhark/IR/MCMem.hs
+++ b/src/Futhark/IR/MCMem.hs
@@ -79,6 +79,9 @@
   mkBodyB stms res = return $ Engine.mkWiseBody () stms res
   mkLetNamesB = mkLetNamesB''
 
+instance TraverseOpStms (Engine.Wise MCMem) where
+  traverseOpStms = traverseMemOpStms (traverseMCOpStms (const pure))
+
 simplifyProg :: Prog MCMem -> PassM (Prog MCMem)
 simplifyProg = simplifyProgGeneric simpleMCMem
 
diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs
--- a/src/Futhark/IR/Mem.hs
+++ b/src/Futhark/IR/Mem.hs
@@ -62,6 +62,7 @@
     RetTypeMem,
     BranchTypeMem,
     MemOp (..),
+    traverseMemOpStms,
     MemInfo (..),
     MemBound,
     MemBind (..),
@@ -181,6 +182,14 @@
     Alloc SubExp Space
   | Inner inner
   deriving (Eq, Ord, Show)
+
+-- | A helper for defining 'TraverseOpStms'.
+traverseMemOpStms ::
+  Monad m =>
+  OpStmsTraverser m inner rep ->
+  OpStmsTraverser m (MemOp inner) rep
+traverseMemOpStms _ _ op@Alloc {} = pure op
+traverseMemOpStms onInner f (Inner inner) = Inner <$> onInner f inner
 
 instance FreeIn inner => FreeIn (MemOp inner) where
   freeIn' (Alloc size _) = freeIn' size
diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs
--- a/src/Futhark/IR/Parse.hs
+++ b/src/Futhark/IR/Parse.hs
@@ -604,8 +604,8 @@
       keyword "scatter"
         *> parens
           ( SOAC.Scatter <$> pSubExp <* pComma
-              <*> pLambda pr <* pComma
-              <*> braces (pVName `sepBy` pComma)
+              <*> braces (pVName `sepBy` pComma) <* pComma
+              <*> pLambda pr
               <*> many (pComma *> pDest)
           )
       where
@@ -616,9 +616,9 @@
         *> parens
           ( SOAC.Hist
               <$> pSubExp <* pComma
+              <*> braces (pVName `sepBy` pComma) <* pComma
               <*> braces (pHistOp `sepBy` pComma) <* pComma
               <*> pLambda pr
-              <*> many (pComma *> pVName)
           )
       where
         pHistOp =
diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs
--- a/src/Futhark/IR/SOACS/SOAC.hs
+++ b/src/Futhark/IR/SOACS/SOAC.hs
@@ -37,6 +37,7 @@
     isScanSOAC,
     isReduceSOAC,
     isMapSOAC,
+    scremaLambda,
     ppScrema,
     ppHist,
     groupScatterResults,
@@ -47,6 +48,7 @@
     SOACMapper (..),
     identitySOACMapper,
     mapSOACM,
+    traverseSOACStms,
   )
 where
 
@@ -115,13 +117,13 @@
     -- will correspond to the first two output values, and so on. For this
     -- example, <lambda> should return a total of 11 values, 8 index values and
     -- 3 output values.
-    Scatter SubExp (Lambda rep) [VName] [(Shape, Int, VName)]
+    Scatter SubExp [VName] (Lambda rep) [(Shape, Int, VName)]
   | -- | @Hist <length> <dest-arrays-and-ops> <bucket fun> <input arrays>@
     --
     -- The first SubExp is the length of the input arrays. The first
     -- list describes the operations to perform.  The t'Lambda' is the
     -- bucket function.  Finally comes the input images.
-    Hist SubExp [HistOp rep] (Lambda rep) [VName]
+    Hist SubExp [VName] [HistOp rep] (Lambda rep)
   | -- | A combination of scan, reduction, and map.  The first
     -- t'SubExp' is the size of the input arrays.
     Screma SubExp [VName] (ScremaForm rep)
@@ -319,6 +321,14 @@
   guard $ null reds
   return map_lam
 
+-- | Return the "main" lambda of the Screma.  For a map, this is
+-- equivalent to 'isMapSOAC'.  Note that the meaning of the return
+-- value of this lambda depends crucially on exactly which Screma this
+-- is.  The parameters will correspond exactly to elements of the
+-- input arrays, however.
+scremaLambda :: ScremaForm rep -> Lambda rep
+scremaLambda (ScremaForm _ _ map_lam) = map_lam
+
 -- | @groupScatterResults <output specification> <results>@
 --
 -- Groups the index values and result values of <results> according to the
@@ -403,11 +413,11 @@
       Parallel o comm <$> mapOnSOACLambda tv lam0
     mapOnStreamForm Sequential =
       pure Sequential
-mapSOACM tv (Scatter len lam ivs as) =
+mapSOACM tv (Scatter w ivs lam as) =
   Scatter
-    <$> mapOnSOACSubExp tv len
-    <*> mapOnSOACLambda tv lam
+    <$> mapOnSOACSubExp tv w
     <*> mapM (mapOnSOACVName tv) ivs
+    <*> mapOnSOACLambda tv lam
     <*> mapM
       ( \(aw, an, a) ->
           (,,) <$> mapM (mapOnSOACSubExp tv) aw
@@ -415,20 +425,20 @@
             <*> mapOnSOACVName tv a
       )
       as
-mapSOACM tv (Hist len ops bucket_fun imgs) =
+mapSOACM tv (Hist w arrs ops bucket_fun) =
   Hist
-    <$> mapOnSOACSubExp tv len
+    <$> mapOnSOACSubExp tv w
+    <*> mapM (mapOnSOACVName tv) arrs
     <*> mapM
-      ( \(HistOp e rf arrs nes op) ->
+      ( \(HistOp e rf op_arrs nes op) ->
           HistOp <$> mapOnSOACSubExp tv e
             <*> mapOnSOACSubExp tv rf
-            <*> mapM (mapOnSOACVName tv) arrs
+            <*> mapM (mapOnSOACVName tv) op_arrs
             <*> mapM (mapOnSOACSubExp tv) nes
             <*> mapOnSOACLambda tv op
       )
       ops
     <*> mapOnSOACLambda tv bucket_fun
-    <*> mapM (mapOnSOACVName tv) imgs
 mapSOACM tv (Screma w arrs (ScremaForm scans reds map_lam)) =
   Screma <$> mapOnSOACSubExp tv w
     <*> mapM (mapOnSOACVName tv) arrs
@@ -448,6 +458,12 @@
             <*> mapOnSOACLambda tv map_lam
         )
 
+-- | A helper for defining 'TraverseOpStms'.
+traverseSOACStms :: Monad m => OpStmsTraverser m (SOAC rep) rep
+traverseSOACStms f = mapSOACM mapper
+  where
+    mapper = identitySOACMapper {mapOnSOACLambda = traverseLambdaStms f}
+
 instance ASTRep rep => FreeIn (SOAC rep) where
   freeIn' = flip execState mempty . mapSOACM free
     where
@@ -483,13 +499,13 @@
     nms = map paramName $ take (1 + length accs) params
     substs = M.fromList $ zip nms (outersize : accs)
     Lambda params _ rtp = lam
-soacType (Scatter _w lam _ivs as) =
+soacType (Scatter _w _ivs lam dests) =
   zipWith arrayOfShape val_ts ws
   where
     indexes = sum $ zipWith (*) ns $ map length ws
     val_ts = drop indexes $ lambdaReturnType lam
-    (ws, ns, _) = unzip3 as
-soacType (Hist _len ops _bucket_fun _imgs) = do
+    (ws, ns, _) = unzip3 dests
+soacType (Hist _ _ ops _bucket_fun) = do
   op <- ops
   map (`arrayOfRow` histWidth op) (lambdaReturnType $ histOp op)
 soacType (Screma w _arrs form) =
@@ -523,7 +539,7 @@
         zip (map paramName $ drop 1 $ lambdaParams lam) (accs ++ map Var arrs)
   consumedInOp (Scatter _ _ _ as) =
     namesFromList $ map (\(_, _, a) -> a) as
-  consumedInOp (Hist _ ops _ _) =
+  consumedInOp (Hist _ _ ops _) =
     namesFromList $ concatMap histDest ops
 
 mapHistOp ::
@@ -549,14 +565,14 @@
       analyseStreamForm (Parallel o comm lam0) =
         Parallel o comm (Alias.analyseLambda aliases lam0)
       analyseStreamForm Sequential = Sequential
-  addOpAliases aliases (Scatter len lam ivs as) =
-    Scatter len (Alias.analyseLambda aliases lam) ivs as
-  addOpAliases aliases (Hist len ops bucket_fun imgs) =
+  addOpAliases aliases (Scatter len arrs lam dests) =
+    Scatter len arrs (Alias.analyseLambda aliases lam) dests
+  addOpAliases aliases (Hist w arrs ops bucket_fun) =
     Hist
-      len
+      w
+      arrs
       (map (mapHistOp (Alias.analyseLambda aliases)) ops)
       (Alias.analyseLambda aliases bucket_fun)
-      imgs
   addOpAliases aliases (Screma w arrs (ScremaForm scans reds map_lam)) =
     Screma w arrs $
       ScremaForm
@@ -666,7 +682,7 @@
   -- arr, so we can later check that aliases of arr are not used inside lam.
   let fake_lamarrs' = map asArg lamarrs'
   TC.checkLambda lam $ asArg inttp : accargs ++ fake_lamarrs'
-typeCheckSOAC (Scatter w lam ivs as) = do
+typeCheckSOAC (Scatter w arrs lam as) = do
   -- Requirements:
   --
   --   0. @lambdaReturnType@ of @lam@ must be a list
@@ -683,7 +699,7 @@
   --      of a requirement, so that e.g. the source is not hoisted out of a
   --      loop, which will mean it cannot be consumed.
   --
-  --   5. Each of ivs must be an array matching a corresponding lambda
+  --   5. Each of arrs must be an array matching a corresponding lambda
   --      parameters.
   --
   -- Code:
@@ -718,10 +734,10 @@
     TC.consume =<< TC.lookupAliases a
 
   -- 5.
-  arrargs <- TC.checkSOACArrayArgs w ivs
+  arrargs <- TC.checkSOACArrayArgs w arrs
   TC.checkLambda lam arrargs
-typeCheckSOAC (Hist len ops bucket_fun imgs) = do
-  TC.require [Prim int64] len
+typeCheckSOAC (Hist w arrs ops bucket_fun) = do
+  TC.require [Prim int64] w
 
   -- Check the operators.
   forM_ ops $ \(HistOp dest_w rf dests nes op) -> do
@@ -746,7 +762,7 @@
       TC.consume =<< TC.lookupAliases dest
 
   -- Types of input arrays must equal parameter types for bucket function.
-  img' <- TC.checkSOACArrayArgs len imgs
+  img' <- TC.checkSOACArrayArgs w arrs
   TC.checkLambda bucket_fun img'
 
   -- Return type of bucket function must be an index for each
@@ -807,9 +823,9 @@
 instance OpMetrics (Op rep) => OpMetrics (SOAC rep) where
   opMetrics (Stream _ _ _ _ lam) =
     inside "Stream" $ lambdaMetrics lam
-  opMetrics (Scatter _len lam _ivs _as) =
+  opMetrics (Scatter _len _ lam _) =
     inside "Scatter" $ lambdaMetrics lam
-  opMetrics (Hist _len ops bucket_fun _imgs) =
+  opMetrics (Hist _ _ ops bucket_fun) =
     inside "Hist" $ mapM_ (lambdaMetrics . histOp) ops >> lambdaMetrics bucket_fun
   opMetrics (Screma _ _ (ScremaForm scans reds map_lam)) =
     inside "Screma" $ do
@@ -841,15 +857,16 @@
                 </> ppTuple' acc <> comma
                 </> ppr lam
             )
-  ppr (Scatter w lam ivs as) =
+  ppr (Scatter w arrs lam dests) =
     "scatter"
       <> parens
         ( ppr w <> comma
+            </> ppTuple' arrs <> comma
             </> ppr lam <> comma
-            </> commasep (ppTuple' ivs : map ppr as)
+            </> commasep (map ppr dests)
         )
-  ppr (Hist len ops bucket_fun imgs) =
-    ppHist len ops bucket_fun imgs
+  ppr (Hist w arrs ops bucket_fun) =
+    ppHist w arrs ops bucket_fun
   ppr (Screma w arrs (ScremaForm scans reds map_lam))
     | null scans,
       null reds =
@@ -907,20 +924,20 @@
 ppHist ::
   (PrettyRep rep, Pretty inp) =>
   SubExp ->
+  [inp] ->
   [HistOp rep] ->
   Lambda rep ->
-  [inp] ->
   Doc
-ppHist len ops bucket_fun imgs =
+ppHist w arrs ops bucket_fun =
   text "hist"
     <> parens
-      ( ppr len <> comma
+      ( ppr w <> comma
+          </> ppTuple' arrs <> comma
           </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops) <> comma
-          </> ppr bucket_fun <> comma
-          </> commasep (map ppr imgs)
+          </> ppr bucket_fun
       )
   where
-    ppOp (HistOp w rf dests nes op) =
-      ppr w <> comma <+> ppr rf <> comma <+> PP.braces (commasep $ map ppr dests) <> comma
-        </> PP.braces (commasep $ map ppr nes) <> comma
+    ppOp (HistOp dest_w rf dests nes op) =
+      ppr dest_w <> comma <+> ppr rf <> comma <+> PP.braces (commasep $ map ppr dests) <> comma
+        </> ppTuple' nes <> comma
         </> ppr op
diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs
--- a/src/Futhark/IR/SOACS/Simplify.hs
+++ b/src/Futhark/IR/SOACS/Simplify.hs
@@ -112,13 +112,13 @@
       return (Parallel o comm lam0', hoisted)
     simplifyStreamForm Sequential =
       return (Sequential, mempty)
-simplifySOAC (Scatter len lam ivs as) = do
-  len' <- Engine.simplify len
+simplifySOAC (Scatter w ivs lam as) = do
+  w' <- Engine.simplify w
   (lam', hoisted) <- Engine.simplifyLambda lam
   ivs' <- mapM Engine.simplify ivs
   as' <- mapM Engine.simplify as
-  return (Scatter len' lam' ivs' as', hoisted)
-simplifySOAC (Hist w ops bfun imgs) = do
+  return (Scatter w' ivs' lam' as', hoisted)
+simplifySOAC (Hist w imgs ops bfun) = do
   w' <- Engine.simplify w
   (ops', hoisted) <- fmap unzip $
     forM ops $ \(HistOp dests_w rf dests nes op) -> do
@@ -130,7 +130,7 @@
       return (HistOp dests_w' rf' dests' nes' op', hoisted)
   imgs' <- mapM Engine.simplify imgs
   (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun
-  return (Hist w' ops' bfun' imgs', mconcat hoisted <> bfun_hoisted)
+  return (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted)
 simplifySOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do
   (scans', scans_hoisted) <- fmap unzip $
     forM scans $ \(Scan lam nes) -> do
@@ -155,6 +155,9 @@
 
 instance BuilderOps (Wise SOACS)
 
+instance TraverseOpStms (Wise SOACS) where
+  traverseOpStms = traverseSOACStms
+
 fixLambdaParams ::
   (MonadBuilder m, Buildable (Rep m), BuilderOps (Rep m)) =>
   AST.Lambda (Rep m) ->
@@ -356,10 +359,10 @@
 
 -- | Like 'removeReplicateMapping', but for 'Scatter'.
 removeReplicateWrite :: TopDownRuleOp (Wise SOACS)
-removeReplicateWrite vtable pat aux (Scatter len lam ivs as)
+removeReplicateWrite vtable pat aux (Scatter w ivs lam as)
   | Just (stms, lam', ivs') <- removeReplicateInput vtable lam ivs = Simplify $ do
     forM_ stms $ \(vs, cs, e) -> certifying cs $ letBindNames vs e
-    auxing aux $ letBind pat $ Op $ Scatter len lam' ivs' as
+    auxing aux $ letBind pat $ Op $ Scatter w ivs' lam' as
 removeReplicateWrite _ _ _ _ = Skip
 
 removeReplicateInput ::
@@ -553,7 +556,7 @@
 
 -- | If we are writing to an array that is never used, get rid of it.
 removeDeadWrite :: BottomUpRuleOp (Wise SOACS)
-removeDeadWrite (_, used) pat aux (Scatter w fun arrs dests) =
+removeDeadWrite (_, used) pat aux (Scatter w arrs fun dests) =
   let (i_ses, v_ses) = unzip $ groupScatterResults' dests $ bodyResult $ lambdaBody fun
       (i_ts, v_ts) = unzip $ groupScatterResults' dests $ lambdaReturnType fun
       isUsed (bindee, _, _, _, _, _) = (`UT.used` used) $ patElemName bindee
@@ -566,15 +569,14 @@
           }
    in if pat /= Pat pat'
         then
-          Simplify $
-            auxing aux $
-              letBind (Pat pat') $ Op $ Scatter w fun' arrs dests'
+          Simplify . auxing aux $
+            letBind (Pat pat') $ Op $ Scatter w arrs fun' dests'
         else Skip
 removeDeadWrite _ _ _ _ = Skip
 
 -- handles now concatenation of more than two arrays
 fuseConcatScatter :: TopDownRuleOp (Wise SOACS)
-fuseConcatScatter vtable pat _ (Scatter _ fun arrs dests)
+fuseConcatScatter vtable pat _ (Scatter _ arrs fun dests)
   | Just (ws@(w' : _), xss, css) <- unzip3 <$> mapM isConcat arrs,
     xivs <- transpose xss,
     all (w' ==) ws = Simplify $ do
@@ -593,7 +595,7 @@
               lambdaReturnType = mix its <> mix vts
             }
     certifying (mconcat css) . letBind pat . Op $
-      Scatter w' fun' (concat xivs) $ map (incWrites r) dests
+      Scatter w' (concat xivs) fun' $ map (incWrites r) dests
   where
     sizeOf :: VName -> Maybe SubExp
     sizeOf x = arraySize 0 . typeOf <$> ST.lookup x vtable
@@ -867,7 +869,7 @@
 -- corresponding to that transformation performed on the rows of the
 -- full array.
 moveTransformToInput :: TopDownRuleOp (Wise SOACS)
-moveTransformToInput vtable pat aux (Screma w arrs (ScremaForm scan reduce map_lam))
+moveTransformToInput vtable pat aux soac@(Screma w arrs (ScremaForm scan reduce map_lam))
   | ops <- map snd $ filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam,
     not $ null ops = Simplify $ do
     (more_arrs, more_params, replacements) <-
@@ -875,18 +877,20 @@
 
     when (null more_arrs) cannotSimplify
 
-    let substs = M.fromList $ zip ops replacements
-        map_lam' =
+    let map_lam' =
           map_lam
             { lambdaParams = lambdaParams map_lam <> more_params,
-              lambdaBody =
-                replaceArrayOps substs $
-                  lambdaBody map_lam
+              lambdaBody = replaceArrayOps (M.fromList replacements) $ lambdaBody map_lam
             }
 
     auxing aux $
       letBind pat $ Op $ Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')
   where
+    -- It is not safe to move the transform if the root array is being
+    -- consumed by the Screma.  This is a bit too conservative - it's
+    -- actually safe if we completely replace the original input, but
+    -- this rule is not that precise.
+    consumed = consumedInOp soac
     map_param_names = map paramName (lambdaParams map_lam)
     topLevelPat = (`elem` fmap stmPat (bodyStms (lambdaBody map_lam)))
     onlyUsedOnce arr =
@@ -915,7 +919,8 @@
       False
 
     mapOverArr op
-      | Just (_, arr) <- find ((== arrayOpArr op) . fst) (zip map_param_names arrs) = do
+      | Just (_, arr) <- find ((== arrayOpArr op) . fst) (zip map_param_names arrs),
+        not $ arr `nameIn` consumed = do
         arr_t <- lookupType arr
         let whole_dim = DimSlice (intConst Int64 0) (arraySize 0 arr_t) (intConst Int64 1)
         arr_transformed <- certifying (arrayOpCerts op) $
@@ -933,11 +938,11 @@
                 BasicOp $ SubExp $ Var arr
         arr_transformed_t <- lookupType arr_transformed
         arr_transformed_row <- newVName $ baseString arr ++ "_transformed_row"
-        return $
+        pure $
           Just
             ( arr_transformed,
               Param mempty arr_transformed_row (rowType arr_transformed_t),
-              ArrayVar mempty arr_transformed_row
+              (op, ArrayVar mempty arr_transformed_row)
             )
     mapOverArr _ = return Nothing
 moveTransformToInput _ _ _ _ =
diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs
--- a/src/Futhark/IR/SegOp.hs
+++ b/src/Futhark/IR/SegOp.hs
@@ -40,6 +40,7 @@
     SegOpMapper (..),
     identitySegOpMapper,
     mapSegOpM,
+    traverseSegOpStms,
 
     -- * Simplification
     simplifySegOp,
@@ -841,6 +842,20 @@
   Array et <$> traverse (mapOnSegOpSubExp tv) shape <*> pure u
 mapOnSegOpType _tv (Mem s) = pure $ Mem s
 
+-- | A helper for defining 'TraverseOpStms'.
+traverseSegOpStms :: Monad m => OpStmsTraverser m (SegOp lvl rep) rep
+traverseSegOpStms f segop = mapSegOpM mapper segop
+  where
+    seg_scope = scopeOfSegSpace (segSpace segop)
+    f' scope = f (seg_scope <> scope)
+    mapper =
+      identitySegOpMapper
+        { mapOnSegOpLambda = traverseLambdaStms f',
+          mapOnSegOpBody = onBody
+        }
+    onBody (KernelBody dec stms res) =
+      KernelBody dec <$> f seg_scope stms <*> pure res
+
 instance
   (ASTRep rep, Substitute lvl) =>
   Substitute (SegOp lvl rep)
@@ -1020,7 +1035,7 @@
         | [v] <- patNames $ stmPat stm,
           BasicOp (Index arr slice) <- stmExp stm,
           length (sliceDims slice) == length excess_is,
-          arr `ST.elem` vtable,
+          arr `ST.available` vtable,
           Just (slice', cs) <- asPrimExpSlice table slice =
           let idx =
                 ST.IndexedArray
diff --git a/src/Futhark/IR/Seq.hs b/src/Futhark/IR/Seq.hs
--- a/src/Futhark/IR/Seq.hs
+++ b/src/Futhark/IR/Seq.hs
@@ -50,9 +50,15 @@
 
 instance BuilderOps Seq
 
+instance TraverseOpStms Seq where
+  traverseOpStms _ = pure
+
 instance PrettyRep Seq
 
 instance BuilderOps (Engine.Wise Seq)
+
+instance TraverseOpStms (Engine.Wise Seq) where
+  traverseOpStms _ = pure
 
 simpleSeq :: Simplify.SimpleOps Seq
 simpleSeq = Simplify.bindableSimpleOps (const $ pure ((), mempty))
diff --git a/src/Futhark/IR/SeqMem.hs b/src/Futhark/IR/SeqMem.hs
--- a/src/Futhark/IR/SeqMem.hs
+++ b/src/Futhark/IR/SeqMem.hs
@@ -61,10 +61,16 @@
   mkBodyB stms res = return $ Body () stms res
   mkLetNamesB = mkLetNamesB' ()
 
+instance TraverseOpStms SeqMem where
+  traverseOpStms _ = pure
+
 instance BuilderOps (Engine.Wise SeqMem) where
   mkExpDecB pat e = return $ Engine.mkWiseExpDec pat () e
   mkBodyB stms res = return $ Engine.mkWiseBody () stms res
   mkLetNamesB = mkLetNamesB''
+
+instance TraverseOpStms (Engine.Wise SeqMem) where
+  traverseOpStms _ = pure
 
 simplifyProg :: Prog SeqMem -> PassM (Prog SeqMem)
 simplifyProg = simplifyProgGeneric simpleSeqMem
diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs
--- a/src/Futhark/IR/Syntax.hs
+++ b/src/Futhark/IR/Syntax.hs
@@ -132,6 +132,7 @@
     OpaqueOp (..),
     DimChange (..),
     ShapeChange,
+    WithAccInput,
     ExpT (..),
     Exp,
     LoopForm (..),
@@ -406,6 +407,12 @@
     UpdateAcc VName [SubExp] [SubExp]
   deriving (Eq, Ord, Show)
 
+-- | The input to a 'WithAcc' construct.  Comprises the index space of
+-- the accumulator, the underlying arrays, and possibly a combining
+-- function.
+type WithAccInput rep =
+  (Shape, [VName], Maybe (Lambda rep, [SubExp]))
+
 -- | The root Futhark expression type.  The v'Op' constructor contains
 -- a rep-specific operation.  Do-loops, branches and function calls
 -- are special.  Everything else is a simple t'BasicOp'.
@@ -423,7 +430,7 @@
     -- write index space.  The corresponding arrays must all have this
     -- shape outermost.  This construct is not part of t'BasicOp'
     -- because we need the @rep@ parameter.
-    WithAcc [(Shape, [VName], Maybe (Lambda rep, [SubExp]))] (Lambda rep)
+    WithAcc [WithAccInput rep] (Lambda rep)
   | Op (Op rep)
 
 deriving instance RepTypes rep => Eq (ExpT rep)
diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs
--- a/src/Futhark/IR/Traversals.hs
+++ b/src/Futhark/IR/Traversals.hs
@@ -32,6 +32,11 @@
     Walker (..),
     identityWalker,
     walkExpM,
+
+    -- * Ops
+    TraverseOpStms (..),
+    OpStmsTraverser,
+    traverseLambdaStms,
   )
 where
 
@@ -344,3 +349,18 @@
     (params, args) = unzip merge
 walkExpM tv (Op op) =
   walkOnOp tv op
+
+-- | A function for monadically traversing any sub-statements of the
+-- given op for some representation.
+type OpStmsTraverser m op rep = (Scope rep -> Stms rep -> m (Stms rep)) -> op -> m op
+
+-- | This representatin supports an 'OpStmsTraverser' for its 'Op'.
+-- This is used for some simplification rules.
+class TraverseOpStms rep where
+  -- | Transform every sub-'Stms' of this op.
+  traverseOpStms :: Monad m => OpStmsTraverser m (Op rep) rep
+
+-- | A helper for defining 'traverseOpStms'.
+traverseLambdaStms :: Monad m => OpStmsTraverser m (Lambda rep) rep
+traverseLambdaStms f (Lambda ps (Body dec stms res) ret) =
+  Lambda ps <$> (Body dec <$> f (scopeOfLParams ps) stms <*> pure res) <*> pure ret
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
@@ -72,8 +72,8 @@
 askEnv :: DefM Env
 askEnv = asks snd
 
-isGlobal :: VName -> DefM a -> DefM a
-isGlobal v = local $ Arrow.first (S.insert v)
+areGlobal :: [VName] -> DefM a -> DefM a
+areGlobal vs = local $ Arrow.first (S.fromList vs <>)
 
 replaceTypeSizes ::
   M.Map VName SizeSubst ->
@@ -883,12 +883,8 @@
             Var
               fname'
               ( Info
-                  ( Scalar $
-                      Arrow mempty Unnamed (fromStruct t1) $
-                        RetType [] $
-                          Scalar $
-                            Arrow mempty Unnamed (fromStruct t2) $
-                              RetType [] rettype
+                  ( Scalar . Arrow mempty Unnamed (fromStruct t1) . RetType [] $
+                      Scalar . Arrow mempty Unnamed (fromStruct t2) $ RetType [] rettype
                   )
               )
               loc
@@ -930,11 +926,7 @@
     DynamicFun _ sv -> do
       let (argtypes', rettype) = dynamicFunType sv argtypes
           restype = foldFunType argtypes' (RetType [] rettype) `setAliases` aliases ret
-          -- FIXME: what if this application returns both a function
-          -- and a value?
-          callret
-            | orderZero ret = AppRes ret ext
-            | otherwise = AppRes restype ext
+          callret = AppRes (combineTypeShapes ret restype) ext
           apply_e = AppExp (Apply e1' e2' d loc) (Info callret)
       return (apply_e, sv)
     -- Propagate the 'IntrinsicsSV' until we reach the outermost application,
@@ -1255,7 +1247,7 @@
 -- transformed result as well as an environment that binds the name of
 -- the value binding to the static value of the transformed body.  The
 -- boolean is true if the function is a 'DynamicFun'.
-defuncValBind :: ValBind -> DefM (ValBind, Env, Bool)
+defuncValBind :: ValBind -> DefM (ValBind, Env)
 -- Eta-expand entry points with a functional return type.
 defuncValBind (ValBind entry name _ (Info (RetType _ rettype, retext)) tparams params body _ attrs loc)
   | Scalar Arrow {} <- rettype = do
@@ -1309,11 +1301,7 @@
       M.singleton name $
         Binding
           (Just (first (map typeParamName) (valBindTypeScheme valbind)))
-          sv,
-      case sv of
-        DynamicFun {} -> True
-        Dynamic {} -> True
-        _ -> False
+          sv
     )
   where
     anyDimIfNotBound bound_sizes (NamedDim v)
@@ -1324,12 +1312,10 @@
 defuncVals :: [ValBind] -> DefM ()
 defuncVals [] = pure ()
 defuncVals (valbind : ds) = do
-  (valbind', env, dyn) <- defuncValBind valbind
+  (valbind', env) <- defuncValBind valbind
   addValBind valbind'
-  localEnv env $
-    if dyn
-      then isGlobal (valBindName valbind') $ defuncVals ds
-      else defuncVals ds
+  let globals = valBindName valbind' : snd (unInfo (valBindRetType valbind'))
+  localEnv env $ areGlobal globals $ defuncVals ds
 
 {-# NOINLINE transformProg #-}
 
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
@@ -40,52 +40,51 @@
 
 internaliseValBind :: E.ValBind -> InternaliseM ()
 internaliseValBind fb@(E.ValBind entry fname retdecl (Info (rettype, _)) tparams params body _ attrs loc) = do
-  localConstsScope $
-    bindingFParams tparams params $ \shapeparams params' -> do
-      let shapenames = map I.paramName shapeparams
+  localConstsScope . bindingFParams tparams params $ \shapeparams params' -> do
+    let shapenames = map I.paramName shapeparams
 
-      msg <- case retdecl of
-        Just dt ->
-          errorMsg
-            . ("Function return value does not match shape of type " :)
-            <$> typeExpForError dt
-        Nothing -> return $ errorMsg ["Function return value does not match shape of declared return type."]
+    msg <- case retdecl of
+      Just dt ->
+        errorMsg
+          . ("Function return value does not match shape of type " :)
+          <$> typeExpForError dt
+      Nothing -> return $ errorMsg ["Function return value does not match shape of declared return type."]
 
-      (body', rettype') <- buildBody $ do
-        body_res <- internaliseExp (baseString fname <> "_res") body
-        rettype' <-
-          fmap zeroExts . internaliseReturnType rettype =<< mapM subExpType body_res
-        body_res' <-
-          ensureResultExtShape msg loc (map I.fromDecl rettype') $ subExpsRes body_res
-        pure
-          ( body_res',
-            replicate (length (shapeContext rettype')) (I.Prim int64) ++ rettype'
-          )
+    (body', rettype') <- buildBody $ do
+      body_res <- internaliseExp (baseString fname <> "_res") body
+      rettype' <-
+        fmap zeroExts . internaliseReturnType rettype =<< mapM subExpType body_res
+      body_res' <-
+        ensureResultExtShape msg loc (map I.fromDecl rettype') $ subExpsRes body_res
+      pure
+        ( body_res',
+          replicate (length (shapeContext rettype')) (I.Prim int64) ++ rettype'
+        )
 
-      let all_params = shapeparams ++ concat params'
+    let all_params = shapeparams ++ concat params'
 
-      attrs' <- internaliseAttrs attrs
+    attrs' <- internaliseAttrs attrs
 
-      let fd =
-            I.FunDef
-              Nothing
-              attrs'
-              (internaliseFunName fname)
-              rettype'
-              all_params
-              body'
+    let fd =
+          I.FunDef
+            Nothing
+            attrs'
+            (internaliseFunName fname)
+            rettype'
+            all_params
+            body'
 
-      if null params'
-        then bindConstant fname fd
-        else
-          bindFunction
-            fname
-            fd
-            ( shapenames,
-              map declTypeOf $ concat params',
-              all_params,
-              applyRetType rettype' all_params
-            )
+    if null params'
+      then bindConstant fname fd
+      else
+        bindFunction
+          fname
+          fd
+          ( shapenames,
+            map declTypeOf $ concat params',
+            all_params,
+            applyRetType rettype' all_params
+          )
 
   case entry of
     Just (Info entry') -> generateEntryPoint entry' fb
@@ -973,11 +972,9 @@
 internaliseDimIndex w (E.DimFix i) = do
   (i', _) <- internaliseDimExp "i" i
   let lowerBound =
-        I.BasicOp $
-          I.CmpOp (I.CmpSle I.Int64) (I.constant (0 :: I.Int64)) i'
+        I.BasicOp $ I.CmpOp (I.CmpSle I.Int64) (I.constant (0 :: I.Int64)) i'
       upperBound =
-        I.BasicOp $
-          I.CmpOp (I.CmpSlt I.Int64) i' w
+        I.BasicOp $ I.CmpOp (I.CmpSlt I.Int64) i' w
   ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound)
   return (I.DimFix i', ok, [ErrorVal int64 i'])
 
@@ -1193,12 +1190,11 @@
       "length of index and value array does not match"
       loc
   buckets'' <-
-    certifying c $
-      letExp (baseString buckets') $
-        I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion w_img] 1 b_shape) buckets'
+    certifying c . letExp (baseString buckets') $
+      I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion w_img] 1 b_shape) buckets'
 
   letValExp' desc . I.Op $
-    I.Hist w_img [HistOp w_hist rf' hist' ne_shp op'] lam' $ buckets'' : img'
+    I.Hist w_img (buckets'' : img') [HistOp w_hist rf' hist' ne_shp op'] lam'
 
 internaliseStreamMap ::
   String ->
@@ -1914,7 +1910,7 @@
           sivs = si' <> svs'
 
       let sa_ws = map (Shape . take dim . arrayDims) sa_ts
-      letTupExp' desc $ I.Op $ I.Scatter si_w lam sivs $ zip3 sa_ws (repeat 1) sas
+      letTupExp' desc $ I.Op $ I.Scatter si_w sivs lam $ zip3 sa_ws (repeat 1) sas
 
 flatIndexHelper :: String -> SrcLoc -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]
 flatIndexHelper desc loc arr offset slices = do
@@ -2164,13 +2160,9 @@
                 ++ I.varsRes (map I.paramName value_params)
         }
   results <-
-    letTupExp "partition_res" $
-      I.Op $
-        I.Scatter
-          w
-          write_lam
-          (classes : all_offsets ++ arrs)
-          $ zip3 (repeat $ Shape [w]) (repeat 1) blanks
+    letTupExp "partition_res" . I.Op $
+      I.Scatter w (classes : all_offsets ++ arrs) write_lam $
+        zip3 (repeat $ Shape [w]) (repeat 1) blanks
   sizes' <-
     letSubExp "partition_sizes" $
       I.BasicOp $
diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs
--- a/src/Futhark/Internalise/Monad.hs
+++ b/src/Futhark/Internalise/Monad.hs
@@ -35,7 +35,6 @@
 import Futhark.IR.SOACS
 import Futhark.MonadFreshNames
 import Futhark.Tools
-import Futhark.Util (takeLast)
 
 type FunInfo =
   ( [VName],
@@ -152,7 +151,7 @@
 bindConstant cname fd = do
   let stms = bodyStms $ funDefBody fd
       substs =
-        takeLast (length (funDefRetType fd)) $
+        drop (length (shapeContext (funDefRetType fd))) $
           map resSubExp $ bodyResult $ funDefBody fd
   addStms stms
   modify $ \s ->
diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs
--- a/src/Futhark/Optimise/Simplify/Engine.hs
+++ b/src/Futhark/Optimise/Simplify/Engine.hs
@@ -423,8 +423,8 @@
         vtables = scanl (flip ST.insertStm) vtable' $ stmsToList stms
 
     hoistable usageInStm (uses', stms) (stm, vtable')
-      | not $ any (`UT.isUsedDirectly` uses') $ provides stm -- Dead statement.
-        =
+      | not $ any (`UT.isUsedDirectly` uses') $ provides stm =
+        -- Dead statement.
         return (uses', stms)
       | otherwise = do
         res <-
@@ -446,7 +446,7 @@
           Just optimstms -> do
             changed
             (uses'', stms') <- simplifyStmsBottomUp' vtable' uses' optimstms
-            return (uses'', stms' ++ stms)
+            return (uses'', stms' <> stms)
 
 blockUnhoistedDeps ::
   ASTRep rep =>
@@ -473,17 +473,17 @@
   Stm rep ->
   UT.UsageTable
 expandUsage usageInStm vtable utable stm@(Let pat _ e) =
-  UT.expand (`ST.lookupAliases` vtable) (usageInStm stm <> usageThroughAliases)
-    <> ( if any (`UT.isSize` utable) (patNames pat)
-           then UT.sizeUsages (freeIn e)
-           else mempty
-       )
-    <> utable
+  stmUsages <> utable
   where
+    stmUsages =
+      UT.expand (`ST.lookupAliases` vtable) (usageInStm stm <> usageThroughAliases)
+        <> ( if any (`UT.isSize` utable) (patNames pat)
+               then UT.sizeUsages (freeIn e)
+               else mempty
+           )
     usageThroughAliases =
-      mconcat $
-        mapMaybe usageThroughBindeeAliases $
-          zip (patNames pat) (patAliases pat)
+      mconcat . mapMaybe usageThroughBindeeAliases $
+        zip (patNames pat) (patAliases pat)
     usageThroughBindeeAliases (name, aliases) = do
       uses <- UT.lookup name utable
       return $ mconcat $ map (`UT.usage` uses) $ namesToList aliases
@@ -542,11 +542,10 @@
 isNotSafe :: ASTRep rep => BlockPred rep
 isNotSafe _ _ = not . safeExp . stmExp
 
-isInPlaceBound :: BlockPred m
-isInPlaceBound _ _ = isUpdate . stmExp
+isConsuming :: Aliased rep => BlockPred rep
+isConsuming _ _ = isUpdate . stmExp
   where
-    isUpdate (BasicOp Update {}) = True
-    isUpdate _ = False
+    isUpdate e = consumedInExp e /= mempty
 
 isNotCheap :: ASTRep rep => BlockPred rep
 isNotCheap _ _ = not . cheapStm
@@ -629,7 +628,7 @@
       block =
         branch_blocker
           `orIf` ((isNotSafe `orIf` isNotCheap) `andAlso` stmIs (not . desirableToHoist))
-          `orIf` isInPlaceBound
+          `orIf` isConsuming
           `orIf` isNotHoistableBnd
 
   rules <- asksEngineEnv envRules
@@ -802,7 +801,11 @@
         nes' <- simplify nes
         return (Just (op_lam', nes'), op_lam_stms)
     (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs)
-  (lam', lam_stms) <- simplifyLambda lam
+  (lam', lam_stms) <-
+    simplifyLambdaWithBody (isFalse True) lam $
+      localVtable (ST.noteAccTokens (zip (map paramName (lambdaParams lam)) inputs')) $
+        simplifyBody (map (const Observe) (lambdaReturnType lam)) $
+          lambdaBody lam
   pure (WithAcc inputs' lam', mconcat inputs_stms <> lam_stms)
 
 -- Special case for simplification of commutative BinOps where we
@@ -850,6 +853,7 @@
     Simplifiable (LParamInfo rep),
     Simplifiable (RetType rep),
     Simplifiable (BranchType rep),
+    TraverseOpStms (Wise rep),
     CanBeWise (Op rep),
     ST.IndexOp (OpWithWisdom (Op rep)),
     BuilderOps (Wise rep),
@@ -974,14 +978,22 @@
   BlockPred (Wise rep) ->
   Lambda rep ->
   SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-simplifyLambdaMaybeHoist blocked lam@(Lambda params body rettype) = do
+simplifyLambdaMaybeHoist blocked lam =
+  simplifyLambdaWithBody blocked lam $
+    simplifyBody (map (const Observe) (lambdaReturnType lam)) $ lambdaBody lam
+
+simplifyLambdaWithBody ::
+  SimplifiableRep rep =>
+  BlockPred (Wise rep) ->
+  Lambda rep ->
+  SimpleM rep (SimplifiedBody rep Result) ->
+  SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
+simplifyLambdaWithBody blocked lam@(Lambda params _body rettype) m = do
   params' <- mapM (traverse simplify) params
   let paramnames = namesFromList $ boundByLambda lam
   ((lamstms, lamres), hoisted) <-
-    enterLoop $
-      bindLParams params' $
-        blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) $
-          simplifyBody (map (const Observe) rettype) body
+    enterLoop . bindLParams params' $
+      blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) m
   body' <- constructBody lamstms lamres
   rettype' <- simplify rettype
   return (Lambda params' body' rettype', hoisted)
diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs
--- a/src/Futhark/Optimise/Simplify/Rules.hs
+++ b/src/Futhark/Optimise/Simplify/Rules.hs
@@ -21,8 +21,9 @@
 where
 
 import Control.Monad
+import Control.Monad.State
 import Data.Either
-import Data.List (find, unzip4, zip4)
+import Data.List (find, insert, unzip4, zip4)
 import qualified Data.Map.Strict as M
 import Data.Maybe
 import Futhark.Analysis.PrimExp.Convert
@@ -45,7 +46,7 @@
     RuleGeneric withAccTopDown
   ]
 
-bottomUpRules :: BuilderOps rep => [BottomUpRule rep]
+bottomUpRules :: (BuilderOps rep, TraverseOpStms rep) => [BottomUpRule rep]
 bottomUpRules =
   [ RuleIf removeDeadBranchResult,
     RuleGeneric withAccBottomUp,
@@ -55,7 +56,7 @@
 -- | A set of standard simplification rules.  These assume pure
 -- functional semantics, and so probably should not be applied after
 -- memory block merging.
-standardRules :: (BuilderOps rep, Aliased rep) => RuleBook rep
+standardRules :: (BuilderOps rep, TraverseOpStms rep, Aliased rep) => RuleBook rep
 standardRules = ruleBook topDownRules bottomUpRules <> loopRules <> basicOpRules
 
 -- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy
@@ -353,8 +354,29 @@
       pure $ Just x
 withAccTopDown _ _ = Skip
 
-withAccBottomUp :: BuilderOps rep => BottomUpRuleGeneric rep
--- Eliminate dead results.
+elimUpdates :: (ASTRep rep, TraverseOpStms rep) => [VName] -> Body rep -> (Body rep, [VName])
+elimUpdates get_rid_of = flip runState mempty . onBody
+  where
+    onBody body = do
+      stms' <- onStms $ bodyStms body
+      pure body {bodyStms = stms'}
+    onStms = traverse onStm
+    onStm (Let pat@(Pat [PatElem _ dec]) aux (BasicOp (UpdateAcc acc _ _)))
+      | Acc c _ _ _ <- typeOf dec,
+        c `elem` get_rid_of = do
+        modify (insert c)
+        pure $ Let pat aux $ BasicOp $ SubExp $ Var acc
+    onStm (Let pat aux e) = Let pat aux <$> onExp e
+    onExp = mapExpM mapper
+      where
+        mapper =
+          identityMapper
+            { mapOnOp = traverseOpStms (\_ stms -> onStms stms),
+              mapOnBody = \_ body -> onBody body
+            }
+
+withAccBottomUp :: (TraverseOpStms rep, BuilderOps rep) => BottomUpRuleGeneric rep
+-- Eliminate dead results.  See Note [Dead Code Elimination for WithAcc]
 withAccBottomUp (_, utable) (Let pat aux (WithAcc inputs lam))
   | not $ all (`UT.used` utable) $ patNames pat = Simplify $ do
     let (acc_res, nonacc_res) =
@@ -365,32 +387,33 @@
           splitAt (length inputs) $ lambdaParams lam
 
     -- Eliminate unused accumulator results
-    let (acc_pes', inputs', param_pairs, acc_res') =
-          unzip4 . filter keepAccRes $
-            zip4
+    let get_rid_of =
+          map snd . filter getRidOf $
+            zip
               (chunks (map inputArrs inputs) acc_pes)
-              inputs
-              (zip cert_params acc_params)
-              acc_res
-        (cert_params', acc_params') = unzip param_pairs
+              $ map paramName cert_params
 
     -- Eliminate unused non-accumulator results
     let (nonacc_pes', nonacc_res') =
           unzip $ filter keepNonAccRes $ zip nonacc_pes nonacc_res
 
-    when (concat acc_pes' == acc_pes && nonacc_pes' == nonacc_pes) cannotSimplify
+    when (null get_rid_of && nonacc_pes' == nonacc_pes) cannotSimplify
 
-    let pes' = concat acc_pes' ++ nonacc_pes'
+    let (body', eliminated) = elimUpdates get_rid_of $ lambdaBody lam
 
-    lam' <- mkLambda (cert_params' ++ acc_params') $ do
-      void $ bodyBind $ lambdaBody lam
-      pure $ acc_res' ++ nonacc_res'
+    when (null eliminated && nonacc_pes' == nonacc_pes) cannotSimplify
 
-    auxing aux $ letBind (Pat pes') $ WithAcc inputs' lam'
+    let pes' = acc_pes ++ nonacc_pes'
+
+    lam' <- mkLambda (cert_params ++ acc_params) $ do
+      void $ bodyBind body'
+      pure $ acc_res ++ nonacc_res'
+
+    auxing aux $ letBind (Pat pes') $ WithAcc inputs lam'
   where
     num_nonaccs = length (lambdaReturnType lam) - length inputs
     inputArrs (_, arrs, _) = length arrs
-    keepAccRes (pes, _, _, _) = any ((`UT.used` utable) . patElemName) pes
+    getRidOf (pes, _) = not $ any ((`UT.used` utable) . patElemName) pes
     keepNonAccRes (pe, _) = patElemName pe `UT.used` utable
 withAccBottomUp _ _ = Skip
 
@@ -403,3 +426,46 @@
 isCt0 :: SubExp -> Bool
 isCt0 (Constant v) = zeroIsh v
 isCt0 _ = False
+
+-- Note [Dead Code Elimination for WithAcc]
+--
+-- Our static semantics for accumulators are basically those of linear
+-- types.  This makes dead code elimination somewhat tricky.  First,
+-- what we consider dead code is when we have a WithAcc where at least
+-- one of the array results (that internally correspond to an
+-- accumulator) are unused.  E.g
+--
+-- let {X',Y'} =
+--   with_acc {X, Y} (\X_p Y_p X_acc Y_acc -> ... {X_acc', Y_acc'})
+--
+-- where X' is not used later.  Note that Y' is still used later.  If
+-- none of the results of the WithAcc are used, then the Stm as a
+-- whole is dead and can be removed.  That's the trivial case, done
+-- implicitly by the simplifier.  The interesting case is exactly when
+-- some of the results are unused.  How do we get rid of them?
+--
+-- Naively, we might just remove them:
+--
+-- let Y' =
+--   with_acc Y (\Y_p Y_acc -> ... Y_acc')
+--
+-- This is safe *only* if X_acc is used *only* in the result (i.e. an
+-- "identity" WithAcc).  Otherwise we end up with references to X_acc,
+-- which no longer exists.  This simple case is actually handled in
+-- the withAccTopDown rule, and is easy enough.
+--
+-- What we actually do when we decide to eliminate X_acc is that we
+-- inspect the body of the WithAcc and eliminate all UpdateAcc
+-- operations that refer to the same accumulator as X_acc (identified
+-- by the X_p token).  I.e. we turn every
+--
+-- let B = update_acc(A, ...)
+--
+-- where 'A' is ultimately decided from X_acc into
+--
+-- let B = A
+--
+-- That's it!  We then let ordinary dead code elimination eventually
+-- simplify the body enough that we have an "identity" WithAcc.  There
+-- is no _guarantee_ that this will happen, but our general dead code
+-- elimination tends to be pretty good.
diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs
--- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs
+++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs
@@ -382,6 +382,14 @@
     cs' /= cs =
     Simplify . certifying (Certs cs') $
       letBind pat $ BasicOp $ SubExp $ Var v
+-- Remove UpdateAccs that contribute the neutral value, which is
+-- always a no-op.
+ruleBasicOp vtable pat aux (UpdateAcc acc _ vs)
+  | Pat [pe] <- pat,
+    Acc token _ _ _ <- patElemType pe,
+    Just (_, _, Just (_, ne)) <- ST.entryAccInput =<< ST.lookup token vtable,
+    vs == ne =
+    Simplify . auxing aux $ letBind pat $ BasicOp $ SubExp $ Var acc
 ruleBasicOp _ _ _ _ =
   Skip
 
diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs
--- a/src/Futhark/Optimise/Simplify/Rules/Index.hs
+++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs
@@ -53,6 +53,7 @@
       | Just inds' <- sliceIndices (Slice inds),
         Just (ST.IndexedArray cs arr inds'') <- ST.index idd inds' vtable,
         all (worthInlining . untyped) inds'',
+        arr `ST.available` vtable,
         all (`ST.elem` vtable) (unCerts cs) ->
         Just $
           IndexResult cs arr . Slice . map DimFix
diff --git a/src/Futhark/Optimise/Simplify/Rules/Loop.hs b/src/Futhark/Optimise/Simplify/Rules/Loop.hs
--- a/src/Futhark/Optimise/Simplify/Rules/Loop.hs
+++ b/src/Futhark/Optimise/Simplify/Rules/Loop.hs
@@ -117,9 +117,7 @@
     checkInvariance
       (pat_name, (mergeParam, mergeInit), resExp)
       (invariant, explpat', merge', resExps)
-        | not (unique (paramDeclType mergeParam))
-            || arrayRank (paramDeclType mergeParam) == 1,
-          isInvariant,
+        | isInvariant,
           -- Also do not remove the condition in a while-loop.
           not $ paramName mergeParam `nameIn` freeIn form =
           let (stm, explpat'') =
diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs
--- a/src/Futhark/Pass/ExtractKernels.hs
+++ b/src/Futhark/Pass/ExtractKernels.hs
@@ -541,7 +541,7 @@
   types <- asksScope scopeForSOACs
   transformStms path . stmsToList . snd
     =<< runBuilderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types
-transformStm _ (Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) = runBuilder_ $ do
+transformStm _ (Let pat (StmAux cs _ _) (Op (Scatter w ivs lam as))) = runBuilder_ $ do
   let lam' = soacsLambdaToGPU lam
   write_i <- newVName "write_i"
   let (as_ws, _, _) = unzip3 as
@@ -567,7 +567,7 @@
   certifying cs $ do
     addStms stms
     letBind pat $ Op $ SegOp kernel
-transformStm _ (Let orig_pat (StmAux cs _ _) (Op (Hist w ops bucket_fun imgs))) = do
+transformStm _ (Let orig_pat (StmAux cs _ _) (Op (Hist w imgs ops bucket_fun))) = do
   let bfun' = soacsLambdaToGPU bucket_fun
 
   -- It is important not to launch unnecessarily many threads for
@@ -604,7 +604,7 @@
       | Op (Screma w _ form) <- stmExp stm,
         Just lam' <- isMapSOAC form =
         mapLike w lam'
-      | Op (Scatter w lam' _ _) <- stmExp stm =
+      | Op (Scatter w _ lam' _) <- stmExp stm =
         mapLike w lam'
       | DoLoop _ _ body <- stmExp stm =
         bodyInterest body * 10
diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs
--- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs
+++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs
@@ -468,7 +468,7 @@
     distributeMapBodyStms acc stms
 
 -- Parallelise segmented scatters.
-maybeDistributeStm stm@(Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) acc =
+maybeDistributeStm stm@(Let pat (StmAux cs _ _) (Op (Scatter w ivs lam as))) acc =
   distributeSingleStm acc stm >>= \case
     Just (kernels, res, nest, acc')
       | Just (perm, pat_unused) <- permutationAndMissing pat res ->
@@ -481,7 +481,7 @@
     _ ->
       addStmToAcc stm acc
 -- Parallelise segmented Hist.
-maybeDistributeStm stm@(Let pat (StmAux cs _ _) (Op (Hist w ops lam as))) acc =
+maybeDistributeStm stm@(Let pat (StmAux cs _ _) (Op (Hist w as ops lam))) acc =
   distributeSingleStm acc stm >>= \case
     Just (kernels, res, nest, acc')
       | Just (perm, pat_unused) <- permutationAndMissing pat res ->
@@ -523,9 +523,8 @@
           localScope (typeEnvFromDistAcc acc') $ do
             nest' <- expandKernelNest pat_unused nest
             map_lam' <- soacsLambda map_lam
-            lam' <- soacsLambda lam
             localScope (typeEnvFromDistAcc acc') $
-              segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs
+              segmentedScanomapKernel nest' perm w lam map_lam' nes arrs
                 >>= kernelOrNot cs stm acc kernels acc'
       _ ->
         addStmToAcc stm acc
@@ -1019,16 +1018,20 @@
   KernelNest ->
   [Int] ->
   SubExp ->
-  Lambda rep ->
+  Lambda SOACS ->
   Lambda rep ->
   [SubExp] ->
   [VName] ->
   DistNestT rep m (Maybe (Stms rep))
 segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do
   mk_lvl <- asks distSegLevel
+  onLambda <- asks distOnSOACSLambda
+  let onLambda' = fmap fst . runBuilder . onLambda
   isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $
     \pat ispace inps nes' _ -> do
-      let scan_op = SegBinOp Noncommutative lam nes' mempty
+      (lam', nes'', shape) <- determineReduceOp lam nes'
+      lam'' <- onLambda' lam'
+      let scan_op = SegBinOp Noncommutative lam'' nes'' shape
       lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt
       addStms =<< traverse renameStm
         =<< segScan lvl pat segment_size [scan_op] map_lam arrs ispace inps
diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs
--- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs
+++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs
@@ -261,7 +261,7 @@
         certifying (stmAuxCerts aux) $
           addStms =<< segRed lvl' pat w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []
         parallelMin [w]
-    Op (Hist w ops bucket_fun arrs) -> do
+    Op (Hist w arrs ops bucket_fun) -> do
       ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do
         (op', nes', shape) <- determineReduceOp op nes
         let op'' = soacsLambdaToGPU op'
@@ -281,7 +281,7 @@
             replaceSets (IntraAcc x y log) =
               IntraAcc (S.map (map replace) x) (S.map (map replace) y) log
         censor replaceSets $ intraGroupStms lvl stream_stms
-    Op (Scatter w lam ivs dests) -> do
+    Op (Scatter w ivs lam dests) -> do
       write_i <- newVName "write_i"
       space <- mkSegSpace [(write_i, w)]
 
diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs
--- a/src/Futhark/Pass/ExtractMulticore.hs
+++ b/src/Futhark/Pass/ExtractMulticore.hs
@@ -320,7 +320,7 @@
     -- anything, so split it up and try again.
     scope <- castScope <$> askScope
     transformStms =<< runBuilderT_ (dissectScrema pat w form arrs) scope
-transformSOAC pat _ (Scatter w lam ivs dests) = do
+transformSOAC pat _ (Scatter w ivs lam dests) = do
   (gtid, space) <- mkSegSpace w
 
   Body () kstms res <- mapLambdaToBody transformBody gtid lam ivs
@@ -340,7 +340,7 @@
         Op $
           ParOp Nothing $
             SegMap () space rets kbody
-transformSOAC pat _ (Hist w hists map_lam arrs) = do
+transformSOAC pat _ (Hist w arrs hists map_lam) = do
   (seq_hist_stms, seq_op) <-
     transformHist DoNotRename sequentialiseBody w hists map_lam arrs
 
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
@@ -87,10 +87,10 @@
 
 -- | How a program can be transformed.
 data StructurePipeline
-  = KernelsPipeline
+  = GpuPipeline
   | SOACSPipeline
-  | SequentialCpuPipeline
-  | GpuPipeline
+  | SeqMemPipeline
+  | GpuMemPipeline
   | NoPipeline
   deriving (Show)
 
@@ -340,9 +340,9 @@
 optimisePipeline :: Parser () -> Parser StructurePipeline
 optimisePipeline sep =
   choice
-    [ lexeme sep "distributed" $> KernelsPipeline,
+    [ lexeme sep "gpu-mem" $> GpuMemPipeline,
       lexeme sep "gpu" $> GpuPipeline,
-      lexeme sep "cpu" $> SequentialCpuPipeline,
+      lexeme sep "seq-mem" $> SeqMemPipeline,
       lexeme sep "internalised" $> NoPipeline,
       pure SOACSPipeline
     ]
diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs
--- a/src/Futhark/Transform/FirstOrderTransform.hs
+++ b/src/Futhark/Transform/FirstOrderTransform.hs
@@ -274,7 +274,7 @@
       mkBodyM mempty $ res ++ subExpsRes mapout_res'
 
   letBind pat $ DoLoop merge loop_form loop_body
-transformSOAC pat (Scatter len lam ivs as) = do
+transformSOAC pat (Scatter len ivs lam as) = do
   iter <- newVName "write_iter"
 
   let (as_ws, as_ns, as_vs) = unzip3 as
@@ -301,7 +301,7 @@
         foldM saveInArray arr indexes'
       return $ resultBody (map Var ress)
   letBind pat $ DoLoop merge (ForLoop iter Int64 len []) loopBody
-transformSOAC pat (Hist len ops bucket_fun imgs) = do
+transformSOAC pat (Hist len imgs ops bucket_fun) = do
   iter <- newVName "iter"
 
   -- Bind arguments to parameters for the merge-variables.
@@ -343,8 +343,7 @@
 
           -- Apply operator.
           h_val' <-
-            bindLambda (histOp op) $
-              map (BasicOp . SubExp) $ h_val ++ val
+            bindLambda (histOp op) $ map (BasicOp . SubExp) $ h_val ++ val
 
           -- Write values back to histograms.
           hist' <- forM (zip hist h_val') $ \(arr, SubExpRes cs v) -> do
