diff --git a/docs/man/futhark-bench.rst b/docs/man/futhark-bench.rst
--- a/docs/man/futhark-bench.rst
+++ b/docs/man/futhark-bench.rst
@@ -124,7 +124,8 @@
 
   Enable profiling for the binary (by passing ``--profiling`` and
   ``--logging``) and store the recorded information in the file
-  indicated by ``--json``, along with the other benchmarking results .
+  indicated by ``--json`` (which is required), along with the other
+  benchmarking results.
 
 --runner=program
 
diff --git a/docs/man/futhark-profile.rst b/docs/man/futhark-profile.rst
--- a/docs/man/futhark-profile.rst
+++ b/docs/man/futhark-profile.rst
@@ -26,12 +26,11 @@
 =====
 
 The first step is to run :ref:`futhark bench<futhark-bench(1)>` on
-your program, while passing ``--profile`` and ``--json``.  This will
+your program, while passing ``--profile`` and ``--json``. This will
 produce a JSON file containing runtime measurements, as well as
-collected profiling information.  If you neglect to pass
-``--profile``, the latter will be missing.  If you neglect to pass
-``--json``, no file will be created.  The information in the JSON file
-is complete, but it is difficult for humans to read.
+collected profiling information. If you neglect to pass ``--profile``,
+the profiling information will be missing. The information in the JSON
+file is complete, but it is difficult for humans to read.
 
 The next step is to run ``futhark profile`` on the JSON file.  For a
 JSON file ``prog.json``, this will create a *top level directory*
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -1,6 +1,6 @@
 cabal-version: 2.4
 name:           futhark
-version:        0.25.6
+version:        0.25.7
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -309,7 +309,7 @@
       Futhark.IR.Prop.Aliases
       Futhark.IR.Prop.Constants
       Futhark.IR.Prop.Names
-      Futhark.IR.Prop.Patterns
+      Futhark.IR.Prop.Pat
       Futhark.IR.Prop.Rearrange
       Futhark.IR.Prop.Reshape
       Futhark.IR.Prop.Scope
diff --git a/rts/c/server.h b/rts/c/server.h
--- a/rts/c/server.h
+++ b/rts/c/server.h
@@ -969,6 +969,7 @@
   size_t obj_size;
   void *data = NULL;
   (void)aux->store(ctx, obj, &data, &obj_size);
+  assert(futhark_context_sync(ctx) == 0);
   fwrite(data, sizeof(char), obj_size, f);
   free(data);
 }
diff --git a/src/Futhark/AD/Rev/Hist.hs b/src/Futhark/AD/Rev/Hist.hs
--- a/src/Futhark/AD/Rev/Hist.hs
+++ b/src/Futhark/AD/Rev/Hist.hs
@@ -8,6 +8,7 @@
   ( diffMinMaxHist,
     diffMulHist,
     diffAddHist,
+    diffVecHist,
     diffHist,
   )
 where
@@ -521,6 +522,64 @@
 
   vs_bar <- letExp (baseString vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar
   updateAdj vs vs_bar
+
+-- Special case for vectorised combining operator. Rewrite
+--   reduce_by_index dst (map2 op) nes is vss
+-- to
+--   map3 (\dst_col vss_col ne ->
+--           reduce_by_index dst_col op ne is vss_col
+--        ) (transpose dst) (transpose vss) nes |> transpose
+-- before differentiating.
+diffVecHist ::
+  VjpOps ->
+  VName ->
+  StmAux () ->
+  SubExp ->
+  Lambda SOACS ->
+  VName ->
+  VName ->
+  VName ->
+  SubExp ->
+  SubExp ->
+  VName ->
+  ADM () ->
+  ADM ()
+diffVecHist ops x aux n op nes is vss w rf dst m = do
+  stms <- collectStms_ $ do
+    rank <- arrayRank <$> lookupType vss
+    let dims = [1, 0] ++ drop 2 [0 .. rank - 1]
+
+    dstT <- letExp "dstT" $ BasicOp $ Rearrange dims dst
+    vssT <- letExp "vssT" $ BasicOp $ Rearrange dims vss
+    t_dstT <- lookupType dstT
+    t_vssT <- lookupType vssT
+    t_nes <- lookupType nes
+
+    dst_col <- newParam "dst_col" $ rowType t_dstT
+    vss_col <- newParam "vss_col" $ rowType t_vssT
+    ne <- newParam "ne" $ rowType t_nes
+
+    f <- mkIdentityLambda (Prim int64 : lambdaReturnType op)
+    map_lam <-
+      mkLambda [dst_col, vss_col, ne] $ do
+        -- TODO Have to copy dst_col, but isn't it already unique?
+        dst_col_cpy <-
+          letExp "dst_col_cpy" . BasicOp $
+            Replicate mempty (Var $ paramName dst_col)
+        fmap (varsRes . pure) . letExp "col_res" $
+          Op $
+            Hist
+              n
+              [is, paramName vss_col]
+              [HistOp (Shape [w]) rf [dst_col_cpy] [Var $ paramName ne] op]
+              f
+    histT <-
+      letExp "histT" $
+        Op $
+          Screma (arraySize 0 t_dstT) [dstT, vssT, nes] $
+            mapSOAC map_lam
+    auxing aux . letBindNames [x] . BasicOp $ Rearrange dims histT
+  foldr (vjpStm ops) m stms
 
 --
 -- a step in the radix sort implementation
diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs
--- a/src/Futhark/AD/Rev/SOAC.hs
+++ b/src/Futhark/AD/Rev/SOAC.hs
@@ -158,6 +158,12 @@
 vjpSOAC ops pat aux (Hist n [is, vs] [histop] f) m
   | isIdentityLambda f,
     [x] <- patNames pat,
+    HistOp (Shape [w]) rf [dst] [Var ne] lam <- histop,
+    -- Note that the operator is vectorised, so `ne` cannot be a 'PrimValue'.
+    Just op <- mapOp lam =
+      diffVecHist ops x aux n op ne is vs w rf dst m
+  | isIdentityLambda f,
+    [x] <- patNames pat,
     HistOp (Shape [w]) rf [dst] [ne] lam <- histop,
     lam' <- nestedMapOp lam,
     Just [(op, _, _, _)] <- lamIsBinOp lam',
diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs
--- a/src/Futhark/Analysis/DataDependencies.hs
+++ b/src/Futhark/Analysis/DataDependencies.hs
@@ -49,7 +49,9 @@
           reductionDependencies deps lam' nes (depsOfArrays' shape arrs)
     grow deps (Let pat _ (Op op)) =
       let op_deps = map (depsOfNames deps) (opDependencies op)
-       in M.fromList (zip (patNames pat) op_deps) `M.union` deps
+          pat_deps = map (depsOfNames deps . freeIn) (patElems pat)
+       in M.fromList (zip (patNames pat) $ zipWith (<>) pat_deps op_deps)
+            `M.union` deps
     grow deps (Let pat _ (Match c cases defbody _)) =
       let cases_deps = map (dataDependencies' deps . caseBody) cases
           defbody_deps = dataDependencies' deps defbody
diff --git a/src/Futhark/CLI/Autotune.hs b/src/Futhark/CLI/Autotune.hs
--- a/src/Futhark/CLI/Autotune.hs
+++ b/src/Futhark/CLI/Autotune.hs
@@ -472,7 +472,7 @@
   T.putStrLn $ "Result of autotuning:\n" <> tuning
 
 supportedBackends :: [String]
-supportedBackends = ["opencl", "cuda"]
+supportedBackends = ["opencl", "cuda", "hip"]
 
 commandLineOptions :: [FunOptDescr AutotuneOptions]
 commandLineOptions =
diff --git a/src/Futhark/CLI/Bench.hs b/src/Futhark/CLI/Bench.hs
--- a/src/Futhark/CLI/Bench.hs
+++ b/src/Futhark/CLI/Bench.hs
@@ -643,4 +643,8 @@
 main = mainWithOptions initialBenchOptions commandLineOptions "options... programs..." $ \progs config ->
   case progs of
     [] -> Nothing
-    _ -> Just $ runBenchmarks (excludeBackend config) progs
+    _
+      | optProfile config && isNothing (optJSON config) ->
+          Just $ optionsError "--profile cannot be used without --json."
+      | otherwise ->
+          Just $ runBenchmarks (excludeBackend config) progs
diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs
--- a/src/Futhark/CLI/Dev.hs
+++ b/src/Futhark/CLI/Dev.hs
@@ -618,6 +618,7 @@
     iplOption [],
     allocateOption "a",
     kernelsMemPassOption doubleBufferGPU [],
+    mcMemPassOption doubleBufferMC [],
     kernelsMemPassOption expandAllocations [],
     kernelsMemPassOption MemoryBlockMerging.optimise [],
     seqMemPassOption LiftAllocations.liftAllocationsSeqMem [],
diff --git a/src/Futhark/CLI/Literate.hs b/src/Futhark/CLI/Literate.hs
--- a/src/Futhark/CLI/Literate.hs
+++ b/src/Futhark/CLI/Literate.hs
@@ -808,7 +808,7 @@
       void $ system "gnuplot" [] script'
 --
 processDirective env (DirectiveVideo e params) = do
-  when (format `notElem` ["webm", "gif"]) $
+  unless (format `elem` ["webm", "gif"]) $
     throwError $
       "Unknown video format: " <> format
 
diff --git a/src/Futhark/CLI/REPL.hs b/src/Futhark/CLI/REPL.hs
--- a/src/Futhark/CLI/REPL.hs
+++ b/src/Futhark/CLI/REPL.hs
@@ -411,7 +411,8 @@
       then
         annotate italicized $
           "\n\nPolymorphic in"
-            <+> mconcat (intersperse " " $ map pretty ps) <> "."
+            <+> mconcat (intersperse " " $ map pretty ps)
+            <> "."
       else mempty
 
 mtypeCommand :: Command
diff --git a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs
--- a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs
+++ b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs
@@ -127,26 +127,26 @@
     "if"
       <+> pretty cond
       <> ":"
-      </> indent 2 "pass"
+        </> indent 2 "pass"
   pretty (If cond [] fbranch) =
     "if"
       <+> pretty cond
       <> ":"
-      </> indent 2 "pass"
-      </> "else:"
-      </> indent 2 (stack $ map pretty fbranch)
+        </> indent 2 "pass"
+        </> "else:"
+        </> indent 2 (stack $ map pretty fbranch)
   pretty (If cond tbranch []) =
     "if"
       <+> pretty cond
       <> ":"
-      </> indent 2 (stack $ map pretty tbranch)
+        </> indent 2 (stack $ map pretty tbranch)
   pretty (If cond tbranch fbranch) =
     "if"
       <+> pretty cond
       <> ":"
-      </> indent 2 (stack $ map pretty tbranch)
-      </> "else:"
-      </> indent 2 (stack $ map pretty fbranch)
+        </> indent 2 (stack $ map pretty tbranch)
+        </> "else:"
+        </> indent 2 (stack $ map pretty fbranch)
   pretty (Try pystms pyexcepts) =
     "try:"
       </> indent 2 (stack $ map pretty pystms)
@@ -155,19 +155,19 @@
     "while"
       <+> pretty cond
       <> ":"
-      </> indent 2 (stack $ map pretty body)
+        </> indent 2 (stack $ map pretty body)
   pretty (For i what body) =
     "for"
       <+> pretty i
       <+> "in"
       <+> pretty what
       <> ":"
-      </> indent 2 (stack $ map pretty body)
+        </> indent 2 (stack $ map pretty body)
   pretty (With what body) =
     "with"
       <+> pretty what
       <> ":"
-      </> indent 2 (stack $ map pretty body)
+        </> indent 2 (stack $ map pretty body)
   pretty (Assign e1 e2) = pretty e1 <+> "=" <+> pretty e2
   pretty (AssignOp op e1 e2) = pretty e1 <+> pretty (op ++ "=") <+> pretty e2
   pretty (Comment s body) = "#" <> pretty s </> stack (map pretty body)
@@ -190,14 +190,14 @@
       <+> pretty fname
       <> parens (commasep $ map pretty params)
       <> ":"
-      </> indent 2 (stack (map pretty body))
+        </> indent 2 (stack (map pretty body))
 
 instance Pretty PyClassDef where
   pretty (Class cname body) =
     "class"
       <+> pretty cname
       <> ":"
-      </> indent 2 (stack (map pretty body))
+        </> indent 2 (stack (map pretty body))
 
 instance Pretty PyExcept where
   pretty (Catch pyexp stms) =
diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs
--- a/src/Futhark/CodeGen/ImpCode.hs
+++ b/src/Futhark/CodeGen/ImpCode.hs
@@ -574,9 +574,15 @@
   pretty (Free name space) =
     "free" <> parens (pretty name) <> pretty space
   pretty (Write name i bt space vol val) =
-    pretty name <> langle <> vol' <> pretty bt <> pretty space <> rangle <> brackets (pretty i)
-      <+> "<-"
-      <+> pretty val
+    pretty name
+      <> langle
+      <> vol'
+      <> pretty bt
+      <> pretty space
+      <> rangle
+      <> brackets (pretty i)
+        <+> "<-"
+        <+> pretty val
     where
       vol' = case vol of
         Volatile -> "volatile "
@@ -584,7 +590,13 @@
   pretty (Read name v is bt space vol) =
     pretty name
       <+> "<-"
-      <+> pretty v <> langle <> vol' <> pretty bt <> pretty space <> rangle <> brackets (pretty is)
+      <+> pretty v
+      <> langle
+      <> vol'
+      <> pretty bt
+      <> pretty space
+      <> rangle
+      <> brackets (pretty is)
     where
       vol' = case vol of
         Volatile -> "volatile "
@@ -602,14 +614,17 @@
       <> (parens . align)
         ( foldMap (brackets . pretty) shape
             <> ","
-            </> p dst dstspace dstoffset dststrides
+              </> p dst dstspace dstoffset dststrides
             <> ","
-            </> p src srcspace srcoffset srcstrides
+              </> p src srcspace srcoffset srcstrides
         )
     where
       p mem space offset strides =
-        pretty mem <> pretty space <> "+" <> pretty offset
-          <+> foldMap (brackets . pretty) strides
+        pretty mem
+          <> pretty space
+          <> "+"
+          <> pretty offset
+            <+> foldMap (brackets . pretty) strides
   pretty (If cond tbranch fbranch) =
     "if"
       <+> pretty cond
@@ -626,7 +641,8 @@
     "call"
       <+> commasep (map pretty dests)
       <+> "<-"
-      <+> pretty fname <> parens (commasep $ map pretty args)
+      <+> pretty fname
+      <> parens (commasep $ map pretty args)
   pretty (Comment s code) =
     "--" <+> pretty s </> pretty code
   pretty (DebugPrint desc (Just e)) =
diff --git a/src/Futhark/CodeGen/ImpCode/GPU.hs b/src/Futhark/CodeGen/ImpCode/GPU.hs
--- a/src/Futhark/CodeGen/ImpCode/GPU.hs
+++ b/src/Futhark/CodeGen/ImpCode/GPU.hs
@@ -105,15 +105,17 @@
   pretty (GetSize dest key size_class) =
     pretty dest
       <+> "<-"
-      <+> "get_size" <> parens (commasep [pretty key, pretty size_class])
+      <+> "get_size"
+      <> parens (commasep [pretty key, pretty size_class])
   pretty (GetSizeMax dest size_class) =
     pretty dest <+> "<-" <+> "get_size_max" <> parens (pretty size_class)
   pretty (CmpSizeLe dest name size_class x) =
     pretty dest
       <+> "<-"
-      <+> "get_size" <> parens (commasep [pretty name, pretty size_class])
-      <+> "<"
-      <+> pretty x
+      <+> "get_size"
+      <> parens (commasep [pretty name, pretty size_class])
+        <+> "<"
+        <+> pretty x
   pretty (CallKernel c) =
     pretty c
 
@@ -211,15 +213,18 @@
   pretty (GetGroupId dest i) =
     pretty dest
       <+> "<-"
-      <+> "get_group_id" <> parens (pretty i)
+      <+> "get_group_id"
+      <> parens (pretty i)
   pretty (GetLocalId dest i) =
     pretty dest
       <+> "<-"
-      <+> "get_local_id" <> parens (pretty i)
+      <+> "get_local_id"
+      <> parens (pretty i)
   pretty (GetLocalSize dest i) =
     pretty dest
       <+> "<-"
-      <+> "get_local_size" <> parens (pretty i)
+      <+> "get_local_size"
+      <> parens (pretty i)
   pretty (GetLockstepWidth dest) =
     pretty dest
       <+> "<-"
@@ -242,68 +247,68 @@
     pretty old
       <+> "<-"
       <+> "atomic_add_"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicFAdd t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_fadd_"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicSMax t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_smax"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicSMin t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_smin"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicUMax t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_umax"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicUMin t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_umin"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicAnd t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_and"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicOr t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_or"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicXor t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_xor"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
   pretty (Atomic _ (AtomicCmpXchg t old arr ind x y)) =
     pretty old
       <+> "<-"
       <+> "atomic_cmp_xchg"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x, pretty y])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x, pretty y])
   pretty (Atomic _ (AtomicXchg t old arr ind x)) =
     pretty old
       <+> "<-"
       <+> "atomic_xchg"
-        <> pretty t
-        <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
+      <> pretty t
+      <> parens (commasep [pretty arr <> brackets (pretty ind), pretty x])
 
 instance FreeIn KernelOp where
   freeIn' (Atomic _ op) = freeIn' op
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
@@ -1345,7 +1345,7 @@
 
     let params =
           [ Imp.MemParam mem (Space "device"),
-            Imp.ScalarParam n int32,
+            Imp.ScalarParam n int64,
             Imp.ScalarParam x $ IntType bt,
             Imp.ScalarParam s $ IntType bt
           ]
@@ -1358,7 +1358,7 @@
       arr <-
         sArray "arr" (IntType bt) shape mem $
           LMAD.iota 0 (map pe64 (shapeDims shape))
-      sIotaKernel arr (sExt64 n') x' s' bt
+      sIotaKernel arr n' x' s' bt
 
   pure fname
 
@@ -1392,8 +1392,9 @@
 compileThreadResult space pe (Returns _ _ what) = do
   let is = map (Imp.le64 . fst) $ unSegSpace space
   copyDWIMFix (patElemName pe) is what []
-compileThreadResult _ pe (WriteReturns _ (Shape rws) _arr dests) = do
-  let rws' = map pe64 rws
+compileThreadResult _ pe (WriteReturns _ arr dests) = do
+  arr_t <- lookupType arr
+  let rws' = map pe64 $ arrayDims arr_t
   forM_ dests $ \(slice, e) -> do
     let slice' = fmap pe64 slice
         write = inBounds slice' rws'
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Group.hs b/src/Futhark/CodeGen/ImpGen/GPU/Group.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/Group.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/Group.hs
@@ -201,10 +201,11 @@
 
 -- Construct the necessary lock arrays for an intra-group histogram.
 prepareIntraGroupSegHist ::
+  Shape ->
   Count GroupSize SubExp ->
   [HistOp GPUMem] ->
   InKernelGen [[Imp.TExp Int64] -> InKernelGen ()]
-prepareIntraGroupSegHist group_size =
+prepareIntraGroupSegHist segments group_size =
   fmap snd . mapAccumLM onOp Nothing
   where
     onOp l op = do
@@ -221,7 +222,7 @@
           locks <- newVName "locks"
 
           let num_locks = pe64 $ unCount group_size
-              dims = map pe64 $ shapeDims (histOpShape op <> histShape op)
+              dims = map pe64 $ shapeDims (segments <> histOpShape op <> histShape op)
               l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims)
               locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness
 
@@ -517,7 +518,7 @@
       sOp $ Imp.Barrier Imp.FenceLocal
 compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do
   compileFlatId space
-  let (ltids, _dims) = unzip $ unSegSpace space
+  let (ltids, dims) = unzip $ unSegSpace space
 
   -- We don't need the red_pes, because it is guaranteed by our type
   -- rules that they occupy the same memory as the destinations for
@@ -527,7 +528,7 @@
         splitAt num_red_res $ patElems pat
 
   group_size <- kernelGroupSizeCount . kernelConstants <$> askEnv
-  ops' <- prepareIntraGroupSegHist group_size ops
+  ops' <- prepareIntraGroupSegHist (Shape $ init dims) group_size ops
 
   -- Ensure that all locks have been initialised.
   sOp $ Imp.Barrier Imp.FenceLocal
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs
@@ -460,8 +460,10 @@
                       dest_shape' = map pe64 $ shapeDims dest_shape
                       flat_bucket = flattenIndex dest_shape' bucket'
                       bucket_in_bounds =
-                        chk_beg .<=. flat_bucket
-                          .&&. flat_bucket .<. (chk_beg + hist_H_chk)
+                        chk_beg
+                          .<=. flat_bucket
+                          .&&. flat_bucket
+                          .<. (chk_beg + hist_H_chk)
                           .&&. inBounds (Slice (map DimFix bucket')) dest_shape'
                       vs_params = takeLast (length vs') $ lambdaParams lam
 
@@ -760,8 +762,10 @@
                       flat_bucket = flattenIndex dest_shape' bucket'
                       bucket_in_bounds =
                         inBounds (Slice (map DimFix bucket')) dest_shape'
-                          .&&. chk_beg .<=. flat_bucket
-                          .&&. flat_bucket .<. (chk_beg + tvExp hist_H_chk)
+                          .&&. chk_beg
+                          .<=. flat_bucket
+                          .&&. flat_bucket
+                          .<. (chk_beg + tvExp hist_H_chk)
                       bucket_is =
                         [sExt64 thread_local_subhisto_i, flat_bucket - chk_beg]
                       vs_params = takeLast (length vs') $ lambdaParams lam
@@ -1025,11 +1029,14 @@
   -- asymptotically efficient.  This mostly matters for the segmented
   -- case.
   let pick_local =
-        hist_Nin .>=. hist_H
+        hist_Nin
+          .>=. hist_H
           .&&. (local_mem_needed .<=. tvExp hist_L)
           .&&. (hist_S .<=. max_S)
-          .&&. hist_C .<=. hist_B
-          .&&. tvExp hist_M .>. 0
+          .&&. hist_C
+          .<=. hist_B
+          .&&. tvExp hist_M
+          .>. 0
 
       run = do
         emit $ Imp.DebugPrint "## Using local memory" Nothing
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
@@ -319,8 +319,8 @@
               .>. 0
               .&&. isActive (init $ zip gtids dims)
               .&&. ltid
-                .<. segment_size
-                  * segments_per_group
+              .<. segment_size
+              * segments_per_group
           )
           in_bounds
           out_of_bounds
@@ -345,8 +345,8 @@
           ( sExt64 group_id'
               * segments_per_group
               + sExt64 ltid
-              .<. num_segments
-              .&&. ltid
+                .<. num_segments
+                .&&. ltid
                 .<. segments_per_group
           )
         $ forM_ (zip segred_pes (concat reds_arrs))
@@ -603,7 +603,7 @@
     is_last_thread =
       Imp.unCount num_elements
         .<. (thread_index + 1)
-          * Imp.unCount elements_per_thread
+        * Imp.unCount elements_per_thread
 
 reductionStageZero ::
   KernelConstants ->
diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs
--- a/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs
+++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs
@@ -18,9 +18,10 @@
   MulticoreGen ()
 writeResult is pe (Returns _ _ se) =
   copyDWIMFix (patElemName pe) (map Imp.le64 is) se []
-writeResult _ pe (WriteReturns _ (Shape rws) _ idx_vals) = do
+writeResult _ pe (WriteReturns _ arr idx_vals) = do
+  arr_t <- lookupType arr
   let (iss, vs) = unzip idx_vals
-      rws' = map pe64 rws
+      rws' = map pe64 $ arrayDims arr_t
   forM_ (zip iss vs) $ \(slice, v) -> do
     let slice' = fmap pe64 slice
     sWhen (inBounds slice' rws') $
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
@@ -110,8 +110,11 @@
 
 instance PP.Pretty KernelGrid where
   pretty (KernelGrid num_groups group_size) =
-    "groups=" <> pretty num_groups <> PP.semi
-      <+> "groupsize=" <> pretty group_size
+    "groups="
+      <> pretty num_groups
+      <> PP.semi
+        <+> "groupsize="
+      <> pretty group_size
 
 instance PP.Pretty SegLevel where
   pretty (SegThread virt grid) =
@@ -219,9 +222,10 @@
   pretty (GetSizeMax size_class) =
     "get_size_max" <> parens (commasep [pretty size_class])
   pretty (CmpSizeLe name size_class x) =
-    "cmp_size" <> parens (commasep [pretty name, pretty size_class])
-      <+> "<="
-      <+> pretty x
+    "cmp_size"
+      <> parens (commasep [pretty name, pretty size_class])
+        <+> "<="
+        <+> pretty x
   pretty (CalcNumGroups w max_num_groups group_size) =
     "calc_num_groups" <> parens (commasep [pretty w, pretty max_num_groups, pretty group_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
@@ -833,10 +833,8 @@
         <*> pure cs
         <*> pSubExp,
       try $
-        flip (SegOp.WriteReturns cs)
+        SegOp.WriteReturns cs
           <$> pVName
-          <* pColon
-          <*> pShape
           <* keyword "with"
           <*> parens (pWrite `sepBy` pComma),
       try "tile"
diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs
--- a/src/Futhark/IR/Pretty.hs
+++ b/src/Futhark/IR/Pretty.hs
@@ -196,7 +196,8 @@
       Array {} -> brackets $ commastack $ map pretty es
       _ -> brackets $ commasep $ map pretty es
       <+> colon
-      <+> "[]" <> pretty rt
+      <+> "[]"
+      <> pretty rt
   pretty (BinOp bop x y) = pretty bop <> parens (pretty x <> comma <+> pretty y)
   pretty (CmpOp op x y) = pretty op <> parens (pretty x <> comma <+> pretty y)
   pretty (ConvOp conv x) =
@@ -270,13 +271,13 @@
   pretty (Match [c] [Case [Just (BoolValue True)] t] f (MatchDec ret ifsort)) =
     "if"
       <> info'
-      <+> pretty c
-      </> "then"
-      <+> maybeNest t
-      <+> "else"
-      <+> maybeNest f
-      </> colon
-      <+> ppTupleLines' (map pretty ret)
+        <+> pretty c
+        </> "then"
+        <+> maybeNest t
+        <+> "else"
+        <+> maybeNest f
+        </> colon
+        <+> ppTupleLines' (map pretty ret)
     where
       info' = case ifsort of
         MatchNormal -> mempty
@@ -300,8 +301,8 @@
     applykw
       <+> pretty (nameToString fname)
       <> apply (map (align . prettyArg) args)
-      </> colon
-      <+> braces (commasep $ map prettyRet ret)
+        </> colon
+        <+> braces (commasep $ map prettyRet ret)
     where
       prettyArg (arg, Consume) = "*" <> pretty arg
       prettyArg (arg, _) = pretty arg
@@ -318,9 +319,11 @@
               ForLoop i it bound ->
                 "for"
                   <+> align
-                    ( pretty i <> ":" <> pretty it
-                        <+> "<"
-                        <+> align (pretty bound)
+                    ( pretty i
+                        <> ":"
+                        <> pretty it
+                          <+> "<"
+                          <+> align (pretty bound)
                     )
               WhileLoop cond ->
                 "while" <+> pretty cond
@@ -335,12 +338,13 @@
     where
       ppInput (shape, arrs, op) =
         parens
-          ( pretty shape <> comma
-              <+> ppTuple' (map pretty arrs)
-                <> case op of
-                  Nothing -> mempty
-                  Just (op', nes) ->
-                    comma </> parens (pretty op' <> comma </> ppTuple' (map pretty nes))
+          ( pretty shape
+              <> comma
+                <+> ppTuple' (map pretty arrs)
+              <> case op of
+                Nothing -> mempty
+                Just (op', nes) ->
+                  comma </> parens (pretty op' <> comma </> ppTuple' (map pretty nes))
           )
 
 instance (PrettyRep rep) => Pretty (Lambda rep) where
@@ -388,9 +392,9 @@
                   <> pretty p_name
                   <> "\""
                   <> comma
-                  </> ppTupleLines' (map pretty p_entry)
+                    </> ppTupleLines' (map pretty p_entry)
                   <> comma
-                  </> ppTupleLines' (map pretty ret_entry)
+                    </> ppTupleLines' (map pretty ret_entry)
               )
 
 instance Pretty OpaqueType where
diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs
--- a/src/Futhark/IR/Prop.hs
+++ b/src/Futhark/IR/Prop.hs
@@ -10,7 +10,7 @@
     module Futhark.IR.Prop.Types,
     module Futhark.IR.Prop.Constants,
     module Futhark.IR.Prop.TypeOf,
-    module Futhark.IR.Prop.Patterns,
+    module Futhark.IR.Prop.Pat,
     module Futhark.IR.Prop.Names,
     module Futhark.IR.RetType,
     module Futhark.IR.Rephrase,
@@ -45,7 +45,7 @@
 import Futhark.IR.Pretty
 import Futhark.IR.Prop.Constants
 import Futhark.IR.Prop.Names
-import Futhark.IR.Prop.Patterns
+import Futhark.IR.Prop.Pat
 import Futhark.IR.Prop.Rearrange
 import Futhark.IR.Prop.Reshape
 import Futhark.IR.Prop.TypeOf
diff --git a/src/Futhark/IR/Prop/Aliases.hs b/src/Futhark/IR/Prop/Aliases.hs
--- a/src/Futhark/IR/Prop/Aliases.hs
+++ b/src/Futhark/IR/Prop/Aliases.hs
@@ -33,7 +33,7 @@
 import Data.Map qualified as M
 import Futhark.IR.Prop (ASTRep, IsOp, NameInfo (..), Scope)
 import Futhark.IR.Prop.Names
-import Futhark.IR.Prop.Patterns
+import Futhark.IR.Prop.Pat
 import Futhark.IR.Prop.Types
 import Futhark.IR.Syntax
 
diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs
--- a/src/Futhark/IR/Prop/Names.hs
+++ b/src/Futhark/IR/Prop/Names.hs
@@ -48,7 +48,7 @@
 import Data.IntSet qualified as IS
 import Data.Map.Strict qualified as M
 import Data.Set qualified as S
-import Futhark.IR.Prop.Patterns
+import Futhark.IR.Prop.Pat
 import Futhark.IR.Syntax
 import Futhark.IR.Traversals
 import Futhark.Util.Pretty
diff --git a/src/Futhark/IR/Prop/Pat.hs b/src/Futhark/IR/Prop/Pat.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/IR/Prop/Pat.hs
@@ -0,0 +1,73 @@
+{-# LANGUAGE TypeFamilies #-}
+
+-- | Inspecing and modifying t'Pat's, function parameters and
+-- pattern elements.
+module Futhark.IR.Prop.Pat
+  ( -- * Function parameters
+    paramIdent,
+    paramType,
+    paramDeclType,
+
+    -- * Pat elements
+    patElemIdent,
+    patElemType,
+    setPatElemDec,
+    patIdents,
+    patNames,
+    patTypes,
+    patSize,
+
+    -- * Pat construction
+    basicPat,
+  )
+where
+
+import Futhark.IR.Prop.Types (DeclTyped (..), Typed (..))
+import Futhark.IR.Syntax
+
+-- | The 'Type' of a parameter.
+paramType :: (Typed dec) => Param dec -> Type
+paramType = typeOf
+
+-- | The 'DeclType' of a parameter.
+paramDeclType :: (DeclTyped dec) => Param dec -> DeclType
+paramDeclType = declTypeOf
+
+-- | An 'Ident' corresponding to a parameter.
+paramIdent :: (Typed dec) => Param dec -> Ident
+paramIdent param = Ident (paramName param) (typeOf param)
+
+-- | An 'Ident' corresponding to a pattern element.
+patElemIdent :: (Typed dec) => PatElem dec -> Ident
+patElemIdent pelem = Ident (patElemName pelem) (typeOf pelem)
+
+-- | The type of a name bound by a t'PatElem'.
+patElemType :: (Typed dec) => PatElem dec -> Type
+patElemType = typeOf
+
+-- | Set the rep of a t'PatElem'.
+setPatElemDec :: PatElem oldattr -> newattr -> PatElem newattr
+setPatElemDec pe x = fmap (const x) pe
+
+-- | Return a list of the 'Ident's bound by the t'Pat'.
+patIdents :: (Typed dec) => Pat dec -> [Ident]
+patIdents = map patElemIdent . patElems
+
+-- | Return a list of the 'Name's bound by the t'Pat'.
+patNames :: Pat dec -> [VName]
+patNames = map patElemName . patElems
+
+-- | Return a list of the typess bound by the pattern.
+patTypes :: (Typed dec) => Pat dec -> [Type]
+patTypes = map identType . patIdents
+
+-- | Return the number of names bound by the pattern.
+patSize :: Pat dec -> Int
+patSize (Pat xs) = length xs
+
+-- | Create a pattern using 'Type' as the attribute.
+basicPat :: [Ident] -> Pat Type
+basicPat values =
+  Pat $ map patElem values
+  where
+    patElem (Ident name t) = PatElem name t
diff --git a/src/Futhark/IR/Prop/Patterns.hs b/src/Futhark/IR/Prop/Patterns.hs
deleted file mode 100644
--- a/src/Futhark/IR/Prop/Patterns.hs
+++ /dev/null
@@ -1,73 +0,0 @@
-{-# LANGUAGE TypeFamilies #-}
-
--- | Inspecing and modifying t'Pat's, function parameters and
--- pattern elements.
-module Futhark.IR.Prop.Patterns
-  ( -- * Function parameters
-    paramIdent,
-    paramType,
-    paramDeclType,
-
-    -- * Pat elements
-    patElemIdent,
-    patElemType,
-    setPatElemDec,
-    patIdents,
-    patNames,
-    patTypes,
-    patSize,
-
-    -- * Pat construction
-    basicPat,
-  )
-where
-
-import Futhark.IR.Prop.Types (DeclTyped (..), Typed (..))
-import Futhark.IR.Syntax
-
--- | The 'Type' of a parameter.
-paramType :: (Typed dec) => Param dec -> Type
-paramType = typeOf
-
--- | The 'DeclType' of a parameter.
-paramDeclType :: (DeclTyped dec) => Param dec -> DeclType
-paramDeclType = declTypeOf
-
--- | An 'Ident' corresponding to a parameter.
-paramIdent :: (Typed dec) => Param dec -> Ident
-paramIdent param = Ident (paramName param) (typeOf param)
-
--- | An 'Ident' corresponding to a pattern element.
-patElemIdent :: (Typed dec) => PatElem dec -> Ident
-patElemIdent pelem = Ident (patElemName pelem) (typeOf pelem)
-
--- | The type of a name bound by a t'PatElem'.
-patElemType :: (Typed dec) => PatElem dec -> Type
-patElemType = typeOf
-
--- | Set the rep of a t'PatElem'.
-setPatElemDec :: PatElem oldattr -> newattr -> PatElem newattr
-setPatElemDec pe x = fmap (const x) pe
-
--- | Return a list of the 'Ident's bound by the t'Pat'.
-patIdents :: (Typed dec) => Pat dec -> [Ident]
-patIdents = map patElemIdent . patElems
-
--- | Return a list of the 'Name's bound by the t'Pat'.
-patNames :: Pat dec -> [VName]
-patNames = map patElemName . patElems
-
--- | Return a list of the typess bound by the pattern.
-patTypes :: (Typed dec) => Pat dec -> [Type]
-patTypes = map identType . patIdents
-
--- | Return the number of names bound by the pattern.
-patSize :: Pat dec -> Int
-patSize (Pat xs) = length xs
-
--- | Create a pattern using 'Type' as the attribute.
-basicPat :: [Ident] -> Pat Type
-basicPat values =
-  Pat $ map patElem values
-  where
-    patElem (Ident name t) = PatElem name t
diff --git a/src/Futhark/IR/Prop/Types.hs b/src/Futhark/IR/Prop/Types.hs
--- a/src/Futhark/IR/Prop/Types.hs
+++ b/src/Futhark/IR/Prop/Types.hs
@@ -103,7 +103,7 @@
   TypeBase newshape u
 modifyArrayShape f (Array t ds u)
   | shapeRank ds' == 0 = Prim t
-  | otherwise = Array t (f ds) u
+  | otherwise = Array t ds' u
   where
     ds' = f ds
 modifyArrayShape _ (Prim t) = Prim t
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
@@ -937,9 +937,9 @@
         ( PP.align $
             pretty lam
               <> comma
-              </> PP.braces (commasep $ map pretty args)
+                </> PP.braces (commasep $ map pretty args)
               <> comma
-              </> PP.braces (commasep $ map pretty vec)
+                </> PP.braces (commasep $ map pretty vec)
         )
   pretty (JVP lam args vec) =
     "jvp"
@@ -947,9 +947,9 @@
         ( PP.align $
             pretty lam
               <> comma
-              </> PP.braces (commasep $ map pretty args)
+                </> PP.braces (commasep $ map pretty args)
               <> comma
-              </> PP.braces (commasep $ map pretty vec)
+                </> PP.braces (commasep $ map pretty vec)
         )
   pretty (Stream size arrs acc lam) =
     ppStream size arrs acc lam
@@ -964,31 +964,31 @@
           <> (parens . align)
             ( pretty w
                 <> comma
-                </> ppTuple' (map pretty arrs)
+                  </> ppTuple' (map pretty arrs)
                 <> comma
-                </> pretty map_lam
+                  </> pretty map_lam
             )
     | null scans =
         "redomap"
           <> (parens . align)
             ( pretty w
                 <> comma
-                </> ppTuple' (map pretty arrs)
+                  </> ppTuple' (map pretty arrs)
                 <> comma
-                </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)
+                  </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)
                 <> comma
-                </> pretty map_lam
+                  </> pretty map_lam
             )
     | null reds =
         "scanomap"
           <> (parens . align)
             ( pretty w
                 <> comma
-                </> ppTuple' (map pretty arrs)
+                  </> ppTuple' (map pretty arrs)
                 <> comma
-                </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)
+                  </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)
                 <> comma
-                </> pretty map_lam
+                  </> pretty map_lam
             )
   pretty (Screma w arrs form) = ppScrema w arrs form
 
@@ -1000,13 +1000,13 @@
     <> (parens . align)
       ( pretty w
           <> comma
-          </> ppTuple' (map pretty arrs)
+            </> ppTuple' (map pretty arrs)
           <> comma
-          </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)
+            </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)
           <> comma
-          </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)
+            </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)
           <> comma
-          </> pretty map_lam
+            </> pretty map_lam
       )
 
 -- | Prettyprint the given Stream.
@@ -1017,11 +1017,11 @@
     <> (parens . align)
       ( pretty size
           <> comma
-          </> ppTuple' (map pretty arrs)
+            </> ppTuple' (map pretty arrs)
           <> comma
-          </> ppTuple' (map pretty acc)
+            </> ppTuple' (map pretty acc)
           <> comma
-          </> pretty lam
+            </> pretty lam
       )
 
 -- | Prettyprint the given Scatter.
@@ -1032,11 +1032,11 @@
     <> (parens . align)
       ( pretty w
           <> comma
-          </> ppTuple' (map pretty arrs)
+            </> ppTuple' (map pretty arrs)
           <> comma
-          </> pretty lam
+            </> pretty lam
           <> comma
-          </> commasep (map pretty dests)
+            </> commasep (map pretty dests)
       )
 
 instance (PrettyRep rep) => Pretty (Scan rep) where
@@ -1052,7 +1052,7 @@
     ppComm comm
       <> pretty red_lam
       <> comma
-      </> PP.braces (commasep $ map pretty red_nes)
+        </> PP.braces (commasep $ map pretty red_nes)
 
 -- | Prettyprint the given histogram operation.
 ppHist ::
@@ -1067,20 +1067,20 @@
     <> parens
       ( pretty w
           <> comma
-          </> ppTuple' (map pretty arrs)
+            </> ppTuple' (map pretty arrs)
           <> comma
-          </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops)
+            </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops)
           <> comma
-          </> pretty bucket_fun
+            </> pretty bucket_fun
       )
   where
     ppOp (HistOp dest_w rf dests nes op) =
       pretty dest_w
         <> comma
-        <+> pretty rf
+          <+> pretty rf
         <> comma
-        <+> PP.braces (commasep $ map pretty dests)
+          <+> PP.braces (commasep $ map pretty dests)
         <> comma
-        </> ppTuple' (map pretty nes)
+          </> ppTuple' (map pretty nes)
         <> comma
-        </> pretty op
+          </> pretty op
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
@@ -53,7 +53,6 @@
 import Control.Monad.Writer
 import Data.Bifunctor (first)
 import Data.Bitraversable
-import Data.Foldable (traverse_)
 import Data.List
   ( elemIndex,
     foldl',
@@ -194,8 +193,7 @@
     Returns ResultManifest Certs SubExp
   | WriteReturns
       Certs
-      Shape -- Size of array.  Must match number of dims.
-      VName -- Which array
+      VName -- Destination array
       [(Slice SubExp, SubExp)]
   | TileReturns
       Certs
@@ -217,20 +215,20 @@
 -- | Get the certs for this 'KernelResult'.
 kernelResultCerts :: KernelResult -> Certs
 kernelResultCerts (Returns _ cs _) = cs
-kernelResultCerts (WriteReturns cs _ _ _) = cs
+kernelResultCerts (WriteReturns cs _ _) = cs
 kernelResultCerts (TileReturns cs _ _) = cs
 kernelResultCerts (RegTileReturns cs _ _) = cs
 
 -- | Get the root t'SubExp' corresponding values for a 'KernelResult'.
 kernelResultSubExp :: KernelResult -> SubExp
 kernelResultSubExp (Returns _ _ se) = se
-kernelResultSubExp (WriteReturns _ _ arr _) = Var arr
+kernelResultSubExp (WriteReturns _ arr _) = Var arr
 kernelResultSubExp (TileReturns _ _ v) = Var v
 kernelResultSubExp (RegTileReturns _ _ v) = Var v
 
 instance FreeIn KernelResult where
   freeIn' (Returns _ cs what) = freeIn' cs <> freeIn' what
-  freeIn' (WriteReturns cs rws arr res) = freeIn' cs <> freeIn' rws <> freeIn' arr <> freeIn' res
+  freeIn' (WriteReturns cs arr res) = freeIn' cs <> freeIn' arr <> freeIn' res
   freeIn' (TileReturns cs dims v) =
     freeIn' cs <> freeIn' dims <> freeIn' v
   freeIn' (RegTileReturns cs dims_n_tiles v) =
@@ -252,10 +250,9 @@
 instance Substitute KernelResult where
   substituteNames subst (Returns manifest cs se) =
     Returns manifest (substituteNames subst cs) (substituteNames subst se)
-  substituteNames subst (WriteReturns cs rws arr res) =
+  substituteNames subst (WriteReturns cs arr res) =
     WriteReturns
       (substituteNames subst cs)
-      (substituteNames subst rws)
       (substituteNames subst arr)
       (substituteNames subst res)
   substituteNames subst (TileReturns cs dims v) =
@@ -296,7 +293,7 @@
 consumedInKernelBody (KernelBody dec stms res) =
   consumedInBody (Body dec stms []) <> mconcat (map consumedByReturn res)
   where
-    consumedByReturn (WriteReturns _ _ a _) = oneName a
+    consumedByReturn (WriteReturns _ a _) = oneName a
     consumedByReturn _ = mempty
 
 checkKernelBody ::
@@ -320,7 +317,7 @@
           <> " values."
     zipWithM_ checkKernelResult kres ts
   where
-    consumeKernelResult (WriteReturns _ _ arr _) =
+    consumeKernelResult (WriteReturns _ arr _) =
       TC.consume =<< TC.lookupAliases arr
     consumeKernelResult _ =
       pure ()
@@ -328,24 +325,20 @@
     checkKernelResult (Returns _ cs what) t = do
       TC.checkCerts cs
       TC.require [t] what
-    checkKernelResult (WriteReturns cs shape arr res) t = do
+    checkKernelResult (WriteReturns cs arr res) t = do
       TC.checkCerts cs
-      mapM_ (TC.require [Prim int64]) $ shapeDims shape
       arr_t <- lookupType arr
+      unless (arr_t == t) $
+        TC.bad . TC.TypeError $
+          "WriteReturns result type annotation for "
+            <> prettyText arr
+            <> " is "
+            <> prettyText t
+            <> ", but inferred as"
+            <> prettyText arr_t
       forM_ res $ \(slice, e) -> do
-        traverse_ (TC.require [Prim int64]) slice
-        TC.require [t] e
-        unless (arr_t == t `arrayOfShape` shape) $
-          TC.bad $
-            TC.TypeError $
-              "WriteReturns returning "
-                <> prettyText e
-                <> " of type "
-                <> prettyText t
-                <> ", shape="
-                <> prettyText shape
-                <> ", but destination array has type "
-                <> prettyText arr_t
+        TC.checkSlice arr_t slice
+        TC.require [t `setArrayShape` sliceShape slice] e
     checkKernelResult (TileReturns cs dims v) t = do
       TC.checkCerts cs
       forM_ dims $ \(dim, tile) -> do
@@ -395,15 +388,10 @@
     hsep $ certAnnots cs <> ["returns (private)" <+> pretty what]
   pretty (Returns ResultMaySimplify cs what) =
     hsep $ certAnnots cs <> ["returns" <+> pretty what]
-  pretty (WriteReturns cs shape arr res) =
+  pretty (WriteReturns cs arr res) =
     hsep $
       certAnnots cs
-        <> [ pretty arr
-               <+> PP.colon
-               <+> pretty shape
-               </> "with"
-               <+> PP.apply (map ppRes res)
-           ]
+        <> [pretty arr </> "with" <+> PP.apply (map ppRes res)]
     where
       ppRes (slice, e) = pretty slice <+> "=" <+> pretty e
   pretty (TileReturns cs dims v) =
@@ -450,6 +438,12 @@
 -- of information.  For example, in GPU backends, it is used to
 -- indicate whether the 'SegOp' is expected to run at the thread-level
 -- or the group-level.
+--
+-- The type list is usually the type of the element returned by a
+-- single thread. The result of the SegOp is then an array of that
+-- type, with the shape of the 'SegSpace' prepended. One exception is
+-- for 'WriteReturns', where the type annotation is the /full/ type of
+-- the result.
 data SegOp lvl rep
   = SegMap lvl SegSpace [Type] (KernelBody rep)
   | -- | The KernelSpace must always have at least two dimensions,
@@ -483,8 +477,8 @@
     SegHist _ _ _ _ body -> body
 
 segResultShape :: SegSpace -> Type -> KernelResult -> Type
-segResultShape _ t (WriteReturns _ shape _ _) =
-  t `arrayOfShape` shape
+segResultShape _ t (WriteReturns {}) =
+  t
 segResultShape space t Returns {} =
   foldr (flip arrayOfRow) t $ segSpaceDims space
 segResultShape _ t (TileReturns _ dims _) =
@@ -874,9 +868,9 @@
   pretty (SegBinOp comm lam nes shape) =
     PP.braces (PP.commasep $ map pretty nes)
       <> PP.comma
-      </> pretty shape
+        </> pretty shape
       <> PP.comma
-      </> comm'
+        </> comm'
       <> pretty lam
     where
       comm' = case comm of
@@ -887,47 +881,47 @@
   pretty (SegMap lvl space ts body) =
     "segmap"
       <> pretty lvl
-      </> PP.align (pretty space)
-      <+> PP.colon
-      <+> ppTuple' (map pretty ts)
-      <+> PP.nestedBlock "{" "}" (pretty body)
+        </> PP.align (pretty space)
+        <+> PP.colon
+        <+> ppTuple' (map pretty ts)
+        <+> PP.nestedBlock "{" "}" (pretty body)
   pretty (SegRed lvl space reds ts body) =
     "segred"
       <> pretty lvl
-      </> PP.align (pretty space)
-      </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty reds)
-      </> PP.colon
-      <+> ppTuple' (map pretty ts)
-      <+> PP.nestedBlock "{" "}" (pretty body)
+        </> PP.align (pretty space)
+        </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty reds)
+        </> PP.colon
+        <+> ppTuple' (map pretty ts)
+        <+> PP.nestedBlock "{" "}" (pretty body)
   pretty (SegScan lvl space scans ts body) =
     "segscan"
       <> pretty lvl
-      </> PP.align (pretty space)
-      </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty scans)
-      </> PP.colon
-      <+> ppTuple' (map pretty ts)
-      <+> PP.nestedBlock "{" "}" (pretty body)
+        </> PP.align (pretty space)
+        </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty scans)
+        </> PP.colon
+        <+> ppTuple' (map pretty ts)
+        <+> PP.nestedBlock "{" "}" (pretty body)
   pretty (SegHist lvl space ops ts body) =
     "seghist"
       <> pretty lvl
-      </> PP.align (pretty space)
-      </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)
-      </> PP.colon
-      <+> ppTuple' (map pretty ts)
-      <+> PP.nestedBlock "{" "}" (pretty body)
+        </> PP.align (pretty space)
+        </> PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)
+        </> PP.colon
+        <+> ppTuple' (map pretty ts)
+        <+> PP.nestedBlock "{" "}" (pretty body)
     where
       ppOp (HistOp w rf dests nes shape op) =
         pretty w
           <> PP.comma
-          <+> pretty rf
+            <+> pretty rf
           <> PP.comma
-          </> PP.braces (PP.commasep $ map pretty dests)
+            </> PP.braces (PP.commasep $ map pretty dests)
           <> PP.comma
-          </> PP.braces (PP.commasep $ map pretty nes)
+            </> PP.braces (PP.commasep $ map pretty nes)
           <> PP.comma
-          </> pretty shape
+            </> pretty shape
           <> PP.comma
-          </> pretty op
+            </> pretty op
 
 instance CanBeAliased (SegOp lvl) where
   addOpAliases aliases = runIdentity . mapSegOpM alias
@@ -1017,10 +1011,9 @@
 instance Engine.Simplifiable KernelResult where
   simplify (Returns manifest cs what) =
     Returns manifest <$> Engine.simplify cs <*> Engine.simplify what
-  simplify (WriteReturns cs ws a res) =
+  simplify (WriteReturns cs a res) =
     WriteReturns
       <$> Engine.simplify cs
-      <*> Engine.simplify ws
       <*> Engine.simplify a
       <*> Engine.simplify res
   simplify (TileReturns cs dims what) =
@@ -1088,7 +1081,7 @@
     scope_vtable = segSpaceSymbolTable space
     bound_here = namesFromList $ M.keys $ scopeOfSegSpace space
 
-    consumedInResult (WriteReturns _ _ arr _) =
+    consumedInResult (WriteReturns _ arr _) =
       [arr]
     consumedInResult _ =
       []
@@ -1428,7 +1421,7 @@
   m [ExpReturns]
 kernelBodyReturns = zipWithM correct . kernelBodyResult
   where
-    correct (WriteReturns _ _ arr _) _ = varReturns arr
+    correct (WriteReturns _ arr _) _ = varReturns arr
     correct _ ret = pure ret
 
 -- | Like 'segOpType', but for memory representations.
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
@@ -455,12 +455,12 @@
 -- | For-loop or while-loop?
 data LoopForm
   = ForLoop
+      -- | The loop iterator var
       VName
-      -- ^ The loop iterator var
+      -- | The type of the loop iterator var
       IntType
-      -- ^ The type of the loop iterator var
+      -- | The number of iterations.
       SubExp
-      -- ^ The number of iterations.
   | WhileLoop VName
   deriving (Eq, Ord, Show)
 
diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs
--- a/src/Futhark/IR/Syntax/Core.hs
+++ b/src/Futhark/IR/Syntax/Core.hs
@@ -425,10 +425,10 @@
 -- | A dimension in a 'FlatSlice'.
 data FlatDimIndex d
   = FlatDimIndex
+      -- | Number of elements in dimension
       d
-      -- ^ Number of elements in dimension
+      -- | Stride of dimension
       d
-      -- ^ Stride of dimension
   deriving (Eq, Ord, Show)
 
 instance Traversable FlatDimIndex where
diff --git a/src/Futhark/IR/TypeCheck.hs b/src/Futhark/IR/TypeCheck.hs
--- a/src/Futhark/IR/TypeCheck.hs
+++ b/src/Futhark/IR/TypeCheck.hs
@@ -26,6 +26,7 @@
     checkExp,
     checkStms,
     checkStm,
+    checkSlice,
     checkType,
     checkExtType,
     matchExtPat,
@@ -517,9 +518,7 @@
 require :: (Checkable rep) => [Type] -> SubExp -> TypeM rep ()
 require ts se = do
   t <- checkSubExp se
-  unless (t `elem` ts) $
-    bad $
-      UnexpectedType (BasicOp $ SubExp se) t ts
+  unless (t `elem` ts) $ bad $ UnexpectedType (BasicOp $ SubExp se) t ts
 
 -- | Variant of 'require' working on variable names.
 requireI :: (Checkable rep) => [Type] -> VName -> TypeM rep ()
@@ -562,7 +561,7 @@
     check _ (OpaqueType _) =
       pure ()
     checkEntryPointType known (TypeOpaque s) =
-      when (s `notElem` known) $
+      unless (s `elem` known) $
         Left . Error [] . TypeError $
           "Opaque not defined before first use: " <> nameToText s
     checkEntryPointType _ (TypeTransparent _) = pure ()
@@ -824,6 +823,13 @@
   where
     bound_here = namesFromList $ M.keys $ scopeOf stms
 
+-- | Check a slicing operation of an array of the provided type.
+checkSlice :: (Checkable rep) => Type -> Slice SubExp -> TypeM rep ()
+checkSlice vt (Slice idxes) = do
+  when (arrayRank vt /= length idxes) . bad $
+    SlicingError (arrayRank vt) (length idxes)
+  mapM_ (traverse $ require [Prim int64]) idxes
+
 checkBasicOp :: (Checkable rep) => BasicOp -> TypeM rep ()
 checkBasicOp (SubExp es) =
   void $ checkSubExp es
@@ -849,26 +855,20 @@
 checkBasicOp (BinOp op e1 e2) = checkBinOpArgs (binOpType op) e1 e2
 checkBasicOp (CmpOp op e1 e2) = checkCmpOp op e1 e2
 checkBasicOp (ConvOp op e) = require [Prim $ fst $ convOpType op] e
-checkBasicOp (Index ident (Slice idxes)) = do
+checkBasicOp (Index ident slice) = do
   vt <- lookupType ident
   observe ident
-  when (arrayRank vt /= length idxes) $
-    bad $
-      SlicingError (arrayRank vt) (length idxes)
-  mapM_ checkDimIndex idxes
-checkBasicOp (Update _ src (Slice idxes) se) = do
+  checkSlice vt slice
+checkBasicOp (Update _ src slice se) = do
   (src_shape, src_pt) <- checkArrIdent src
-  when (shapeRank src_shape /= length idxes) $
-    bad $
-      SlicingError (shapeRank src_shape) (length idxes)
 
   se_aliases <- subExpAliasesM se
   when (src `nameIn` se_aliases) $
     bad $
       TypeError "The target of an Update must not alias the value to be written."
 
-  mapM_ checkDimIndex idxes
-  require [arrayOf (Prim src_pt) (Shape (sliceDims (Slice idxes))) NoUniqueness] se
+  checkSlice (arrayOf (Prim src_pt) src_shape NoUniqueness) slice
+  require [arrayOf (Prim src_pt) (sliceShape slice) NoUniqueness] se
   consume =<< lookupAliases src
 checkBasicOp (FlatIndex ident slice) = do
   vt <- lookupType ident
@@ -1236,13 +1236,6 @@
 checkFlatSlice (FlatSlice offset idxs) = do
   require [Prim int64] offset
   mapM_ checkFlatDimIndex idxs
-
-checkDimIndex ::
-  (Checkable rep) =>
-  DimIndex SubExp ->
-  TypeM rep ()
-checkDimIndex (DimFix i) = require [Prim int64] i
-checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int64]) [i, n, s]
 
 checkStm ::
   (Checkable rep) =>
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
@@ -5,6 +5,7 @@
 import Control.Monad.Identity
 import Control.Monad.Reader
 import Control.Monad.State
+import Data.Bifoldable (bifoldMap)
 import Data.Bifunctor
 import Data.Bitraversable
 import Data.Foldable
@@ -783,6 +784,11 @@
   where
     defunc' = mapM defuncExp'
 
+envFromDimNames :: [VName] -> Env
+envFromDimNames = M.fromList . flip zip (repeat d)
+  where
+    d = Binding Nothing $ Dynamic $ Scalar $ Prim $ Signed Int64
+
 -- | Defunctionalize a let-bound function, while preserving parameters
 -- that have order 0 types (i.e., non-functional).
 defuncLet ::
@@ -975,12 +981,14 @@
       apply_e = mkApply f' [(d, argext, arg')] callret
   pure (apply_e, sv)
 --
-defuncApplyArg fname_s (_, sv) _ =
+defuncApplyArg fname_s (_, sv) ((_, arg), _) =
   error $
     "defuncApplyArg: cannot apply StaticVal\n"
       <> show sv
       <> "\nFunction name: "
       <> prettyString fname_s
+      <> "\nArgument: "
+      <> prettyString arg
 
 updateReturn :: AppRes -> Exp -> Exp
 updateReturn (AppRes ret1 ext1) (AppExp apply (Info (AppRes ret2 ext2))) =
@@ -1059,11 +1067,6 @@
   PatLit {} -> mempty
   PatConstr _ _ ps _ -> foldMap envFromPat ps
 
-envFromDimNames :: [VName] -> Env
-envFromDimNames = M.fromList . flip zip (repeat d)
-  where
-    d = Binding Nothing $ Dynamic $ Scalar $ Prim $ Signed Int64
-
 -- | Given a closure environment, construct a record pattern that
 -- binds the closed over variables.  Insert wildcard for any patterns
 -- that would otherwise clash with size parameters.
@@ -1140,8 +1143,11 @@
       then dim_env <> M.singleton vn (Binding Nothing $ Dynamic t)
       else dim_env <> M.singleton vn (Binding Nothing sv)
   where
-    dim_env =
-      M.fromList $ map (,i64) $ S.toList $ fvVars $ freeInType t
+    -- Extract all sizes that are potentially bound here. This is
+    -- different from all free variables (see #2040).
+    dim_env = bifoldMap onDim (const mempty) t
+    onDim (Var v _ _) = M.singleton (qualLeaf v) i64
+    onDim _ = mempty
     i64 = Binding Nothing $ Dynamic $ Scalar $ Prim $ Signed Int64
 matchPatSV (Wildcard _ _) _ = pure mempty
 matchPatSV (PatAscription pat _ _) sv = matchPatSV pat sv
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
@@ -1699,9 +1699,12 @@
       old_dim <- I.arraysSize 0 <$> mapM lookupType arrs
       dim_ok <-
         letSubExp "dim_ok" <=< toExp $
-          pe64 old_dim .==. pe64 n' * pe64 m'
-            .&&. pe64 n' .>=. 0
-            .&&. pe64 m' .>=. 0
+          pe64 old_dim .==. pe64 n'
+            * pe64 m'
+              .&&. pe64 n'
+              .>=. 0
+              .&&. pe64 m'
+              .>=. 0
       dim_ok_cert <-
         assert
           "dim_ok_cert"
diff --git a/src/Futhark/Internalise/FullNormalise.hs b/src/Futhark/Internalise/FullNormalise.hs
--- a/src/Futhark/Internalise/FullNormalise.hs
+++ b/src/Futhark/Internalise/FullNormalise.hs
@@ -345,15 +345,7 @@
       _ -> Info $ AppRes (typeOf e) []
 
     f body (PatBind sizes p expr) =
-      AppExp
-        ( LetPat
-            sizes
-            p
-            expr
-            body
-            mempty
-        )
-        appRes
+      AppExp (LetPat sizes p expr body mempty) appRes
     f body (FunBind vn infos) =
       AppExp (LetFun vn infos body mempty) appRes
 
diff --git a/src/Futhark/Internalise/LiftLambdas.hs b/src/Futhark/Internalise/LiftLambdas.hs
--- a/src/Futhark/Internalise/LiftLambdas.hs
+++ b/src/Futhark/Internalise/LiftLambdas.hs
@@ -169,9 +169,10 @@
         <*> pure case_loc
 transformExp (AppExp (Loop sizes pat args form body loc) appres) = do
   args' <- transformExp args
-  form' <- astMap transformSubExps form
-  body' <- bindingParams sizes [pat] $ bindingForm form' $ transformExp body
-  pure $ AppExp (Loop sizes pat args' form' body' loc) appres
+  bindingParams sizes [pat] $ do
+    form' <- astMap transformSubExps form
+    body' <- bindingForm form' $ transformExp body
+    pure $ AppExp (Loop sizes pat args' form' body' loc) appres
 transformExp e@(Var v _ _) =
   -- Note that function-typed variables can only occur in expressions,
   -- not in other places where VNames/QualNames can occur.
diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs
--- a/src/Futhark/Internalise/Monomorphise.hs
+++ b/src/Futhark/Internalise/Monomorphise.hs
@@ -425,8 +425,10 @@
 transformType :: TypeBase Size u -> MonoM (TypeBase Size u)
 transformType typ =
   case typ of
-    Scalar scalar -> Scalar <$> transformScalarSizes scalar
-    Array u shape scalar -> Array u <$> mapM onDim shape <*> transformScalarSizes scalar
+    Scalar scalar ->
+      Scalar <$> transformScalarSizes scalar
+    Array u shape scalar ->
+      Array u <$> mapM onDim shape <*> transformScalarSizes scalar
   where
     transformScalarSizes :: ScalarTypeBase Size u -> MonoM (ScalarTypeBase Size u)
     transformScalarSizes (Record fs) =
@@ -511,10 +513,11 @@
         let (bs_local, bs_prop) = Seq.partition ((== fname) . fst) bs
         pure (unfoldLetFuns (map snd $ toList bs_local) e', const bs_prop)
   | otherwise = do
-      body' <- scoping (S.fromList (foldMap patNames params)) $ transformExp body
-      ret' <- transformRetTypeSizes (S.fromList (foldMap patNames params)) ret
+      params' <- mapM transformPat params
+      body' <- scoping (S.fromList (foldMap patNames params')) $ transformExp body
+      ret' <- transformRetTypeSizes (S.fromList (foldMap patNames params')) ret
       AppExp
-        <$> ( LetFun fname (tparams, params, retdecl, Info ret', body')
+        <$> ( LetFun fname (tparams, params', retdecl, Info ret', body')
                 <$> scoping (S.singleton fname) (transformExp e)
                 <*> pure loc
             )
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
@@ -581,7 +581,7 @@
         & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)
         & Slice
     resultSlice ixf = IxFun.slice ixf $ fullSlice (IxFun.shape ixf) thread_slice
-makeSegMapCoals _ _ td_env _ _ x (_, _, WriteReturns _ _ return_name _) =
+makeSegMapCoals _ _ td_env _ _ x (_, _, WriteReturns _ return_name _) =
   case getScopeMemInfo return_name $ scope td_env of
     Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem
     Nothing -> error "Should not happen?"
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs
@@ -119,14 +119,14 @@
 -- | Coalesced Access Entry
 data Coalesced
   = Coalesced
+      -- | the kind of coalescing
       CoalescedKind
-      -- ^ the kind of coalescing
-      ArrayMemBound
-      -- ^ destination mem_block info @f_m_x[i]@ (must be ArrayMem)
+      -- | destination mem_block info @f_m_x[i]@ (must be ArrayMem)
       -- (Maybe IxFun) -- the inverse ixfun of a coalesced array, such that
       --                     --  ixfuns can be correctly constructed for aliases;
+      ArrayMemBound
+      -- | substitutions for free vars in index function
       FreeVarSubsts
-      -- ^ substitutions for free vars in index function
 
 data CoalsEntry = CoalsEntry
   { -- | destination memory block
@@ -222,26 +222,27 @@
 instance Pretty Coalesced where
   pretty (Coalesced knd mbd _) =
     "(Kind:"
-      <+> pretty knd <> ", membds:"
-      <+> pretty mbd -- <> ", subs:" <+> pretty subs
-        <> ")"
-      <+> "\n"
+      <+> pretty knd
+      <> ", membds:"
+        <+> pretty mbd -- <> ", subs:" <+> pretty subs
+      <> ")"
+        <+> "\n"
 
 instance Pretty CoalsEntry where
   pretty etry =
     "{"
       <+> "Dstmem:"
       <+> pretty (dstmem etry)
-        <> ", AliasMems:"
-      <+> pretty (alsmem etry)
-      <+> ", optdeps:"
-      <+> pretty (M.toList $ optdeps etry)
-      <+> ", memrefs:"
-      <+> pretty (memrefs etry)
-      <+> ", vartab:"
-      <+> pretty (M.toList $ vartab etry)
-      <+> "}"
-      <+> "\n"
+      <> ", AliasMems:"
+        <+> pretty (alsmem etry)
+        <+> ", optdeps:"
+        <+> pretty (M.toList $ optdeps etry)
+        <+> ", memrefs:"
+        <+> pretty (memrefs etry)
+        <+> ", vartab:"
+        <+> pretty (M.toList $ vartab etry)
+        <+> "}"
+        <+> "\n"
 
 -- | Compute the union of two 'CoalsEntry'. If two 'CoalsEntry' do not refer to
 -- the same destination memory and use the same index function, the first
diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs
--- a/src/Futhark/Optimise/BlkRegTiling.hs
+++ b/src/Futhark/Optimise/BlkRegTiling.hs
@@ -57,7 +57,6 @@
     SegLevel,
     [Int],
     (VName, SubExp, VName, SubExp, SubExp),
-    (SubExp, SubExp),
     (VName, VName),
     (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType),
     (Lambda GPU, Lambda GPU)
@@ -72,7 +71,6 @@
     segthd_lvl,
     var_dims,
     (gtid_x, width_B, gtid_y, height_A, common_dim),
-    (a_loc_sz, b_loc_sz),
     (iii, jjj),
     (load_A, inp_A, pt_A, load_B, inp_B, pt_B),
     (map_lam, red_lam)
@@ -84,11 +82,11 @@
     kk <- letExp "kk" =<< toExp (le64 kk0 * pe64 tk)
     -- copy A to local memory
     (a_loc, aCopyLoc2Reg) <-
-      copyGlb2ShMem kk (gtid_y, iii, map_t1, height_A, inp_A, load_A, a_loc_sz, a_loc_init')
+      copyGlb2ShMem kk (gtid_y, iii, map_t1, height_A, inp_A, load_A, a_loc_init')
 
     -- copy B from global to shared memory
     (b_loc, bCopyLoc2Reg) <-
-      copyGlb2ShMem kk (gtid_x, jjj, map_t2, width_B, inp_B, load_B, b_loc_sz, b_loc_init')
+      copyGlb2ShMem kk (gtid_x, jjj, map_t2, width_B, inp_B, load_B, b_loc_init')
 
     -- inner loop updating this thread's accumulator (loop k in mmm_kernels).
     thd_acc <- forLoop tk [thd_res_merge] $ \k [acc_merge] ->
@@ -173,10 +171,16 @@
                           -- is garbage anyways and should not be written.
                           -- so fits_ij should be always true!!!
 
-                            le64 iii + le64 i + pe64 ry * le64 ltid_y
-                              .<. pe64 height_A
-                              .&&. le64 jjj + le64 j + pe64 rx * le64 ltid_x
-                                .<. pe64 width_B
+                            le64 iii
+                              + le64 i
+                              + pe64 ry
+                                * le64 ltid_y
+                                  .<. pe64 height_A
+                                  .&&. le64 jjj
+                              + le64 j
+                              + pe64 rx
+                                * le64 ltid_x
+                                  .<. pe64 width_B
                     )
                     ( do
                         a <- index "a" as [i]
@@ -204,14 +208,14 @@
       --
       copyGlb2ShMem ::
         VName ->
-        (VName, VName, PrimType, SubExp, VName, Stm GPU, SubExp, VName) ->
+        (VName, VName, PrimType, SubExp, VName, Stm GPU, VName) ->
         Builder GPU (VName, VName -> VName -> Builder GPU VName)
-      copyGlb2ShMem kk (gtid, ii, ptp_X_el, parlen_X, inp_X, load_X, loc_sz_X, x_loc_init') = do
+      copyGlb2ShMem kk (gtid, ii, ptp_X_el, parlen_X, inp_X, load_X, x_loc_init') = do
         let (t_par, r_par, tseq_div_tpar) = (tx, rx, tk_div_tx)
             is_inner_coal = isInnerCoal env inp_X load_X
             str_A = baseString inp_X
         x_loc <-
-          segScatter2D (str_A ++ "_glb2loc") loc_sz_X x_loc_init' [r_par, tseq_div_tpar] (t_par, t_par) $
+          segScatter2D (str_A ++ "_glb2loc") x_loc_init' [r_par, tseq_div_tpar] (t_par, t_par) $
             scatterFun is_inner_coal
 
         pure (x_loc, copyLoc2Reg is_inner_coal str_A x_loc)
@@ -255,7 +259,8 @@
               letSubExp (str_A ++ "_elem")
                 =<< eIf
                   ( toExp $
-                      le64 gtid .<. pe64 parlen_X
+                      le64 gtid
+                        .<. pe64 parlen_X
                         .&&. if epilogue
                           then le64 a_seqdim_idx .<. pe64 common_dim
                           else true
@@ -355,7 +360,6 @@
                   segthd_lvl,
                   var_dims,
                   (gtid_x, width_B, gtid_y, height_A, common_dim),
-                  (a_loc_sz, b_loc_sz),
                   (iii, jjj),
                   (load_A, inp_A, map_t1, load_B, inp_B, map_t2),
                   (map_lam, red_lam)
@@ -376,8 +380,10 @@
             letTupExp "redomap_res_if"
               =<< eIf
                 ( toExp $
-                    le64 full_tiles .==. pe64 rk
-                      .||. pe64 common_dim .==. (pe64 tk * le64 full_tiles + le64 ttt)
+                    le64 full_tiles
+                      .==. pe64 rk
+                      .||. pe64 common_dim
+                      .==. (pe64 tk * le64 full_tiles + le64 ttt)
                 )
                 (resultBodyM $ map Var prologue_res_list)
                 ( do
@@ -460,8 +466,10 @@
                 letSubExp "res_elem"
                   =<< eIf
                     ( toExp $
-                        le64 gtid_y .<. pe64 height_A
-                          .&&. le64 gtid_x .<. pe64 width_B
+                        le64 gtid_y
+                          .<. pe64 height_A
+                          .&&. le64 gtid_x
+                          .<. pe64 width_B
                     )
                     ( do
                         addStms code2_subs
@@ -539,7 +547,6 @@
                   segthd_lvl,
                   var_dims,
                   (gtid_x, width_B, gtid_y, height_A, common_dim),
-                  (a_loc_sz, b_loc_sz),
                   (iii, jjj),
                   (load_A, inp_A, map_t1, load_B, inp_B, map_t2),
                   (map_lam, red_lam)
@@ -608,8 +615,10 @@
                       letSubExp "res_elem"
                         =<< eIf
                           ( toExp $
-                              le64 gtid_y .<. pe64 height_A
-                                .&&. le64 gtid_x .<. pe64 width_B
+                              le64 gtid_y
+                                .<. pe64 height_A
+                                .&&. le64 gtid_x
+                                .<. pe64 width_B
                           )
                           ( do
                               addStms code2'
@@ -1153,14 +1162,14 @@
                           -- y_tp  <- subExpType y_elm
                           pure (y_elm, y_ind)
 
-                        let ret = WriteReturns mempty (Shape [rz]) loc_Y_nm [(Slice [DimFix res_i], res_v)]
+                        let ret = WriteReturns mempty loc_Y_nm [(Slice [DimFix res_i], res_v)]
                         let body = KernelBody () stms [ret]
+                        loc_Y_nm_t <- lookupType loc_Y_nm
 
                         res_nms <-
                           letTupExp "Y_glb2loc" <=< renameExp $
-                            Op $
-                              SegOp $
-                                SegMap segthd_lvl segspace [Prim ptp_Y] body
+                            Op . SegOp $
+                              SegMap segthd_lvl segspace [loc_Y_nm_t] body
                         let res_nm : _ = res_nms
                         pure res_nm
                     resultBodyM $ map Var loc_arr_merge2_nms'
@@ -1251,9 +1260,12 @@
                     letTupExp' "res_elem"
                       =<< eIf
                         ( toExp $
-                            le64 gtid_y .<. pe64 d_Ky
-                              .&&. le64 gtid_x .<. pe64 d_Kx
-                              .&&. le64 gtid_z .<. pe64 d_M
+                            le64 gtid_y
+                              .<. pe64 d_Ky
+                              .&&. le64 gtid_x
+                              .<. pe64 d_Kx
+                              .&&. le64 gtid_z
+                              .<. pe64 d_M
                         )
                         ( do
                             addStms code2'
diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs
--- a/src/Futhark/Optimise/CSE.hs
+++ b/src/Futhark/Optimise/CSE.hs
@@ -109,11 +109,7 @@
     f stms =
       fst $
         runReader
-          ( cseInStms
-              (consumedInStms stms)
-              (stmsToList stms)
-              (pure ())
-          )
+          (cseInStms (consumedInStms stms) (stmsToList stms) (pure ()))
           -- It is never safe to CSE arrays in stms in isolation,
           -- because we might introduce additional aliasing.
           (newCSEState False)
diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs
--- a/src/Futhark/Optimise/DoubleBuffer.hs
+++ b/src/Futhark/Optimise/DoubleBuffer.hs
@@ -89,20 +89,12 @@
 import Futhark.Transform.Substitute
 import Futhark.Util (mapAccumLM, maybeHead)
 
--- | The pass for GPU kernels.
-doubleBufferGPU :: Pass GPUMem GPUMem
-doubleBufferGPU = doubleBuffer optimiseGPUOp
-
--- | The pass for multicore
-doubleBufferMC :: Pass MCMem MCMem
-doubleBufferMC = doubleBuffer optimiseMCOp
-
 -- | The double buffering pass definition.
-doubleBuffer :: (Mem rep inner) => OptimiseOp rep -> Pass rep rep
-doubleBuffer onOp =
+doubleBuffer :: (Mem rep inner) => String -> String -> OptimiseOp rep -> Pass rep rep
+doubleBuffer name desc onOp =
   Pass
-    { passName = "Double buffer",
-      passDescription = "Perform double buffering for merge parameters of sequential loops.",
+    { passName = name,
+      passDescription = desc,
       passFunction = intraproceduralTransformation optimise
     }
   where
@@ -113,6 +105,22 @@
 
     env = Env mempty doNotTouchLoop onOp
     doNotTouchLoop pat merge body = pure (mempty, pat, merge, body)
+
+-- | The pass for GPU kernels.
+doubleBufferGPU :: Pass GPUMem GPUMem
+doubleBufferGPU =
+  doubleBuffer
+    "Double buffer GPU"
+    "Double buffer memory in sequential loops (GPU rep)."
+    optimiseGPUOp
+
+-- | The pass for multicore
+doubleBufferMC :: Pass MCMem MCMem
+doubleBufferMC =
+  doubleBuffer
+    "Double buffer MC"
+    "Double buffer memory in sequential loops (MC rep)."
+    optimiseMCOp
 
 type OptimiseLoop rep =
   Pat (LetDec rep) ->
diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
--- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
+++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
@@ -11,7 +11,7 @@
 import Control.Monad
 import Control.Monad.Writer
 import Data.Either
-import Data.List (find, unzip4)
+import Data.List (find, unzip5)
 import Data.Maybe (isNothing, mapMaybe)
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.Construct
@@ -84,12 +84,12 @@
   updates
     | all ((`elem` patNames pat) . updateValue) updates,
       not source_used_in_kbody = do
-        mk <- lowerUpdatesIntoSegMap scope pat updates space kbody
+        mk <- lowerUpdatesIntoSegMap scope pat ts updates space kbody
         Just $ do
-          (pat', kbody', poststms) <- mk
+          (pat', ts', kbody', poststms) <- mk
           let cs = stmAuxCerts aux <> foldMap updateCerts updates
           pure $
-            certify cs (Let pat' aux $ Op $ SegOp $ SegMap lvl space ts kbody')
+            certify cs (Let pat' aux $ Op $ SegOp $ SegMap lvl space ts' kbody')
               : stmsToList poststms
     where
       -- This check is a bit more conservative than ideal.  In a perfect
@@ -106,24 +106,27 @@
   (MonadFreshNames m) =>
   Scope (Aliases GPU) ->
   Pat (LetDec (Aliases GPU)) ->
+  [Type] ->
   [DesiredUpdate (LetDec (Aliases GPU))] ->
   SegSpace ->
   KernelBody (Aliases GPU) ->
   Maybe
     ( m
         ( Pat (LetDec (Aliases GPU)),
+          [Type],
           KernelBody (Aliases GPU),
           Stms (Aliases GPU)
         )
     )
-lowerUpdatesIntoSegMap scope pat updates kspace kbody = do
+lowerUpdatesIntoSegMap scope pat ret_ts updates kspace kbody = do
   -- The updates are all-or-nothing.  Being more liberal would require
   -- changes to the in-place-lowering pass itself.
-  mk <- zipWithM onRet (patElems pat) (kernelBodyResult kbody)
+  mk <- mapM onRet (zip3 (patElems pat) ret_ts (kernelBodyResult kbody))
   pure $ do
-    (pes, bodystms, krets, poststms) <- unzip4 <$> sequence mk
+    (pes, ret_ts', bodystms, krets, poststms) <- unzip5 <$> sequence mk
     pure
       ( Pat pes,
+        ret_ts',
         kbody
           { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms,
             kernelBodyResult = krets
@@ -133,7 +136,7 @@
   where
     (gtids, _dims) = unzip $ unSegSpace kspace
 
-    onRet (PatElem v v_dec) ret
+    onRet (PatElem v v_dec, _, ret)
       | Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <-
           find ((== v) . updateValue) updates = do
           Returns _ cs se <- Just ret
@@ -153,13 +156,13 @@
                   fixSlice (fmap pe64 slice) $
                     map (pe64 . Var) gtids
 
-            let res_dims = take (length slice') $ arrayDims $ snd bindee_dec
-                ret' = WriteReturns cs (Shape res_dims) src [(Slice $ map DimFix slice', se)]
+            let ret' = WriteReturns cs src [(fullSlice (typeOf bindee_dec) (map DimFix slice'), se)]
 
             v_aliased <- newName v
 
             pure
               ( PatElem bindee_nm bindee_dec,
+                typeOf bindee_dec,
                 bodystms,
                 ret',
                 stmsFromList
@@ -167,8 +170,8 @@
                     mkLet [Ident v $ typeOf v_dec] $ BasicOp $ Replicate mempty $ Var v_aliased
                   ]
               )
-    onRet pe ret =
-      Just $ pure (pe, mempty, ret, mempty)
+    onRet (pe, ret_t, ret) =
+      Just $ pure (pe, ret_t, mempty, ret, mempty)
 
 lowerUpdateIntoLoop ::
   ( Buildable rep,
diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs
--- a/src/Futhark/Optimise/TileLoops/Shared.hs
+++ b/src/Futhark/Optimise/TileLoops/Shared.hs
@@ -160,13 +160,12 @@
 
 segScatter2D ::
   String ->
-  SubExp ->
   VName ->
   [SubExp] -> -- dims of sequential loop on top
   (SubExp, SubExp) -> -- (dim_y, dim_x)
   ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)) -> -- f
   Builder GPU VName
-segScatter2D desc arr_size updt_arr seq_dims (dim_x, dim_y) f = do
+segScatter2D desc updt_arr seq_dims (dim_x, dim_y) f = do
   ltid_flat <- newVName "ltid_flat"
   ltid_y <- newVName "ltid_y"
   ltid_x <- newVName "ltid_x"
@@ -179,17 +178,15 @@
         SegThreadInGroup
           (SegNoVirtFull (SegSeqDims [0 .. length seq_dims - 1]))
 
-  ((t_v, res_v, res_i), stms) <- runBuilder $ do
-    (res_v, res_i) <-
-      localScope (scopeOfSegSpace segspace) $
-        f seq_is (ltid_y, ltid_x)
-    t_v <- subExpType res_v
-    pure (t_v, res_v, res_i)
+  ((res_v, res_i), stms) <-
+    runBuilder . localScope (scopeOfSegSpace segspace) $
+      f seq_is (ltid_y, ltid_x)
 
-  let ret = WriteReturns mempty (Shape [arr_size]) updt_arr [(Slice [DimFix res_i], res_v)]
+  let ret = WriteReturns mempty updt_arr [(Slice [DimFix res_i], res_v)]
   let body = KernelBody () stms [ret]
 
-  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [t_v] body
+  updt_arr_t <- lookupType updt_arr
+  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [updt_arr_t] body
 
 -- | The variance table keeps a mapping from a variable name
 -- (something produced by a 'Stm') to the kernel thread indices
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
@@ -448,19 +448,66 @@
   types <- asksScope scopeForSOACs
   transformStms path . stmsToList . snd
     =<< runBuilderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types
+--
+-- When we are scattering into a multidimensional array, we want to
+-- fully parallelise, such that we do not have threads writing
+-- potentially large rows. We do this by fissioning the scatter into a
+-- map part and a scatter part, where the former is flattened as
+-- usual, and the latter has a thread per primitive element to be
+-- written.
+--
+-- TODO: this could be slightly smarter. If we are dealing with a
+-- horizontally fused Scatter that targets both single- and
+-- multi-dimensional arrays, we could handle the former in the map
+-- stage. This would save us from having to store all the intermediate
+-- results to memory. Troels suspects such cases are very rare, but
+-- they may appear some day.
+transformStm path (Let pat aux (Op (Scatter w arrs lam as)))
+  | not $ all primType $ lambdaReturnType lam = do
+      -- Produce map stage.
+      map_pat <- fmap Pat $ forM (lambdaReturnType lam) $ \t ->
+        PatElem <$> newVName "scatter_tmp" <*> pure (t `arrayOfRow` w)
+      map_stms <- onMap path $ MapLoop map_pat aux w lam arrs
+
+      -- Now do the scatters.
+      runBuilder_ $ do
+        addStms map_stms
+        zipWithM_ doScatter (patElems pat) $ groupScatterResults as $ patNames map_pat
+  where
+    -- Generate code for a scatter where each thread writes only a scalar.
+    doScatter res_pe (scatter_space, arr, is_vs) = do
+      kernel_i <- newVName "write_i"
+      arr_t <- lookupType arr
+      val_t <- stripArray (shapeRank scatter_space) <$> lookupType arr
+      val_is <- replicateM (arrayRank val_t) (newVName "val_i")
+      (kret, kstms) <- collectStms $ do
+        is_vs' <- forM is_vs $ \(is, v) -> do
+          v' <- letSubExp (baseString v <> "_elem") $ BasicOp $ Index v $ Slice $ map (DimFix . Var) $ kernel_i : val_is
+          is' <- forM is $ \i' ->
+            letSubExp (baseString i' <> "_i") $ BasicOp $ Index i' $ Slice [DimFix $ Var kernel_i]
+          pure (Slice $ map DimFix $ is' <> map Var val_is, v')
+        pure $ WriteReturns mempty arr is_vs'
+      (kernel, stms) <-
+        mapKernel
+          segThreadCapped
+          ((kernel_i, w) : zip val_is (arrayDims val_t))
+          mempty
+          [arr_t]
+          (KernelBody () kstms [kret])
+      addStms stms
+      letBind (Pat [res_pe]) $ Op $ SegOp kernel
+--
 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
-      kstms = bodyStms $ lambdaBody lam'
-      krets = do
-        (a_w, a, is_vs) <- groupScatterResults as $ bodyResult $ lambdaBody lam'
+  let krets = do
+        (_a_w, a, is_vs) <- groupScatterResults as $ bodyResult $ lambdaBody lam'
         let res_cs =
               foldMap (foldMap resCerts . fst) is_vs
                 <> foldMap (resCerts . snd) is_vs
             is_vs' = [(Slice $ map (DimFix . resSubExp) is, resSubExp v) | (is, v) <- is_vs]
-        pure $ WriteReturns res_cs a_w a is_vs'
-      body = KernelBody () kstms krets
+        pure $ WriteReturns res_cs a is_vs'
+      body = KernelBody () (bodyStms $ lambdaBody lam') krets
       inputs = do
         (p, p_a) <- zip (lambdaParams lam') ivs
         pure $ KernelInput (paramName p) (paramType p) p_a [Var write_i]
@@ -469,7 +516,7 @@
       segThreadCapped
       [(write_i, w)]
       inputs
-      (zipWith (stripArray . length) as_ws $ patTypes pat)
+      (patTypes pat)
       body
   certifying cs $ do
     addStms stms
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
@@ -33,7 +33,6 @@
 import Control.Monad.Reader
 import Control.Monad.Trans.Maybe
 import Control.Monad.Writer.Strict
-import Data.Function ((&))
 import Data.List (find, partition, tails)
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Map qualified as M
@@ -54,7 +53,6 @@
 import Futhark.Transform.CopyPropagate
 import Futhark.Transform.FirstOrderTransform qualified as FOT
 import Futhark.Transform.Rename
-import Futhark.Util
 import Futhark.Util.Log
 
 scopeForSOACs :: (SameScope rep SOACS) => Scope rep -> Scope SOACS
@@ -788,27 +786,21 @@
 
   mk_lvl <- mkSegLevel
 
-  let rts =
-        concatMap (take 1) $
-          chunks as_ns $
-            drop (sum indexes) $
-              lambdaReturnType lam
-      (is, vs) = splitAt (sum indexes) $ bodyResult $ lambdaBody lam
-
+  let (is, vs) = splitAt (sum indexes) $ bodyResult $ lambdaBody lam
   (is', k_body_stms) <- runBuilder $ do
     addStms $ bodyStms $ lambdaBody lam
     pure is
 
-  let k_body =
-        groupScatterResults (zip3 as_ws as_ns as_inps) (is' ++ vs)
-          & map (inPlaceReturn ispace)
-          & KernelBody () k_body_stms
+  let grouped = groupScatterResults (zip3 as_ws as_ns as_inps) (is' ++ vs)
+      (_, dest_arrs, _) = unzip3 grouped
+      k_body = KernelBody () k_body_stms (map (inPlaceReturn ispace) grouped)
       -- Remove unused kernel inputs, since some of these might
       -- reference the array we are scattering into.
       kernel_inps' =
         filter ((`nameIn` freeIn k_body) . kernelInputName) kernel_inps
 
-  (k, k_stms) <- mapKernel mk_lvl ispace kernel_inps' rts k_body
+  dest_arrs_ts <- mapM (lookupType . kernelInputArray) dest_arrs
+  (k, k_stms) <- mapKernel mk_lvl ispace kernel_inps' dest_arrs_ts k_body
 
   traverse renameStm <=< runBuilder_ $ do
     addStms k_stms
@@ -825,18 +817,17 @@
       maybe bad pure $ find ((== a) . kernelInputName) kernel_inps
     bad = error "Ill-typed nested scatter encountered."
 
-    inPlaceReturn ispace (aw, inp, is_vs) =
+    inPlaceReturn ispace (_aw, inp, is_vs) =
       WriteReturns
         ( foldMap (foldMap resCerts . fst) is_vs
             <> foldMap (resCerts . snd) is_vs
         )
-        (Shape (init ws ++ shapeDims aw))
         (kernelInputArray inp)
         [ (Slice $ map DimFix $ map Var (init gtids) ++ map resSubExp is, resSubExp v)
           | (is, v) <- is_vs
         ]
       where
-        (gtids, ws) = unzip ispace
+        (gtids, _ws) = unzip ispace
 
 segmentedUpdateKernel ::
   (MonadFreshNames m, LocalScope rep m, DistRep rep) =>
@@ -854,7 +845,7 @@
 
   let ispace = base_ispace ++ zip slice_gtids slice_dims
 
-  ((res_t, res), kstms) <- runBuilder $ do
+  ((dest_t, res), kstms) <- runBuilder $ do
     -- Compute indexes into full array.
     v' <-
       certifying cs . letSubExp "v" . BasicOp . Index v $
@@ -869,10 +860,9 @@
           maybe (error "incorrectly typed Update") kernelInputArray $
             find ((== arr) . kernelInputName) kernel_inps
     arr_t <- lookupType arr'
-    v_t <- subExpType v'
     pure
-      ( v_t,
-        WriteReturns mempty (arrayShape arr_t) arr' [(Slice $ map DimFix write_is, v')]
+      ( arr_t,
+        WriteReturns mempty arr' [(Slice $ map DimFix write_is, v')]
       )
 
   -- Remove unused kernel inputs, since some of these might
@@ -882,18 +872,12 @@
 
   mk_lvl <- mkSegLevel
   (k, prestms) <-
-    mapKernel mk_lvl ispace kernel_inps' [res_t] $
+    mapKernel mk_lvl ispace kernel_inps' [dest_t] $
       KernelBody () kstms [res]
 
   traverse renameStm <=< runBuilder_ $ do
     addStms prestms
-
-    let pat =
-          Pat . rearrangeShape perm $
-            patElems $
-              loopNestingPat $
-                fst nest
-
+    let pat = Pat . rearrangeShape perm $ patElems $ loopNestingPat $ fst nest
     letBind pat $ Op $ segOp k
 
 segmentedGatherKernel ::
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
@@ -267,6 +267,11 @@
           certifying (stmAuxCerts aux) $
             addStms =<< segRed lvl pat mempty w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []
           parallelMin [w]
+    Op (Screma w arrs form) ->
+      -- This screma is too complicated for us to immediately do
+      -- anything, so split it up and try again.
+      mapM_ intraGroupStm . fmap (certify (stmAuxCerts aux)) . snd
+        =<< runBuilderT (dissectScrema pat w form arrs) (scopeForSOACs scope)
     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
@@ -292,15 +297,14 @@
       space <- mkSegSpace [(write_i, w)]
 
       let lam' = soacsLambdaToGPU lam
-          (dests_ws, _, _) = unzip3 dests
           krets = do
-            (a_w, a, is_vs) <-
+            (_a_w, a, is_vs) <-
               groupScatterResults dests $ bodyResult $ lambdaBody lam'
             let cs =
                   foldMap (foldMap resCerts . fst) is_vs
                     <> foldMap (resCerts . snd) is_vs
                 is_vs' = [(Slice $ map (DimFix . resSubExp) is, resSubExp v) | (is, v) <- is_vs]
-            pure $ WriteReturns cs a_w a is_vs'
+            pure $ WriteReturns cs a is_vs'
           inputs = do
             (p, p_a) <- zip (lambdaParams lam') ivs
             pure $ KernelInput (paramName p) (paramType p) p_a [Var write_i]
@@ -311,9 +315,8 @@
           addStms $ bodyStms $ lambdaBody lam'
 
       certifying (stmAuxCerts aux) $ do
-        let ts = zipWith (stripArray . length) dests_ws $ patTypes pat
-            body = KernelBody () kstms krets
-        letBind pat $ Op $ SegOp $ SegMap lvl space ts body
+        let body = KernelBody () kstms krets
+        letBind pat $ Op $ SegOp $ SegMap lvl space (patTypes pat) body
 
       parallelMin [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
@@ -27,7 +27,6 @@
 import Futhark.Pass.ExtractKernels.ToGPU (injectSOACS)
 import Futhark.Tools
 import Futhark.Transform.Rename (Rename, renameSomething)
-import Futhark.Util (takeLast)
 import Futhark.Util.Log
 
 newtype ExtractM a = ExtractM (ReaderT (Scope MC) (State VNameSource) a)
@@ -258,21 +257,15 @@
 
   Body () kstms res <- mapLambdaToBody transformBody gtid lam ivs
 
-  let rets = takeLast (length dests) $ lambdaReturnType lam
-      kres = do
-        (a_w, a, is_vs) <- groupScatterResults dests res
-        let cs =
-              foldMap (foldMap resCerts . fst) is_vs
-                <> foldMap (resCerts . snd) is_vs
-            is_vs' = [(Slice $ map (DimFix . resSubExp) is, resSubExp v) | (is, v) <- is_vs]
-        pure $ WriteReturns cs a_w a is_vs'
-      kbody = KernelBody () kstms kres
-  pure $
-    oneStm $
-      Let pat (defAux ()) $
-        Op $
-          ParOp Nothing $
-            SegMap () space rets kbody
+  (rets, kres) <- fmap unzip $ forM (groupScatterResults dests res) $ \(_a_w, a, is_vs) -> do
+    a_t <- lookupType a
+    let cs =
+          foldMap (foldMap resCerts . fst) is_vs
+            <> foldMap (resCerts . snd) is_vs
+        is_vs' = [(fullSlice a_t $ map (DimFix . resSubExp) is, resSubExp v) | (is, v) <- is_vs]
+    pure (a_t, WriteReturns cs a is_vs')
+  pure . oneStm . Let pat (defAux ()) . Op . ParOp Nothing $
+    SegMap () space rets (KernelBody () kstms kres)
 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/Passes.hs b/src/Futhark/Passes.hs
--- a/src/Futhark/Passes.hs
+++ b/src/Futhark/Passes.hs
@@ -93,8 +93,6 @@
         tileLoops,
         simplifyGPU,
         histAccsGPU,
-        babysitKernels,
-        simplifyGPU,
         unstreamGPU,
         performCSE True,
         simplifyGPU,
@@ -105,7 +103,12 @@
         mergeGPUBodies,
         simplifyGPU, -- Cleanup merged GPUBody kernels.
         sinkGPU, -- Sink reads within GPUBody kernels.
-        inPlaceLoweringGPU
+        inPlaceLoweringGPU,
+        babysitKernels,
+        -- Important to simplify after babysitting in order to fix up
+        -- redundant manifests.
+        simplifyGPU,
+        performCSE True
       ]
 
 -- | The pipeline used by the sequential backends.  Turns all
diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs
--- a/src/Futhark/Test.hs
+++ b/src/Futhark/Test.hs
@@ -300,11 +300,11 @@
 testRunReferenceOutput prog entry tr =
   "data"
     </> takeBaseName prog
-    <> ":"
-    <> T.unpack entry
-    <> "-"
-    <> map clean (T.unpack (runDescription tr))
-    <.> "out"
+      <> ":"
+      <> T.unpack entry
+      <> "-"
+      <> map clean (T.unpack (runDescription tr))
+        <.> "out"
   where
     clean '/' = '_' -- Would this ever happen?
     clean ' ' = '_'
diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs
--- a/src/Futhark/Transform/Rename.hs
+++ b/src/Futhark/Transform/Rename.hs
@@ -35,7 +35,7 @@
 import Data.Maybe
 import Futhark.FreshNames hiding (newName)
 import Futhark.IR.Prop.Names
-import Futhark.IR.Prop.Patterns
+import Futhark.IR.Prop.Pat
 import Futhark.IR.Syntax
 import Futhark.IR.Traversals
 import Futhark.MonadFreshNames (MonadFreshNames (..), modifyNameSource, newName)
diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs
--- a/src/Futhark/Util.hs
+++ b/src/Futhark/Util.hs
@@ -447,9 +447,10 @@
 -- constructing a set of corresponding values.
 invertMap :: (Ord v, Ord k) => M.Map k v -> M.Map v (S.Set k)
 invertMap m =
-  M.toList m
-    & fmap (swap . first S.singleton)
-    & foldr (uncurry $ M.insertWith (<>)) mempty
+  foldr
+    (uncurry (M.insertWith (<>)) . swap . first S.singleton)
+    mempty
+    (M.toList m)
 
 -- | Compute the cartesian product of two foldable collections, using the given
 -- combinator function.
diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs
--- a/src/Language/Futhark/Interpreter.hs
+++ b/src/Language/Futhark/Interpreter.hs
@@ -1402,7 +1402,7 @@
               <+> dquotes (prettyValue x)
               <+> "and"
               <+> dquotes (prettyValue y)
-                <> "."
+              <> "."
       where
         bopDef' (valf, retf, op) (x, y) = do
           x' <- valf x
@@ -1419,7 +1419,7 @@
           bad noLoc mempty . docText $
             "Cannot apply function to argument"
               <+> dquotes (prettyValue x)
-                <> "."
+              <> "."
       where
         unopDef' (valf, retf, op) x = do
           x' <- valf x
@@ -1436,7 +1436,8 @@
         _ ->
           bad noLoc mempty . docText $
             "Cannot apply operator to argument"
-              <+> dquotes (prettyValue v) <> "."
+              <+> dquotes (prettyValue v)
+              <> "."
 
     def "!" =
       Just $
@@ -1995,7 +1996,7 @@
           "Entry point "
             <> dquotes (prettyName entry)
             <> " expects input of type(s)"
-            </> indent 2 (stack (map pretty param_ts))
+              </> indent 2 (stack (map pretty param_ts))
 
 -- | Execute the named function on the given arguments; may fail
 -- horribly if these are ill-typed.
diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs
--- a/src/Language/Futhark/Pretty.hs
+++ b/src/Language/Futhark/Pretty.hs
@@ -232,8 +232,8 @@
       ( hsep (map (brackets . prettyName) sizeparams ++ [pretty pat])
           <+> equals
           <+> pretty initexp
-          <+> pretty form
-          <+> "do"
+          </> pretty form
+          </> "do"
       )
     </> indent 2 (pretty loopbody)
 prettyAppExp _ (Index e idxs _) =
@@ -258,9 +258,9 @@
   "let"
     <+> hsep (prettyName fname : map pretty tparams ++ map pretty params)
     <> retdecl'
-    <+> equals
-    </> indent 2 (pretty e)
-    </> letBody body
+      <+> equals
+      </> indent 2 (pretty e)
+      </> letBody body
   where
     retdecl' = case (pretty <$> unAnnot rettype) `mplus` (pretty <$> retdecl) of
       Just rettype' -> colon <+> align rettype'
@@ -270,9 +270,9 @@
       "let"
         <+> pretty dest
         <> list (map pretty idxs)
-        <+> equals
-        <+> align (pretty ve)
-        </> letBody body
+          <+> equals
+          <+> align (pretty ve)
+          </> letBody body
   | otherwise =
       "let"
         <+> pretty dest
@@ -374,8 +374,8 @@
     "\\"
       <> hsep (map pretty params)
       <> ppAscription rettype
-      <+> "->"
-      </> indent 2 (align (pretty body))
+        <+> "->"
+        </> indent 2 (align (pretty body))
 prettyExp _ (OpSection binop _ _) =
   parens $ pretty binop
 prettyExp _ (OpSectionLeft binop _ x _ _ _) =
@@ -484,8 +484,8 @@
     "\\"
       <> pretty param
       <> maybe_sig'
-      <+> "->"
-      </> indent 2 (pretty body)
+        <+> "->"
+        </> indent 2 (pretty body)
   where
     maybe_sig' = case maybe_sig of
       Nothing -> mempty
@@ -501,10 +501,11 @@
 
 instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where
   pretty (TypeBind name l params te rt _ _) =
-    "type" <> pretty l
-      <+> hsep (prettyName name : map pretty params)
-      <+> equals
-      <+> maybe (pretty te) pretty (unAnnot rt)
+    "type"
+      <> pretty l
+        <+> hsep (prettyName name : map pretty params)
+        <+> equals
+        <+> maybe (pretty te) pretty (unAnnot rt)
 
 instance (Eq vn, IsName vn) => Pretty (TypeParamBase vn) where
   pretty (TypeParamDim name _) = brackets $ prettyName name
@@ -514,16 +515,16 @@
   pretty (ValBind entry name retdecl rettype tparams args body _ attrs _) =
     mconcat (map ((<> line) . prettyAttr) attrs)
       <> fun
-      <+> align
-        ( sep
-            ( prettyName name
-                : map pretty tparams
-                ++ map pretty args
-                ++ retdecl'
-                ++ ["="]
-            )
-        )
-      </> indent 2 (pretty body)
+        <+> align
+          ( sep
+              ( prettyName name
+                  : map pretty tparams
+                  ++ map pretty args
+                  ++ retdecl'
+                  ++ ["="]
+              )
+          )
+        </> indent 2 (pretty body)
     where
       fun
         | isJust entry = "entry"
diff --git a/src/Language/Futhark/Semantic.hs b/src/Language/Futhark/Semantic.hs
--- a/src/Language/Futhark/Semantic.hs
+++ b/src/Language/Futhark/Semantic.hs
@@ -177,9 +177,9 @@
       renderTypeBind (name, TypeAbbr l tps tp) =
         p l
           <+> prettyName name
-            <> mconcat (map ((" " <>) . pretty) tps)
-            <> " ="
-          <+> pretty tp
+          <> mconcat (map ((" " <>) . pretty) tps)
+          <> " ="
+            <+> pretty tp
         where
           p Lifted = "type^"
           p SizeLifted = "type~"
@@ -187,9 +187,9 @@
       renderValBind (name, BoundV tps t) =
         "val"
           <+> prettyName name
-            <> mconcat (map ((" " <>) . pretty) tps)
-            <> " ="
-          <+> pretty t
+          <> mconcat (map ((" " <>) . pretty) tps)
+          <> " ="
+            <+> pretty t
       renderModType (name, _sig) =
         "module type" <+> prettyName name
       renderMod (name, mod) =
diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs
--- a/src/Language/Futhark/TypeChecker.hs
+++ b/src/Language/Futhark/TypeChecker.hs
@@ -160,8 +160,8 @@
       <+> pretty space
       <+> prettyName name
       <> "."
-      </> "Previously defined at"
-      <+> pretty (locStr loc2)
+        </> "Previously defined at"
+        <+> pretty (locStr loc2)
       <> "."
 
 checkForDuplicateDecs :: [DecBase NoInfo Name] -> TypeM ()
diff --git a/src/Language/Futhark/TypeChecker/Consumption.hs b/src/Language/Futhark/TypeChecker/Consumption.hs
--- a/src/Language/Futhark/TypeChecker/Consumption.hs
+++ b/src/Language/Futhark/TypeChecker/Consumption.hs
@@ -168,7 +168,8 @@
 returnAliased name loc =
   addError loc mempty . withIndexLink "return-aliased" $
     "Unique-typed return value is aliased to"
-      <+> dquotes (prettyName name) <> ", which is not consumable."
+      <+> dquotes (prettyName name)
+      <> ", which is not consumable."
 
 uniqueReturnAliased :: SrcLoc -> CheckM ()
 uniqueReturnAliased loc =
@@ -251,7 +252,7 @@
     bind env =
       env
         { envVtable =
-            foldr (uncurry M.insert) (envVtable env) (fmap f (matchPat p t))
+            foldr (uncurry M.insert . f) (envVtable env) (matchPat p t)
         }
       where
         f (v, (_, als)) = (v, Consumable $ second (S.insert (AliasBound v)) als)
@@ -264,7 +265,7 @@
     bind env =
       env
         { envVtable =
-            foldr (uncurry M.insert) (envVtable env) (fmap f (patternMap p))
+            foldr (uncurry M.insert . f) (envVtable env) (patternMap p)
         }
     f (v, t)
       | diet t == Consume = (v, Consumable $ t `setAliases` S.singleton (AliasBound v))
@@ -305,8 +306,10 @@
     v' <- describeVar v
     addError rloc mempty . withIndexLink "use-after-consume" $
       "Using"
-        <+> v' <> ", but this was consumed at"
-        <+> pretty (locStrRel rloc wloc) <> ".  (Possibly through aliases.)"
+        <+> v'
+        <> ", but this was consumed at"
+          <+> pretty (locStrRel rloc wloc)
+        <> ".  (Possibly through aliases.)"
 
 consumed :: Consumed -> CheckM ()
 consumed vs = modify $ \s -> s {stateConsumed = stateConsumed s <> vs}
@@ -597,7 +600,8 @@
             "Return value for consuming loop parameter"
               <+> dquotes (prettyName pat_v)
               <+> "aliases"
-              <+> dquotes (prettyName v) <> "."
+              <+> dquotes (prettyName v)
+              <> "."
         (cons, obs) <- get
         unless (S.null $ aliases t `S.intersection` cons) $
           lift . addError loop_loc mempty $
@@ -668,7 +672,7 @@
       "Loop body uses"
         <+> v'
         <> " (or an alias),"
-        </> "but this is consumed by the initial loop argument."
+          </> "but this is consumed by the initial loop argument."
 
   v <- VName "internal_loop_result" <$> incCounter
   modify $ \s -> s {stateNames = M.insert v (NameLoopRes (srclocOf loop_loc)) $ stateNames s}
@@ -959,9 +963,9 @@
       "Function result aliases the free variable "
         <> dquotes (prettyName v)
         <> "."
-        </> "Use"
-        <+> dquotes "copy"
-        <+> "to break the aliasing."
+          </> "Use"
+          <+> dquotes "copy"
+          <+> "to break the aliasing."
 
 -- | Type-check a value definition.  This also infers a new return
 -- type that may be more unique than previously.
diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs
--- a/src/Language/Futhark/TypeChecker/Modules.hs
+++ b/src/Language/Futhark/TypeChecker/Modules.hs
@@ -207,7 +207,10 @@
         else
           typeError loc mempty $
             "Cannot refine a type having"
-              <+> tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "."
+              <+> tpMsg ps
+              <> " with a type having "
+              <> tpMsg cur_ps
+              <> "."
   | otherwise =
       typeError loc mempty $ dquotes (pretty tname) <+> "is not an abstract type in the module type."
   where
@@ -367,14 +370,16 @@
 ppTypeAbbr abs name (l, ps, RetType [] (Scalar (TypeVar _ tn args)))
   | qualLeaf tn `elem` abs,
     map typeParamToArg ps == args =
-      "type" <> pretty l
-        <+> pretty name
-        <+> hsep (map pretty ps)
+      "type"
+        <> pretty l
+          <+> pretty name
+          <+> hsep (map pretty ps)
 ppTypeAbbr _ name (l, ps, t) =
-  "type" <> pretty l
-    <+> hsep (pretty name : map pretty ps)
-    <+> equals
-    <+> nest 2 (align (pretty t))
+  "type"
+    <> pretty l
+      <+> hsep (pretty name : map pretty ps)
+      <+> equals
+      <+> nest 2 (align (pretty t))
 
 -- | Return new renamed/abstracted env, as well as a mapping from
 -- names in the signature to names in the new env.  This is used for
diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs
--- a/src/Language/Futhark/TypeChecker/Monad.hs
+++ b/src/Language/Futhark/TypeChecker/Monad.hs
@@ -153,7 +153,7 @@
   typeError loc mempty $
     "Use of"
       <+> dquotes (pretty name)
-        <> ": variables prefixed with underscore may not be accessed."
+      <> ": variables prefixed with underscore may not be accessed."
 
 -- | A mapping from import import names to 'Env's.  This is used to
 -- resolve @import@ declarations.
diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs
--- a/src/Language/Futhark/TypeChecker/Terms.hs
+++ b/src/Language/Futhark/TypeChecker/Terms.hs
@@ -388,7 +388,8 @@
             "Field"
               <+> dquotes (pretty f)
               <+> "previously defined at"
-              <+> pretty (locStrRel rloc sloc) <> "."
+              <+> pretty (locStrRel rloc sloc)
+              <> "."
         Nothing -> pure ()
 checkExp (ArrayLit all_es _ loc) =
   -- Construct the result type and unify all elements with it.  We
@@ -1067,13 +1068,13 @@
           <+> fname'
           <+> "to argument #"
           <> pretty (prev_applied + 1)
-          <+> dquotes (shorten $ group $ pretty argexp)
+            <+> dquotes (shorten $ group $ pretty argexp)
           <> ","
-          </> "as"
-          <+> fname'
-          <+> "only takes"
-          <+> pretty prev_applied
-          <+> arguments
+            </> "as"
+            <+> fname'
+            <+> "only takes"
+            <+> pretty prev_applied
+            <+> arguments
           <> "."
   where
     arguments
@@ -1234,19 +1235,19 @@
           <+> "needed for type of"
           <+> what
           <> colon
-          </> indent 2 (pretty t)
-          </> "But"
-          <+> dquotes (prettyName d)
-          <+> "is computed at"
-          <+> pretty (locStrRel loc dloc)
+            </> indent 2 (pretty t)
+            </> "But"
+            <+> dquotes (prettyName d)
+            <+> "is computed at"
+            <+> pretty (locStrRel loc dloc)
           <> "."
-          </> ""
-          </> "Hint:"
-          <+> align
-            ( textwrap "Bind the expression producing"
-                <+> dquotes (prettyName d)
-                <+> "with 'let' beforehand."
-            )
+            </> ""
+            </> "Hint:"
+            <+> align
+              ( textwrap "Bind the expression producing"
+                  <+> dquotes (prettyName d)
+                  <+> "with 'let' beforehand."
+              )
 
 -- | Traverse the expression, emitting warnings and errors for various
 -- problems:
@@ -1393,7 +1394,7 @@
             "Type is ambiguous (could be one of"
               <+> commasep (map pretty ots)
               <> ")."
-              </> "Add a type annotation to disambiguate the type."
+                </> "Add a type annotation to disambiguate the type."
     fixOverloaded (v, NoConstraint _ usage) = do
       -- See #1552.
       unify usage (Scalar (TypeVar mempty (qualName v) [])) $
@@ -1416,7 +1417,7 @@
         "Type is ambiguous (must be a sum type with constructors:"
           <+> pretty (Sum cs)
           <> ")."
-          </> "Add a type annotation to disambiguate the type."
+            </> "Add a type annotation to disambiguate the type."
     fixOverloaded (v, Size Nothing (Usage Nothing loc)) =
       typeError loc mempty . withIndexLink "ambiguous-size" $
         "Ambiguous size" <+> dquotes (prettyName v) <> "."
@@ -1540,10 +1541,10 @@
               </> "refers to size"
               <+> dquotes (prettyName d)
               <> comma
-              </> textwrap "which will not be accessible to the caller"
+                </> textwrap "which will not be accessible to the caller"
               <> comma
-              </> textwrap "possibly because it is nested in a tuple or record."
-              </> textwrap "Consider ascribing an explicit type that does not reference "
+                </> textwrap "possibly because it is nested in a tuple or record."
+                </> textwrap "Consider ascribing an explicit type that does not reference "
               <> dquotes (prettyName d)
               <> "."
       | otherwise = verifyParams forbidden' ps
@@ -1638,7 +1639,7 @@
               <+> "in parameter of"
               <+> dquotes (prettyName defname)
               <> ", which is inferred as:"
-              </> indent 2 (pretty t)
+                </> indent 2 (pretty t)
       | k `S.member` produced_sizes =
           pure $ Just $ Right k
     closeOver (_, _) =
diff --git a/src/Language/Futhark/TypeChecker/Terms/Monad.hs b/src/Language/Futhark/TypeChecker/Terms/Monad.hs
--- a/src/Language/Futhark/TypeChecker/Terms/Monad.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/Monad.hs
@@ -112,12 +112,14 @@
         case f of
           Nothing ->
             "Cannot apply function to"
-              <+> dquotes (shorten $ group $ pretty e) <> " (invalid type)."
+              <+> dquotes (shorten $ group $ pretty e)
+              <> " (invalid type)."
           Just fname ->
             "Cannot apply"
               <+> dquotes (pretty fname)
               <+> "to"
-              <+> dquotes (align $ shorten $ group $ pretty e) <> " (invalid type)."
+              <+> dquotes (align $ shorten $ group $ pretty e)
+              <> " (invalid type)."
   pretty (CheckingReturn expected actual) =
     "Function body does not have expected type."
       </> "Expected:"
@@ -159,24 +161,24 @@
     "Type mismatch when updating record field"
       <+> dquotes fs'
       <> "."
-      </> "Existing:"
-      <+> align (pretty expected)
-      </> "New:     "
-      <+> align (pretty actual)
+        </> "Existing:"
+        <+> align (pretty expected)
+        </> "New:     "
+        <+> align (pretty actual)
     where
       fs' = mconcat $ punctuate "." $ map pretty fs
   pretty (CheckingRequired [expected] actual) =
     "Expression must must have type"
       <+> pretty expected
       <> "."
-      </> "Actual type:"
-      <+> align (pretty actual)
+        </> "Actual type:"
+        <+> align (pretty actual)
   pretty (CheckingRequired expected actual) =
     "Type of expression must must be one of "
       <+> expected'
       <> "."
-      </> "Actual type:"
-      <+> align (pretty actual)
+        </> "Actual type:"
+        <+> align (pretty actual)
     where
       expected' = commasep (map pretty expected)
   pretty (CheckingBranches t1 t2) =
diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
--- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
@@ -263,11 +263,13 @@
   | Just ts <- M.lookup n cs = do
       when (length ps /= length ts) $
         typeError loc mempty $
-          "Pattern #" <> pretty n <> " expects"
-            <+> pretty (length ps)
-            <+> "constructor arguments, but type provides"
-            <+> pretty (length ts)
-            <+> "arguments."
+          "Pattern #"
+            <> pretty n
+            <> " expects"
+              <+> pretty (length ps)
+              <+> "constructor arguments, but type provides"
+              <+> pretty (length ts)
+              <+> "arguments."
       ps' <- zipWithM (checkPat' sizes) ps $ map Ascribed ts
       pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc
 checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed t) = do
diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs
--- a/src/Language/Futhark/TypeChecker/Types.hs
+++ b/src/Language/Futhark/TypeChecker/Types.hs
@@ -259,7 +259,8 @@
           <+> "requires"
           <+> pretty (length ps)
           <+> "arguments, but provided"
-          <+> pretty (length targs) <> "."
+          <+> pretty (length targs)
+          <> "."
     else do
       (targs', dims, substs) <- unzip3 <$> zipWithM checkArgApply ps targs
       pure
@@ -314,7 +315,8 @@
         "Type argument"
           <+> pretty a
           <+> "not valid for a type parameter"
-          <+> pretty p <> "."
+          <+> pretty p
+          <> "."
 
 -- | Check a type expression, producing:
 --
@@ -359,7 +361,8 @@
               "Name"
                 <+> dquotes (pretty v)
                 <+> "also bound at"
-                <+> pretty (locStr prev_loc) <> "."
+                <+> pretty (locStr prev_loc)
+                <> "."
         Nothing ->
           modify $ M.insert (ns, v) loc
 
@@ -379,7 +382,8 @@
         "Name"
           <+> dquotes (pretty v)
           <+> "also bound at"
-          <+> pretty (locStr prev_loc) <> "."
+          <+> pretty (locStr prev_loc)
+          <> "."
 
     check seen (TEArrow (Just v) t1 t2 loc)
       | Just prev_loc <- M.lookup v seen =
@@ -433,7 +437,8 @@
               "Type parameter"
                 <+> dquotes (pretty v)
                 <+> "previously defined at"
-                <+> pretty (locStr prev) <> "."
+                <+> pretty (locStr prev)
+                <> "."
         Nothing -> do
           modify $ M.insert (ns, v) loc
           lift $ checkName ns v loc
diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs
--- a/src/Language/Futhark/TypeChecker/Unify.hs
+++ b/src/Language/Futhark/TypeChecker/Unify.hs
@@ -58,7 +58,8 @@
       </> indent 2 (pretty t2)
   pretty (MatchingFields fields) =
     "When matching types of record field"
-      <+> dquotes (mconcat $ punctuate "." $ map pretty fields) <> dot
+      <+> dquotes (mconcat $ punctuate "." $ map pretty fields)
+      <> dot
   pretty (MatchingConstructor c) =
     "When matching types of constructor" <+> dquotes (pretty c) <> dot
   pretty (Matching s) =
@@ -188,12 +189,14 @@
 prettySource :: SrcLoc -> SrcLoc -> RigidSource -> Doc ()
 prettySource ctx loc (RigidRet Nothing) =
   "is unknown size returned by function at"
-    <+> pretty (locStrRel ctx loc) <> "."
+    <+> pretty (locStrRel ctx loc)
+    <> "."
 prettySource ctx loc (RigidRet (Just fname)) =
   "is unknown size returned by"
     <+> dquotes (pretty fname)
     <+> "at"
-    <+> pretty (locStrRel ctx loc) <> "."
+    <+> pretty (locStrRel ctx loc)
+    <> "."
 prettySource ctx loc (RigidArg fname arg) =
   "is value of argument"
     </> indent 2 (shorten (pretty arg))
@@ -209,7 +212,7 @@
     </> indent 2 (shorten (pretty slice))
     </> d_desc
     <> "at"
-    <+> pretty (locStrRel ctx loc)
+      <+> pretty (locStrRel ctx loc)
     <> "."
   where
     d_desc = case d of
@@ -231,21 +234,22 @@
     <> " going out of scope at "
     <> pretty (locStrRel ctx loc)
     <> "."
-    </> "Originally bound at "
+      </> "Originally bound at "
     <> pretty (locStrRel ctx boundloc)
     <> "."
 prettySource ctx loc RigidCoerce =
   "is an unknown size arising from empty dimension in coercion at"
-    <+> pretty (locStrRel ctx loc) <> "."
+    <+> pretty (locStrRel ctx loc)
+    <> "."
 prettySource _ _ RigidUnify =
   "is an artificial size invented during unification of functions with anonymous sizes."
 prettySource ctx loc (RigidCond t1 t2) =
   "is unknown due to conditional expression at "
     <> pretty (locStrRel ctx loc)
     <> "."
-    </> "One branch returns array of type: "
+      </> "One branch returns array of type: "
     <> align (pretty t1)
-    </> "The other an array of type:       "
+      </> "The other an array of type:       "
     <> align (pretty t2)
 
 -- | Retrieve notes describing the purpose or origin of the given
@@ -585,7 +589,8 @@
       "Occurs check: cannot instantiate"
         <+> prettyName vn
         <+> "with"
-        <+> pretty tp <> "."
+        <+> pretty tp
+        <> "."
 
 scopeCheck ::
   (MonadUnify m) =>
@@ -857,7 +862,9 @@
       unifyError usage mempty noBreadCrumbs $
         "Cannot unify type"
           <+> dquotes (pretty t)
-          <+> "with any of " <> commasep (map pretty ts) <> "."
+          <+> "with any of "
+          <> commasep (map pretty ts)
+          <> "."
 
 linkVarToTypes :: (MonadUnify m) => Usage -> VName -> [PrimType] -> m ()
 linkVarToTypes usage vn ts = do
@@ -872,22 +879,23 @@
               <+> "but also one of"
               <+> commasep (map pretty vn_ts)
               <+> "due to"
-              <+> pretty vn_usage <> "."
+              <+> pretty vn_usage
+              <> "."
         ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage)
     Just (_, HasConstrs _ _ vn_usage) ->
       unifyError usage mempty noBreadCrumbs $
         "Type constrained to one of"
           <+> commasep (map pretty ts)
-            <> ", but also inferred to be sum type due to"
-          <+> pretty vn_usage
-            <> "."
+          <> ", but also inferred to be sum type due to"
+            <+> pretty vn_usage
+          <> "."
     Just (_, HasFields _ _ vn_usage) ->
       unifyError usage mempty noBreadCrumbs $
         "Type constrained to one of"
           <+> commasep (map pretty ts)
-            <> ", but also inferred to be record due to"
-          <+> pretty vn_usage
-            <> "."
+          <> ", but also inferred to be record due to"
+            <+> pretty vn_usage
+          <> "."
     Just (lvl, _) -> modifyConstraints $ M.insert vn (lvl, Overloaded ts usage)
     Nothing ->
       unifyError usage mempty noBreadCrumbs $
@@ -1118,7 +1126,8 @@
             "Attempt to access field"
               <+> dquotes (pretty l)
               <+> " of value of type"
-              <+> pretty (toStructural t) <> "."
+              <+> pretty (toStructural t)
+              <> "."
     _ -> do
       unify usage t $ Scalar $ Record $ M.singleton l l_type
       pure l_type