diff --git a/docs/installation.rst b/docs/installation.rst
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -197,7 +197,7 @@
 have installed otherwise.
 
 * On NixOS, for OpenCL, you should import ``opencl-headers`` and
-  ``opencl-icd``.  You also need some form of OpenCL backend.  If you
+  ``ocl-icd``.  You also need some form of OpenCL backend.  If you
   have an AMD GPU and use ROCm, you may also need
   ``rocm-opencl-runtime``.
 
@@ -212,4 +212,4 @@
 
 These can be easily made available with e.g::
 
-  nix-shell -p opencl-headers -p opencl-icd
+  nix-shell -p opencl-headers -p ocl-icd
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -1,6 +1,6 @@
 cabal-version: 2.4
 name:           futhark
-version:        0.22.2
+version:        0.22.3
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs
--- a/src/Futhark/AD/Rev/Scan.hs
+++ b/src/Futhark/AD/Rev/Scan.hs
@@ -79,12 +79,7 @@
 -- but insert explicit indexing to reverse inside the map.
 mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName
 mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do
-  nm1 <- letSubExp "nm1" =<< toExp (pe64 w - 1)
-  y_adj_last <-
-    letExp (baseString y_adj ++ "_last") $
-      BasicOp $
-        Index y_adj $
-          fullSlice arr_tp [DimFix nm1]
+  y_adj_last <- letExp (baseString y_adj <> "_last") =<< eLast y_adj
 
   par_i <- newParam "i" $ Prim int64
   lam <- mkLambda [par_i] $ do
@@ -121,9 +116,8 @@
           (resultBodyM $ map (Var . paramName) par_r)
           ( buildBody_ $ do
               im1 <- letSubExp "im1" =<< toExp (le64 i - 1)
-              ys_im1 <- forM ys $ \y -> do
-                y_t <- lookupType y
-                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1]
+              ys_im1 <- forM ys $ \y ->
+                letSubExp (baseString y <> "_im1") =<< eIndex y (eSubExp im1)
 
               lam_res <-
                 mapM (letExp "const" . BasicOp . SubExp . resSubExp)
diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs
--- a/src/Futhark/Actions.hs
+++ b/src/Futhark/Actions.hs
@@ -4,6 +4,7 @@
   ( printAction,
     printAliasesAction,
     printLastUseGPU,
+    printFusionGraph,
     printInterferenceGPU,
     printMemAliasGPU,
     callGraphAction,
@@ -54,6 +55,7 @@
 import Futhark.IR.Prop.Aliases
 import Futhark.IR.SOACS (SOACS)
 import Futhark.IR.SeqMem (SeqMem)
+import Futhark.Optimise.Fusion.GraphRep qualified
 import Futhark.Util (runProgramWithExitCode, unixEnvironment)
 import Futhark.Version (versionString)
 import System.Directory
@@ -86,6 +88,22 @@
     { actionName = "print last use gpu",
       actionDescription = "Print last use information on gpu.",
       actionProcedure = liftIO . print . LastUse.analyseGPUMem
+    }
+
+-- | Print fusion graph to stdout.
+printFusionGraph :: Action SOACS
+printFusionGraph =
+  Action
+    { actionName = "print fusion graph",
+      actionDescription = "Print fusion graph in Graphviz format.",
+      actionProcedure =
+        liftIO
+          . mapM_
+            ( putStrLn
+                . Futhark.Optimise.Fusion.GraphRep.pprg
+                . Futhark.Optimise.Fusion.GraphRep.mkDepGraphForFun
+            )
+          . progFuns
     }
 
 -- | Print interference information to stdout.
diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs
--- a/src/Futhark/Analysis/Interference.hs
+++ b/src/Futhark/Analysis/Interference.hs
@@ -16,7 +16,7 @@
 import Futhark.Analysis.LastUse qualified as LastUse
 import Futhark.Analysis.MemAlias qualified as MemAlias
 import Futhark.IR.GPUMem
-import Futhark.Util (invertMap)
+import Futhark.Util (cartesian, invertMap)
 
 -- | The set of 'VName' currently in use.
 type InUse = Names
@@ -36,12 +36,6 @@
 makeEdge v1 v2
   | v1 == v2 = mempty
   | otherwise = S.singleton (min v1 v2, max v1 v2)
-
--- | Compute the cartesian product of two foldable collections, using the given
--- combinator function.
-cartesian :: (Monoid m, Foldable t) => (a -> a -> m) -> t a -> t a -> m
-cartesian f xs ys =
-  foldMap (uncurry f) [(x, y) | x <- toList xs, y <- toList ys]
 
 analyseStm ::
   LocalScope GPUMem m =>
diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs
--- a/src/Futhark/Analysis/LastUse.hs
+++ b/src/Futhark/Analysis/LastUse.hs
@@ -136,6 +136,16 @@
         used_acc <> unAliases aliases
       )
 
+    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem {}, Constant _) = (lumap_acc, used_acc)
+    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem name _, Var body_res) =
+      -- Any aliases of `name` should have the same last-use as `name`
+      ( case M.lookup name lumap_acc of
+          Just name' ->
+            insertNames name' (oneName body_res) lumap_acc
+          Nothing -> lumap_acc,
+        used_acc <> oneName body_res
+      )
+
     pat_name = patElemName $ head $ patElems pat
 
     analyseExp (lumap, used) (BasicOp _) = do
@@ -149,16 +159,21 @@
         bimap mconcat mconcat . unzip
           <$> mapM (analyseBody lumap used . caseBody) cases
       (lumap_defbody, used_defbody) <- analyseBody lumap used defbody
-      let used' = used_cases <> used_defbody
+      let (lumap', used') =
+            (lumap_defbody <> lumap_cases, used_cases <> used_defbody)
+              & flip (foldl addAliasesFromBodyRes) (zip (patElems pat) (map resSubExp $ bodyResult defbody))
           nms = (freeIn ses <> freeIn dec) `namesSubtract` used'
       pure
-        ( insertNames pat_name nms (lumap_cases <> lumap_defbody),
+        ( insertNames pat_name nms lumap',
           used' <> nms
         )
     analyseExp (lumap, used) (DoLoop merge form body) = do
-      (lumap', used') <- analyseBody lumap used body
-      let nms = (freeIn merge <> freeIn form) `namesSubtract` used'
-      pure (insertNames pat_name nms lumap', used' <> nms)
+      let (lumap', used') =
+            zip (patElems pat) (map resSubExp $ bodyResult body)
+              & foldl addAliasesFromBodyRes (lumap, used)
+      (lumap'', used'') <- analyseBody lumap' used' body
+      let nms = (freeIn merge <> freeIn form) `namesSubtract` used''
+      pure (insertNames pat_name nms lumap'', used'' <> nms)
     analyseExp (lumap, used) (Op op) = do
       onOp <- asks envLastUseOp
       onOp pat_name (lumap, used) op
diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs
--- a/src/Futhark/Analysis/PrimExp.hs
+++ b/src/Futhark/Analysis/PrimExp.hs
@@ -228,6 +228,20 @@
   | oneIshExp y = y
   | zeroIshExp x = y
   | zeroIshExp y = x
+constFoldPrimExp (UnOpExp Abs {} x)
+  | not $ negativeIshExp x = x
+constFoldPrimExp (BinOpExp UMod {} x y)
+  | sameIshExp x y,
+    IntType it <- primExpType x =
+      ValueExp $ IntValue $ intValue it (0 :: Integer)
+constFoldPrimExp (BinOpExp SMod {} x y)
+  | sameIshExp x y,
+    IntType it <- primExpType x =
+      ValueExp $ IntValue $ intValue it (0 :: Integer)
+constFoldPrimExp (BinOpExp SRem {} x y)
+  | sameIshExp x y,
+    IntType it <- primExpType x =
+      ValueExp $ IntValue $ intValue it (0 :: Integer)
 constFoldPrimExp e = e
 
 -- | The class of numeric types that can be used for constructing
@@ -336,8 +350,8 @@
     | otherwise = numBad "*" (x, y)
 
   abs (TPrimExp x)
-    | IntType t <- primExpType x = TPrimExp $ UnOpExp (Abs t) x
-    | FloatType t <- primExpType x = TPrimExp $ UnOpExp (FAbs t) x
+    | IntType t <- primExpType x = TPrimExp $ constFoldPrimExp $ UnOpExp (Abs t) x
+    | FloatType t <- primExpType x = TPrimExp $ constFoldPrimExp $ UnOpExp (FAbs t) x
     | otherwise = numBad "abs" x
 
   signum (TPrimExp x)
@@ -419,7 +433,7 @@
 
   TPrimExp x `mod` TPrimExp y
     | Just z <- msum [asIntOp (`SMod` Unsafe) x y] =
-        TPrimExp z
+        TPrimExp $ constFoldPrimExp z
     | otherwise = numBad "mod" (x, y)
 
   TPrimExp x `quot` TPrimExp y
@@ -588,6 +602,15 @@
 oneIshExp :: PrimExp v -> Bool
 oneIshExp (ValueExp v) = oneIsh v
 oneIshExp _ = False
+
+-- | Is the expression a constant negative of some sort?
+negativeIshExp :: PrimExp v -> Bool
+negativeIshExp (ValueExp v) = negativeIsh v
+negativeIshExp _ = False
+
+sameIshExp :: PrimExp v -> PrimExp v -> Bool
+sameIshExp (ValueExp v1) (ValueExp v2) = v1 == v2
+sameIshExp _ _ = False
 
 -- | If the given 'PrimExp' is a constant of the wrong integer type,
 -- coerce it to the given integer type.  This is a workaround for an
diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs
--- a/src/Futhark/Analysis/UsageTable.hs
+++ b/src/Futhark/Analysis/UsageTable.hs
@@ -23,6 +23,7 @@
     consumedU,
     presentU,
     usageInStm,
+    usageInPat,
   )
 where
 
@@ -119,7 +120,7 @@
 -- | Construct a usage table where the given names have been used as
 -- an array or memory size.
 sizeUsages :: Names -> UsageTable
-sizeUsages = UsageTable . IM.map (const sizeU) . namesIntMap
+sizeUsages = UsageTable . IM.map (const (sizeU <> presentU)) . namesIntMap
 
 -- | A description of how a single variable has been used.
 newtype Usages = Usages Int -- Bitmap representation for speed.
@@ -155,20 +156,17 @@
 usageInStm :: (ASTRep rep, Aliased rep) => Stm rep -> UsageTable
 usageInStm (Let pat rep e) =
   mconcat
-    [ usageInPat,
-      usageInExpDec,
+    [ usageInPat pat `without` patNames pat,
+      usages $ freeIn rep,
       usageInExp e,
       usages (freeIn e)
     ]
-  where
-    usageInPat =
-      usages
-        ( mconcat (map freeIn $ patElems pat)
-            `namesSubtract` namesFromList (patNames pat)
-        )
-        <> sizeUsages (foldMap (freeIn . patElemType) (patElems pat))
-    usageInExpDec =
-      usages $ freeIn rep
+
+-- | Usage table reflecting use in pattern.  In particular, free
+-- variables in the decorations are considered used as sizes, even if
+-- they are also bound in this pattern.
+usageInPat :: FreeIn t => Pat t -> UsageTable
+usageInPat = sizeUsages . foldMap freeIn . patElems
 
 usageInExp :: Aliased rep => Exp rep -> UsageTable
 usageInExp (Apply _ args _ _) =
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
@@ -462,6 +462,11 @@
       "Print the resulting IR with aliases.",
     Option
       []
+      ["fusion-graph"]
+      (NoArg $ Right $ \opts -> opts {futharkAction = SOACSAction printFusionGraph})
+      "Print fusion graph.",
+    Option
+      []
       ["print-last-use-gpu"]
       ( NoArg $
           Right $ \opts ->
diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
@@ -282,6 +282,7 @@
                  $ty:et' *arr = calloc($exp:num_elems, $id:info.size);
                  assert(arr != NULL);
                  assert($id:values_array(ctx, $exp:e, arr) == 0);
+                 assert(futhark_context_sync(ctx) == 0);
                  write_array(stdout, binary_output, &$id:info, arr,
                              $id:shape_array(ctx, $exp:e), $int:rank);
                  free(arr);
@@ -304,10 +305,6 @@
       printstms =
         printResult manifest $ zip (map outputType outputs) output_vals
 
-      ctx_ty = [C.cty|struct futhark_context|]
-      sync_ctx = "futhark_context_sync" :: T.Text
-      error_ctx = "futhark_context_get_error" :: T.Text
-
       cli_entry_point_function_name = "futrts_cli_entry_" ++ T.unpack entry_point_name
 
       pause_profiling = "futhark_context_pause_profiling" :: T.Text
@@ -320,8 +317,8 @@
                 int r;
                 // Run the program once.
                 $stms:pack_input
-                if ($id:sync_ctx(ctx) != 0) {
-                  futhark_panic(1, "%s", $id:error_ctx(ctx));
+                if (futhark_context_sync(ctx) != 0) {
+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));
                 };
                 // Only profile last run.
                 if (profile_run) {
@@ -332,10 +329,10 @@
                              $args:(map addrOf output_vals),
                              $args:input_args);
                 if (r != 0) {
-                  futhark_panic(1, "%s", $id:error_ctx(ctx));
+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));
                 }
-                if ($id:sync_ctx(ctx) != 0) {
-                  futhark_panic(1, "%s", $id:error_ctx(ctx));
+                if (futhark_context_sync(ctx) != 0) {
+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));
                 };
                 if (profile_run) {
                   $id:pause_profiling(ctx);
@@ -349,7 +346,7 @@
                 $stms:free_input
               |]
    in ( [C.cedecl|
-   static int $id:cli_entry_point_function_name($ty:ctx_ty *ctx) {
+   static int $id:cli_entry_point_function_name(struct futhark_context *ctx) {
      typename int64_t t_start, t_end;
      int time_runs = 0, profile_run = 0;
      int retval = 0;
diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs
--- a/src/Futhark/CodeGen/ImpGen.hs
+++ b/src/Futhark/CodeGen/ImpGen.hs
@@ -109,6 +109,7 @@
     sWrite,
     sUpdate,
     sLoopNest,
+    sCopy,
     sLoopSpace,
     (<--),
     (<~~),
@@ -1497,7 +1498,7 @@
 
   pure fname
 
--- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'.
+-- | Use 'sCopy' if possible, otherwise 'copyElementWise'.
 defaultCopy :: CopyCompiler rep r op
 defaultCopy pt dest src
   | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <-
@@ -1524,17 +1525,7 @@
       destspace <- entryMemSpace <$> lookupMemory destmem
       if isScalarSpace srcspace || isScalarSpace destspace
         then copyElementWise pt dest src
-        else
-          emit
-            $ Imp.Copy
-              pt
-              destmem
-              (bytes destoffset)
-              destspace
-              srcmem
-              (bytes srcoffset)
-              srcspace
-            $ num_elems `withElemType` pt
+        else sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt
   | otherwise =
       copyElementWise pt dest src
   where
@@ -1926,6 +1917,33 @@
   ([Imp.TExp Int64] -> ImpM rep r op ()) ->
   ImpM rep r op ()
 sLoopNest = sLoopSpace . map pe64 . shapeDims
+
+sCopy ::
+  VName ->
+  Imp.TExp Int64 ->
+  Space ->
+  VName ->
+  Imp.TExp Int64 ->
+  Space ->
+  Count Elements (Imp.TExp Int64) ->
+  PrimType ->
+  ImpM rep r op ()
+sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt =
+  if destmem == srcmem
+    then sUnless (destoffset .==. srcoffset) the_copy
+    else the_copy
+  where
+    the_copy =
+      emit
+        $ Imp.Copy
+          pt
+          destmem
+          (bytes destoffset)
+          destspace
+          srcmem
+          (bytes srcoffset)
+          srcspace
+        $ num_elems `withElemType` pt
 
 -- | Untyped assignment.
 (<~~) :: VName -> Imp.Exp -> ImpM rep r op ()
diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU.hs
@@ -16,7 +16,6 @@
 import Data.List (foldl')
 import Data.Map qualified as M
 import Data.Maybe
-import Futhark.CodeGen.ImpCode.GPU (bytes)
 import Futhark.CodeGen.ImpCode.GPU qualified as Imp
 import Futhark.CodeGen.ImpGen hiding (compileProg)
 import Futhark.CodeGen.ImpGen qualified
@@ -166,12 +165,12 @@
   CallKernelGen ()
 segOpCompiler pat (SegMap lvl space _ kbody) =
   compileSegMap pat lvl space kbody
-segOpCompiler pat (SegRed lvl@SegThread {} space reds _ kbody) =
+segOpCompiler pat (SegRed lvl@(SegThread _ _) space reds _ kbody) =
   compileSegRed pat lvl space reds kbody
-segOpCompiler pat (SegScan lvl@SegThread {} space scans _ kbody) =
+segOpCompiler pat (SegScan lvl@(SegThread _ _) space scans _ kbody) =
   compileSegScan pat lvl space scans kbody
-segOpCompiler pat (SegHist (SegThread num_groups group_size _) space ops _ kbody) =
-  compileSegHist pat num_groups group_size space ops kbody
+segOpCompiler pat (SegHist lvl@(SegThread _ _) space ops _ kbody) =
+  compileSegHist pat lvl space ops kbody
 segOpCompiler pat segop =
   compilerBugS $ "segOpCompiler: unexpected " ++ prettyString (segLevel segop) ++ " for rhs of pattern " ++ prettyString pat
 
@@ -298,17 +297,16 @@
       let num_elems = Imp.elements $ product $ map pe64 srcshape
       srcspace <- entryMemSpace <$> lookupMemory srcmem
       destspace <- entryMemSpace <$> lookupMemory destmem
-      emit
-        $ Imp.Copy
-          bt
-          destmem
-          (bytes $ sExt64 destoffset)
-          destspace
-          srcmem
-          (bytes $ sExt64 srcoffset)
-          srcspace
-        $ num_elems `Imp.withElemType` bt
-  | otherwise = sCopy bt destloc srcloc
+      sCopy
+        destmem
+        (sExt64 destoffset)
+        destspace
+        srcmem
+        (sExt64 srcoffset)
+        srcspace
+        num_elems
+        bt
+  | otherwise = sCopyKernel bt destloc srcloc
 
 mapTransposeForType :: PrimType -> CallKernelGen Name
 mapTransposeForType bt = do
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
@@ -19,6 +19,8 @@
     sKernelThread,
     KernelAttrs (..),
     defKernelAttrs,
+    sCopyKernel,
+    lvlKernelAttrs,
     allocLocal,
     kernelAlloc,
     compileThreadResult,
@@ -1076,6 +1078,33 @@
       kAttrGroupSize = group_size
     }
 
+getSize :: String -> SizeClass -> CallKernelGen (TV Int64)
+getSize desc size_class = do
+  v <- dPrim desc int64
+  fname <- askFunction
+  let v_key = keyWithEntryPoint fname $ nameFromString $ prettyString $ tvVar v
+  sOp $ Imp.GetSize (tvVar v) v_key size_class
+  pure v
+
+-- | Compute kernel attributes from 'SegLevel'; including synthesising
+-- group-size and thread count if no grid is provided.
+lvlKernelAttrs :: SegLevel -> CallKernelGen KernelAttrs
+lvlKernelAttrs lvl =
+  case lvl of
+    SegThread _ Nothing -> mkGrid
+    SegThread _ (Just (KernelGrid num_groups group_size)) ->
+      pure $ defKernelAttrs num_groups group_size
+    SegGroup _ Nothing -> mkGrid
+    SegGroup _ (Just (KernelGrid num_groups group_size)) ->
+      pure $ defKernelAttrs num_groups group_size
+    SegThreadInGroup {} ->
+      error "lvlKernelAttrs: SegThreadInGroup"
+  where
+    mkGrid = do
+      group_size <- getSize "group_size" Imp.SizeGroup
+      num_groups <- getSize "num_groups" Imp.SizeNumGroups
+      pure $ defKernelAttrs (Count $ tvSize num_groups) (Count $ tvSize group_size)
+
 sKernel ::
   Operations GPUMem KernelEnv Imp.KernelOp ->
   (KernelConstants -> Imp.TExp Int32) ->
@@ -1318,8 +1347,8 @@
           [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s]
     else sIotaKernel arr n x s et
 
-sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp
-sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do
+sCopyKernel :: CopyCompiler GPUMem HostEnv Imp.HostOp
+sCopyKernel pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do
   -- Note that the shape of the destination and the source are
   -- necessarily the same.
   let shape = map pe64 srcdims
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
@@ -195,14 +195,8 @@
   bimap (map fst) (map fst) $
     partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..])
 
-sanityCheckLevel :: SegLevel -> InKernelGen ()
-sanityCheckLevel SegThread {} = pure ()
-sanityCheckLevel SegGroup {} =
-  error "compileGroupOp: unexpected group-level SegOp."
-
-compileFlatId :: SegLevel -> SegSpace -> InKernelGen ()
-compileFlatId lvl space = do
-  sanityCheckLevel lvl
+compileFlatId :: SegSpace -> InKernelGen ()
+compileFlatId space = do
   ltid <- kernelLocalThreadId . kernelConstants <$> askEnv
   dPrimV_ (segFlat space) ltid
 
@@ -342,7 +336,7 @@
 compileGroupOp pat (Alloc size space) =
   kernelAlloc pat size space
 compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do
-  compileFlatId lvl space
+  compileFlatId space
 
   groupCoverSegSpace (segVirt lvl) space $
     compileStms mempty (kernelBodyStms body) $
@@ -350,7 +344,7 @@
         kernelBodyResult body
   sOp $ Imp.ErrorSync Imp.FenceLocal
 compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do
-  compileFlatId lvl space
+  compileFlatId space
 
   let (ltids, dims) = unzip $ unSegSpace space
       dims' = map pe64 dims
@@ -397,7 +391,7 @@
         (segBinOpLambda scan)
         arrs_flat
 compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do
-  compileFlatId lvl space
+  compileFlatId space
 
   let dims' = map pe64 dims
       mkTempArr t =
@@ -512,7 +506,7 @@
 
       sOp $ Imp.Barrier Imp.FenceLocal
 compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do
-  compileFlatId lvl space
+  compileFlatId space
   let (ltids, _dims) = unzip $ unSegSpace space
 
   -- We don't need the red_pes, because it is guaranteed by our type
@@ -522,7 +516,8 @@
       (_red_pes, map_pes) =
         splitAt num_red_res $ patElems pat
 
-  ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops
+  group_size <- kernelGroupSizeCount . kernelConstants <$> askEnv
+  ops' <- prepareIntraGroupSegHist 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
@@ -1046,13 +1046,13 @@
 -- | Generate code for a segmented histogram called from the host.
 compileSegHist ::
   Pat LetDecMem ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
+  SegLevel ->
   SegSpace ->
   [HistOp GPUMem] ->
   KernelBody GPUMem ->
   CallKernelGen ()
-compileSegHist (Pat pes) num_groups group_size space ops kbody = do
+compileSegHist (Pat pes) lvl space ops kbody = do
+  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl
   -- Most of this function is not the histogram part itself, but
   -- rather figuring out whether to use a local or global memory
   -- strategy, as well as collapsing the subhistograms produced (which
@@ -1152,7 +1152,7 @@
 
         flat_gtid <- newVName "flat_gtid"
 
-        let lvl = SegThread num_groups group_size SegVirt
+        let grid = KernelGrid num_groups group_size
             segred_space =
               SegSpace flat_gtid $
                 segment_dims
@@ -1161,7 +1161,7 @@
                   ++ [(subhistogram_id, Var $ tvVar num_histos)]
 
         let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty
-        compileSegRed' (Pat red_pes) lvl segred_space [segred_op] $ \red_cont ->
+        compileSegRed' (Pat red_pes) grid segred_space [segred_op] $ \red_cont ->
           red_cont . flip map subhistos $ \subhisto ->
             ( Var subhisto,
               map Imp.le64 $
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs
@@ -23,10 +23,11 @@
   KernelBody GPUMem ->
   CallKernelGen ()
 compileSegMap pat lvl space kbody = do
+  attrs <- lvlKernelAttrs lvl
+
   let (is, dims) = unzip $ unSegSpace space
       dims' = map pe64 dims
-      group_size' = pe64 <$> segGroupSize lvl
-      attrs = defKernelAttrs (segNumGroups lvl) (segGroupSize lvl)
+      group_size' = pe64 <$> kAttrGroupSize attrs
 
   emit $ Imp.DebugPrint "\n# SegMap" Nothing
   case lvl of
@@ -58,4 +59,6 @@
             compileStms mempty (kernelBodyStms kbody) $
               zipWithM_ (compileGroupResult space) (patElems pat) $
                 kernelBodyResult kbody
+    SegThreadInGroup {} ->
+      error "compileSegMap: SegThreadInGroup"
   emit $ Imp.DebugPrint "" 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
@@ -82,8 +82,10 @@
   [SegBinOp GPUMem] ->
   KernelBody GPUMem ->
   CallKernelGen ()
-compileSegRed pat lvl space reds body =
-  compileSegRed' pat lvl space reds $ \red_cont ->
+compileSegRed pat lvl space reds body = do
+  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl
+  let grid = KernelGrid num_groups group_size
+  compileSegRed' pat grid space reds $ \red_cont ->
     compileStms mempty (kernelBodyStms body) $ do
       let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body
 
@@ -96,12 +98,12 @@
 -- | Like 'compileSegRed', but where the body is a monadic action.
 compileSegRed' ::
   Pat LetDecMem ->
-  SegLevel ->
+  KernelGrid ->
   SegSpace ->
   [SegBinOp GPUMem] ->
   DoSegBody ->
   CallKernelGen ()
-compileSegRed' pat lvl space reds body
+compileSegRed' pat grid space reds body
   | genericLength reds > maxNumOps =
       compilerLimitationS $
         "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."
@@ -116,8 +118,8 @@
         (smallSegmentsReduction pat num_groups group_size space reds body)
         (largeSegmentsReduction pat num_groups group_size space reds body)
   where
-    num_groups = segNumGroups lvl
-    group_size = segGroupSize lvl
+    num_groups = gridNumGroups grid
+    group_size = gridGroupSize grid
 
 -- | Prepare intermediate arrays for the reduction.  Prim-typed
 -- arguments go in local memory (so we need to do the allocation of
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs
@@ -215,6 +215,7 @@
   KernelBody GPUMem ->
   CallKernelGen ()
 compileSegScan pat lvl space scanOp kbody = do
+  attrs <- lvlKernelAttrs lvl
   let Pat all_pes = pat
       scanOpNe = segBinOpNeutral scanOp
       tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scanOp
@@ -233,7 +234,7 @@
       mem_constraint = max k_mem sumT `div` maxT
       reg_constraint = (k_reg - 1 - sumT') `div` (2 * sumT')
 
-      group_size = segGroupSize lvl
+      group_size = kAttrGroupSize attrs
       group_size' = pe64 $ unCount group_size
 
   num_groups <-
@@ -274,7 +275,7 @@
     constants <- kernelConstants <$> askEnv
 
     (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <-
-      createLocalArrays (segGroupSize lvl) (intConst Int64 m) tys
+      createLocalArrays (kAttrGroupSize attrs) (intConst Int64 m) tys
 
     dynamicId <- dPrim "dynamic_id" int32
     sWhen (kernelLocalThreadId constants .==. 0) $ do
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
@@ -491,6 +491,8 @@
   KernelBody GPUMem ->
   CallKernelGen ()
 compileSegScan pat lvl space scans kbody = do
+  attrs <- lvlKernelAttrs lvl
+
   -- Since stage 2 involves a group size equal to the number of groups
   -- used for stage 1, we have to cap this number to the maximum group
   -- size.
@@ -501,14 +503,13 @@
     fmap (Imp.Count . tvSize) $
       dPrimV "stage1_num_groups" $
         sMin64 (tvExp stage1_max_num_groups) $
-          pe64 $
-            Imp.unCount $
-              segNumGroups lvl
+          pe64 . Imp.unCount . kAttrNumGroups $
+            attrs
 
   (stage1_num_threads, elems_per_group, crossesSegment) <-
-    scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody
+    scanStage1 pat stage1_num_groups (kAttrGroupSize attrs) space scans kbody
 
   emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group
 
   scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans
-  scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans
+  scanStage3 pat (kAttrNumGroups attrs) (kAttrGroupSize attrs) elems_per_group crossesSegment space scans
diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs
--- a/src/Futhark/Construct.hs
+++ b/src/Futhark/Construct.hs
@@ -86,6 +86,8 @@
     eAny,
     eDimInBounds,
     eOutOfBounds,
+    eIndex,
+    eLast,
 
     -- * Other building blocks
     asIntZ,
@@ -413,6 +415,22 @@
           BasicOp $
             BinOp LogOr less_than_zero greater_than_size
   foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is'
+
+-- | The array element at this index.
+eIndex :: MonadBuilder m => VName -> m (Exp (Rep m)) -> m (Exp (Rep m))
+eIndex arr i = do
+  i' <- letSubExp "i" =<< i
+  arr_t <- lookupType arr
+  pure $ BasicOp $ Index arr $ fullSlice arr_t [DimFix i']
+
+-- | The last element of the given array.
+eLast :: MonadBuilder m => VName -> m (Exp (Rep m))
+eLast arr = do
+  n <- arraySize 0 <$> lookupType arr
+  nm1 <-
+    letSubExp "nm1" . BasicOp $
+      BinOp (Sub Int64 OverflowUndef) n (intConst Int64 1)
+  eIndex arr (eSubExp nm1)
 
 -- | Construct an unspecified value of the given type.
 eBlank :: MonadBuilder m => Type -> m (Exp (Rep m))
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
@@ -12,6 +12,10 @@
 
     -- * SegOp refinements
     SegLevel (..),
+    segVirt,
+    SegVirt (..),
+    SegSeqDims (..),
+    KernelGrid (..),
 
     -- * Reexports
     module Futhark.IR.GPU.Sizes,
@@ -44,68 +48,118 @@
   )
 import Futhark.Util.Pretty qualified as PP
 
+-- | These dimensions (indexed from 0, outermost) of the corresponding
+-- 'SegSpace' should not be parallelised, but instead iterated
+-- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a
+-- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop
+-- with @n@ iterations, while the @m@ dimension will be parallelised.
+--
+-- Semantically, this has no effect, but it may allow reductions in
+-- memory usage or other low-level optimisations.  Operationally, the
+-- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads
+-- running at any given moment will always have the same indexes along
+-- the dimensions specified by @[i,j,k]@.
+--
+-- At the moment, this is only supported for 'SegNoVirtFull'
+-- intra-group parallelism in GPU code, as we have not yet found it
+-- useful anywhere else.
+newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}
+  deriving (Eq, Ord, Show)
+
+-- | Do we need group-virtualisation when generating code for the
+-- segmented operation?  In most cases, we do, but for some simple
+-- kernels, we compute the full number of groups in advance, and then
+-- virtualisation is an unnecessary (but generally very small)
+-- overhead.  This only really matters for fairly trivial but very
+-- wide @map@ kernels where each thread performs constant-time work on
+-- scalars.
+data SegVirt
+  = SegVirt
+  | SegNoVirt
+  | -- | Not only do we not need virtualisation, but we _guarantee_
+    -- that all physical threads participate in the work.  This can
+    -- save some checks in code generation.
+    SegNoVirtFull SegSeqDims
+  deriving (Eq, Ord, Show)
+
+-- | The actual, physical grid dimensions used for the GPU kernel
+-- running this 'SegOp'.
+data KernelGrid = KernelGrid
+  { gridNumGroups :: Count NumGroups SubExp,
+    gridGroupSize :: Count GroupSize SubExp
+  }
+  deriving (Eq, Ord, Show)
+
 -- | At which level the *body* of a t'SegOp' executes.
 data SegLevel
-  = SegThread
-      { segNumGroups :: Count NumGroups SubExp,
-        segGroupSize :: Count GroupSize SubExp,
-        segVirt :: SegVirt
-      }
-  | SegGroup
-      { segNumGroups :: Count NumGroups SubExp,
-        segGroupSize :: Count GroupSize SubExp,
-        segVirt :: SegVirt
-      }
+  = SegThread SegVirt (Maybe KernelGrid)
+  | SegGroup SegVirt (Maybe KernelGrid)
+  | SegThreadInGroup SegVirt
   deriving (Eq, Ord, Show)
 
+-- | The 'SegVirt' of the 'SegLevel'.
+segVirt :: SegLevel -> SegVirt
+segVirt (SegThread v _) = v
+segVirt (SegGroup v _) = v
+segVirt (SegThreadInGroup v) = v
+
+instance PP.Pretty SegVirt where
+  pretty SegNoVirt = mempty
+  pretty (SegNoVirtFull dims) = "full" <+> pretty (segSeqDims dims)
+  pretty SegVirt = "virtualise"
+
+instance PP.Pretty KernelGrid where
+  pretty (KernelGrid num_groups group_size) =
+    "groups=" <> pretty num_groups <> PP.semi
+      <+> "groupsize=" <> pretty group_size
+
 instance PP.Pretty SegLevel where
-  pretty lvl =
-    PP.parens
-      ( lvl' <> PP.semi
-          <+> "#groups=" <> pretty (segNumGroups lvl) <> PP.semi
-          <+> "groupsize=" <> pretty (segGroupSize lvl) <> virt
-      )
-    where
-      lvl' = case lvl of
-        SegThread {} -> "thread"
-        SegGroup {} -> "group"
-      virt = case segVirt lvl of
-        SegNoVirt -> mempty
-        SegNoVirtFull dims -> PP.semi <+> "full" <+> pretty (segSeqDims dims)
-        SegVirt -> PP.semi <+> "virtualise"
+  pretty (SegThread virt grid) =
+    PP.parens ("thread" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)
+  pretty (SegGroup virt grid) =
+    PP.parens ("group" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)
+  pretty (SegThreadInGroup virt) =
+    PP.parens ("ingroup" <> PP.semi <+> pretty virt)
 
-instance Engine.Simplifiable SegLevel where
-  simplify (SegThread num_groups group_size virt) =
-    SegThread
-      <$> traverse Engine.simplify num_groups
-      <*> traverse Engine.simplify group_size
-      <*> pure virt
-  simplify (SegGroup num_groups group_size virt) =
-    SegGroup
+instance Engine.Simplifiable KernelGrid where
+  simplify (KernelGrid num_groups group_size) =
+    KernelGrid
       <$> traverse Engine.simplify num_groups
       <*> traverse Engine.simplify group_size
-      <*> pure virt
 
-instance Substitute SegLevel where
-  substituteNames substs (SegThread num_groups group_size virt) =
-    SegThread
-      (substituteNames substs num_groups)
-      (substituteNames substs group_size)
-      virt
-  substituteNames substs (SegGroup num_groups group_size virt) =
-    SegGroup
+instance Engine.Simplifiable SegLevel where
+  simplify (SegThread virt grid) =
+    SegThread virt <$> Engine.simplify grid
+  simplify (SegGroup virt grid) =
+    SegGroup virt <$> Engine.simplify grid
+  simplify (SegThreadInGroup virt) =
+    pure $ SegThreadInGroup virt
+
+instance Substitute KernelGrid where
+  substituteNames substs (KernelGrid num_groups group_size) =
+    KernelGrid
       (substituteNames substs num_groups)
       (substituteNames substs group_size)
-      virt
 
+instance Substitute SegLevel where
+  substituteNames substs (SegThread virt grid) =
+    SegThread virt (substituteNames substs grid)
+  substituteNames substs (SegGroup virt grid) =
+    SegGroup virt (substituteNames substs grid)
+  substituteNames _ (SegThreadInGroup virt) =
+    SegThreadInGroup virt
+
 instance Rename SegLevel where
   rename = substituteRename
 
+instance FreeIn KernelGrid where
+  freeIn' (KernelGrid num_groups group_size) =
+    freeIn' (num_groups, group_size)
+
 instance FreeIn SegLevel where
-  freeIn' (SegThread num_groups group_size _) =
-    freeIn' num_groups <> freeIn' group_size
-  freeIn' (SegGroup num_groups group_size _) =
-    freeIn' num_groups <> freeIn' group_size
+  freeIn' (SegThread _virt grid) = freeIn' grid
+  freeIn' (SegGroup _virt grid) = freeIn' grid
+  freeIn' (SegThreadInGroup _virt) = mempty
 
 -- | A simple size-level query or computation.
 data SizeOp
@@ -305,22 +359,34 @@
   opMetrics (SizeOp op) = opMetrics op
   opMetrics (GPUBody _ body) = inside "GPUBody" $ bodyMetrics body
 
+checkGrid :: TC.Checkable rep => KernelGrid -> TC.TypeM rep ()
+checkGrid grid = do
+  TC.require [Prim int64] $ unCount $ gridNumGroups grid
+  TC.require [Prim int64] $ unCount $ gridGroupSize grid
+
 checkSegLevel ::
   TC.Checkable rep =>
   Maybe SegLevel ->
   SegLevel ->
   TC.TypeM rep ()
-checkSegLevel Nothing lvl = do
-  TC.require [Prim int64] $ unCount $ segNumGroups lvl
-  TC.require [Prim int64] $ unCount $ segGroupSize lvl
+checkSegLevel (Just SegGroup {}) (SegThreadInGroup _virt) =
+  pure ()
+checkSegLevel _ (SegThreadInGroup _virt) =
+  TC.bad $ TC.TypeError "ingroup SegOp not in group SegOp."
 checkSegLevel (Just SegThread {}) _ =
   TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level."
-checkSegLevel (Just x) y
-  | x == y = TC.bad $ TC.TypeError $ "Already at at level " <> prettyText x
-  | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y =
-      TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel."
-  | otherwise =
-      pure ()
+checkSegLevel (Just SegThreadInGroup {}) _ =
+  TC.bad $ TC.TypeError "SegOps cannot occur when already at ingroup level."
+checkSegLevel _ (SegThread _virt Nothing) =
+  pure ()
+checkSegLevel (Just _) SegThread {} =
+  TC.bad $ TC.TypeError "thread-level SegOp cannot be nested"
+checkSegLevel Nothing (SegThread _virt grid) =
+  mapM_ checkGrid grid
+checkSegLevel (Just _) SegGroup {} =
+  TC.bad $ TC.TypeError "group-level SegOp cannot be nested"
+checkSegLevel Nothing (SegGroup _virt grid) =
+  mapM_ checkGrid grid
 
 typeCheckHostOp ::
   TC.Checkable rep =>
diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs
--- a/src/Futhark/IR/MC/Op.hs
+++ b/src/Futhark/IR/MC/Op.hs
@@ -28,8 +28,7 @@
 import Futhark.Transform.Rename
 import Futhark.Transform.Substitute
 import Futhark.Util.Pretty
-  ( Pretty,
-    nestedBlock,
+  ( nestedBlock,
     pretty,
     (<+>),
     (</>),
diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs
--- a/src/Futhark/IR/Mem.hs
+++ b/src/Futhark/IR/Mem.hs
@@ -131,7 +131,7 @@
 import Futhark.Transform.Rename
 import Futhark.Transform.Substitute
 import Futhark.Util
-import Futhark.Util.Pretty (docText, indent, ppTuple', pretty, (<+>), (</>))
+import Futhark.Util.Pretty (docText, indent, ppTupleLines', pretty, (<+>), (</>))
 import Futhark.Util.Pretty qualified as PP
 import Prelude hiding (id, (.))
 
@@ -748,9 +748,9 @@
       bad s =
         TC.bad . TC.TypeError . docText $
           "Return type"
-            </> indent 2 (ppTuple' $ map pretty rettype)
+            </> indent 2 (ppTupleLines' $ map pretty rettype)
             </> "cannot match returns of results"
-            </> indent 2 (ppTuple' $ map pretty ts)
+            </> indent 2 (ppTupleLines' $ map pretty ts)
             </> pretty s
 
   either bad pure =<< runExceptT (zipWithM_ checkReturn rettype ts)
@@ -772,9 +772,9 @@
 
   unless ok . TC.bad . TC.TypeError . docText $
     "Expression type:"
-      </> indent 2 (ppTuple' $ map pretty rt)
+      </> indent 2 (ppTupleLines' $ map pretty rt)
       </> "cannot match pattern type:"
-      </> indent 2 (ppTuple' $ map pretty val_ts)
+      </> indent 2 (ppTupleLines' $ map pretty val_ts)
   where
     matches _ _ (MemPrim x) (MemPrim y) = x == y
     matches _ _ (MemMem x_space) (MemMem y_space) =
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
@@ -49,7 +49,7 @@
   lexeme . fmap nameFromString $
     (:) <$> satisfy leading <*> many (satisfy constituent)
   where
-    leading c = isAlpha c || c `elem` ("_+-*/%=!<>|&^." :: String)
+    leading c = isAlpha c || c `elem` ("_+-*/%=!<>|&^.#" :: String)
 
 pVName :: Parser VName
 pVName = lexeme $ do
@@ -910,23 +910,36 @@
 
 pSegLevel :: Parser GPU.SegLevel
 pSegLevel =
-  parens $
-    choice
-      [ keyword "thread" $> GPU.SegThread,
-        keyword "group" $> GPU.SegGroup
-      ]
-      <*> (pSemi *> lexeme "#groups=" $> GPU.Count <*> pSubExp)
-      <*> (pSemi *> lexeme "groupsize=" $> GPU.Count <*> pSubExp)
-      <*> choice
-        [ pSemi
-            *> choice
-              [ keyword "full"
-                  $> SegOp.SegNoVirtFull
-                  <*> (SegOp.SegSeqDims <$> brackets (pInt `sepBy` pComma)),
-                keyword "virtualise" $> SegOp.SegVirt
-              ],
-          pure SegOp.SegNoVirt
+  parens . choice $
+    [ "thread"
+        $> GPU.SegThread
+        <* pSemi
+        <*> pSegVirt
+        <* pSemi
+        <*> optional pKernelGrid,
+      "group"
+        $> GPU.SegGroup
+        <* pSemi
+        <*> pSegVirt
+        <* pSemi
+        <*> optional pKernelGrid,
+      "ingroup" $> GPU.SegThreadInGroup <* pSemi <*> pSegVirt
+    ]
+  where
+    pSegVirt =
+      choice
+        [ choice
+            [ keyword "full"
+                $> GPU.SegNoVirtFull
+                <*> (GPU.SegSeqDims <$> brackets (pInt `sepBy` pComma)),
+              keyword "virtualise" $> GPU.SegVirt
+            ],
+          pure GPU.SegNoVirt
         ]
+    pKernelGrid =
+      GPU.KernelGrid
+        <$> (lexeme "groups=" $> GPU.Count <*> pSubExp <* pSemi)
+        <*> (lexeme "groupsize=" $> GPU.Count <*> pSubExp)
 
 pHostOp :: PR rep -> Parser op -> Parser (GPU.HostOp rep op)
 pHostOp pr pOther =
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
@@ -278,7 +278,7 @@
       <+> "else"
       <+> maybeNest f
       </> colon
-      <+> ppTuple' (map pretty ret)
+      <+> ppTupleLines' (map pretty ret)
     where
       info' = case ifsort of
         MatchNormal -> mempty
@@ -291,7 +291,7 @@
       <+> "->"
       <+> maybeNest defb
       </> colon
-      <+> ppTuple' (map pretty ret)
+      <+> ppTupleLines' (map pretty ret)
     where
       info' = case ifsort of
         MatchNormal -> mempty
@@ -359,7 +359,7 @@
   pretty (Lambda params body rettype) =
     "\\"
       <+> ppTuple' (map pretty params)
-      </> indent 2 (colon <+> ppTupleLines' rettype <+> "->")
+      </> indent 2 (colon <+> ppTupleLines' (map pretty rettype) <+> "->")
       </> indent 2 (pretty body)
 
 instance Pretty Signedness where
@@ -385,8 +385,8 @@
     annot (attrAnnots attrs) $
       fun
         </> indent 2 (pretty (nameToString name))
-        <+> apply (map pretty fparams)
-        </> indent 2 (colon <+> align (ppTuple' $ map pretty rettype))
+        <+> parens (commastack $ map pretty fparams)
+        </> indent 2 (colon <+> align (ppTupleLines' $ map pretty rettype))
         <+> equals
         <+> nestedBlock "{" "}" (pretty body)
     where
@@ -396,8 +396,8 @@
           "entry"
             <> (parens . align)
               ( "\"" <> pretty p_name <> "\"" <> comma
-                  </> ppTupleLines' p_entry <> comma
-                  </> ppTupleLines' ret_entry
+                  </> ppTupleLines' (map pretty p_entry) <> comma
+                  </> ppTupleLines' (map pretty ret_entry)
               )
 
 instance Pretty OpaqueType where
@@ -420,7 +420,3 @@
 instance Pretty d => Pretty (DimIndex d) where
   pretty (DimFix i) = pretty i
   pretty (DimSlice i n s) = pretty i <+> ":+" <+> pretty n <+> "*" <+> pretty s
-
--- | Like 'prettyTupleLines', but produces a 'Doc'.
-ppTupleLines' :: Pretty a => [a] -> Doc b
-ppTupleLines' = braces . stack . punctuate comma . map (align . pretty)
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
@@ -54,7 +54,6 @@
 import Futhark.Transform.Rename (Rename, Renameable)
 import Futhark.Transform.Substitute (Substitutable, Substitute)
 import Futhark.Util (maybeNth)
-import Futhark.Util.Pretty
 
 -- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'.
 isBuiltInFunction :: Name -> Bool
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
@@ -47,6 +47,7 @@
 import Data.IntMap.Strict qualified as IM
 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.Scope
 import Futhark.IR.Syntax
@@ -213,6 +214,9 @@
   freeIn' (a, b, c, d) = freeIn' a <> freeIn' b <> freeIn' c <> freeIn' d
 
 instance FreeIn a => FreeIn [a] where
+  freeIn' = foldMap freeIn'
+
+instance FreeIn a => FreeIn (S.Set a) where
   freeIn' = foldMap freeIn'
 
 instance
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
@@ -69,7 +69,7 @@
 import Futhark.Transform.Rename
 import Futhark.Transform.Substitute
 import Futhark.Util (chunks, maybeNth)
-import Futhark.Util.Pretty (Doc, Pretty, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>))
+import Futhark.Util.Pretty (Doc, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>))
 import Futhark.Util.Pretty qualified as PP
 import Prelude hiding (id, (.))
 
diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs
--- a/src/Futhark/IR/SOACS/Simplify.hs
+++ b/src/Futhark/IR/SOACS/Simplify.hs
@@ -83,12 +83,12 @@
   Simplify.SimplifiableRep rep =>
   Simplify.SimplifyOp rep (SOAC (Wise rep))
 simplifySOAC (VJP lam arr vec) = do
-  (lam', hoisted) <- Engine.simplifyLambda lam
+  (lam', hoisted) <- Engine.simplifyLambda mempty lam
   arr' <- mapM Engine.simplify arr
   vec' <- mapM Engine.simplify vec
   pure (VJP lam' arr' vec', hoisted)
 simplifySOAC (JVP lam arr vec) = do
-  (lam', hoisted) <- Engine.simplifyLambda lam
+  (lam', hoisted) <- Engine.simplifyLambda mempty lam
   arr' <- mapM Engine.simplify arr
   vec' <- mapM Engine.simplify vec
   pure (JVP lam' arr' vec', hoisted)
@@ -96,11 +96,11 @@
   outerdim' <- Engine.simplify outerdim
   nes' <- mapM Engine.simplify nes
   arr' <- mapM Engine.simplify arr
-  (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam
+  (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam
   pure (Stream outerdim' arr' nes' lam', lam_hoisted)
 simplifySOAC (Scatter w ivs lam as) = do
   w' <- Engine.simplify w
-  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam
+  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam
   ivs' <- mapM Engine.simplify ivs
   as' <- mapM Engine.simplify as
   pure (Scatter w' ivs' lam' as', hoisted)
@@ -112,25 +112,25 @@
       rf' <- Engine.simplify rf
       dests' <- Engine.simplify dests
       nes' <- mapM Engine.simplify nes
-      (op', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda op
+      (op', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty op
       pure (HistOp dests_w' rf' dests' nes' op', hoisted)
   imgs' <- mapM Engine.simplify imgs
-  (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda bfun
+  (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty bfun
   pure (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted)
 simplifySOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do
   (scans', scans_hoisted) <- fmap unzip $
     forM scans $ \(Scan lam nes) -> do
-      (lam', hoisted) <- Engine.simplifyLambda lam
+      (lam', hoisted) <- Engine.simplifyLambda mempty lam
       nes' <- Engine.simplify nes
       pure (Scan lam' nes', hoisted)
 
   (reds', reds_hoisted) <- fmap unzip $
     forM reds $ \(Reduce comm lam nes) -> do
-      (lam', hoisted) <- Engine.simplifyLambda lam
+      (lam', hoisted) <- Engine.simplifyLambda mempty lam
       nes' <- Engine.simplify nes
       pure (Reduce comm lam' nes', hoisted)
 
-  (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda map_lam
+  (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty map_lam
 
   (,)
     <$> ( Screma
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
@@ -6,8 +6,6 @@
 -- over @iota@s (so there will be explicit indexing inside them).
 module Futhark.IR.SegOp
   ( SegOp (..),
-    SegVirt (..),
-    SegSeqDims (..),
     segLevel,
     segBody,
     segSpace,
@@ -88,7 +86,6 @@
 import Futhark.Util (chunks, maybeNth)
 import Futhark.Util.Pretty
   ( Doc,
-    Pretty,
     apply,
     hsep,
     parens,
@@ -436,40 +433,6 @@
       onDim (dim, blk_tile, reg_tile) =
         pretty dim <+> "/" <+> parens (pretty blk_tile <+> "*" <+> pretty reg_tile)
 
--- | These dimensions (indexed from 0, outermost) of the corresponding
--- 'SegSpace' should not be parallelised, but instead iterated
--- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a
--- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop
--- with @n@ iterations, while the @m@ dimension will be parallelised.
---
--- Semantically, this has no effect, but it may allow reductions in
--- memory usage or other low-level optimisations.  Operationally, the
--- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads
--- running at any given moment will always have the same indexes along
--- the dimensions specified by @[i,j,k]@.
---
--- At the moment, this is only supported for 'SegNoVirtFull'
--- intra-group parallelism in GPU code, as we have not yet found it
--- useful anywhere else.
-newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}
-  deriving (Eq, Ord, Show)
-
--- | Do we need group-virtualisation when generating code for the
--- segmented operation?  In most cases, we do, but for some simple
--- kernels, we compute the full number of groups in advance, and then
--- virtualisation is an unnecessary (but generally very small)
--- overhead.  This only really matters for fairly trivial but very
--- wide @map@ kernels where each thread performs constant-time work on
--- scalars.
-data SegVirt
-  = SegVirt
-  | SegNoVirt
-  | -- | Not only do we not need virtualisation, but we _guarantee_
-    -- that all physical threads participate in the work.  This can
-    -- save some checks in code generation.
-    SegNoVirtFull SegSeqDims
-  deriving (Eq, Ord, Show)
-
 -- | Index space of a 'SegOp'.
 data SegSpace = SegSpace
   { -- | Flat physical index corresponding to the
@@ -1134,9 +1097,10 @@
 
 simplifyLambda ::
   Engine.SimplifiableRep rep =>
+  Names ->
   Lambda (Wise rep) ->
   Engine.SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-simplifyLambda = Engine.blockMigrated . Engine.simplifyLambda
+simplifyLambda bound = Engine.blockMigrated . Engine.simplifyLambda bound
 
 segSpaceSymbolTable :: ASTRep rep => SegSpace -> ST.SymbolTable rep
 segSpaceSymbolTable (SegSpace flat gtids_and_dims) =
@@ -1146,12 +1110,13 @@
 
 simplifySegBinOp ::
   Engine.SimplifiableRep rep =>
+  VName ->
   SegBinOp (Wise rep) ->
   Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep))
-simplifySegBinOp (SegBinOp comm lam nes shape) = do
+simplifySegBinOp phys_id (SegBinOp comm lam nes shape) = do
   (lam', hoisted) <-
     Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $
-      simplifyLambda lam
+      simplifyLambda (oneName phys_id) lam
   shape' <- Engine.simplify shape
   nes' <- mapM Engine.simplify nes
   pure (SegBinOp comm lam' nes' shape', hoisted)
@@ -1175,7 +1140,7 @@
   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)
   (reds', reds_hoisted) <-
     Engine.localVtable (<> scope_vtable) $
-      unzip <$> mapM simplifySegBinOp reds
+      unzip <$> mapM (simplifySegBinOp (segFlat space)) reds
   (kbody', body_hoisted) <- simplifyKernelBody space kbody
 
   pure
@@ -1189,7 +1154,7 @@
   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)
   (scans', scans_hoisted) <-
     Engine.localVtable (<> scope_vtable) $
-      unzip <$> mapM simplifySegBinOp scans
+      unzip <$> mapM (simplifySegBinOp (segFlat space)) scans
   (kbody', body_hoisted) <- simplifyKernelBody space kbody
 
   pure
@@ -1213,7 +1178,7 @@
         (lam', op_hoisted) <-
           Engine.localVtable (<> scope_vtable) $
             Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $
-              simplifyLambda lam
+              simplifyLambda (oneName (segFlat space)) lam
         pure
           ( HistOp w' rf' arrs' nes' dims' lam',
             op_hoisted
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
@@ -100,6 +100,7 @@
   ( module Language.Futhark.Core,
     prettyString,
     prettyText,
+    Pretty,
     module Futhark.IR.Rep,
     module Futhark.IR.Syntax.Core,
 
@@ -171,7 +172,7 @@
 import Data.Traversable (fmapDefault, foldMapDefault)
 import Futhark.IR.Rep
 import Futhark.IR.Syntax.Core
-import Futhark.Util.Pretty (prettyString, prettyText)
+import Futhark.Util.Pretty (Pretty, prettyString, prettyText)
 import Language.Futhark.Core
 import Prelude hiding (id, (.))
 
@@ -440,7 +441,14 @@
 
 -- | For-loop or while-loop?
 data LoopForm rep
-  = ForLoop VName IntType SubExp [(LParam rep, VName)]
+  = ForLoop
+      VName
+      -- ^ The loop iterator var
+      IntType
+      -- ^ The type of the loop iterator var
+      SubExp
+      -- ^ The number of iterations.
+      [(LParam rep, VName)]
   | WhileLoop VName
 
 deriving instance RepTypes rep => Eq (LoopForm rep)
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
@@ -57,7 +57,7 @@
 import Futhark.Construct (instantiateShapes)
 import Futhark.IR.Aliases hiding (lookupAliases)
 import Futhark.Util
-import Futhark.Util.Pretty (Pretty, align, docText, indent, ppTuple', pretty, (<+>), (</>))
+import Futhark.Util.Pretty (align, docText, indent, ppTuple', pretty, (<+>), (</>))
 
 -- | Information about an error during type checking.  The 'Show'
 -- instance for this type produces a human-readable description.
@@ -384,9 +384,9 @@
 checkConsumption (ConsumptionError e) = bad $ TypeError e
 checkConsumption (Consumption os) = pure os
 
--- | Type check two mutually control flow branches.  Think @if@.  This
--- interacts with consumption checking, as it is OK for an array to be
--- consumed in both branches.
+-- | Type check two mutually exclusive control flow branches.  Think
+-- @if@.  This interacts with consumption checking, as it is OK for an
+-- array to be consumed in both branches.
 alternative :: TypeM rep a -> TypeM rep b -> TypeM rep (a, b)
 alternative m1 m2 = do
   (x, os1) <- collectOccurences m1
@@ -394,6 +394,10 @@
   tell $ Consumption $ os1 `altOccurences` os2
   pure (x, y)
 
+alternatives :: [TypeM rep ()] -> TypeM rep ()
+alternatives [] = pure ()
+alternatives (x : xs) = void $ x `alternative` alternatives xs
+
 -- | Permit consumption of only the specified names.  If one of these
 -- names is consumed, the consumption will be rewritten to be a
 -- consumption of the corresponding alias set.  Consumption of
@@ -989,8 +993,9 @@
 checkExp (BasicOp op) = checkBasicOp op
 checkExp (Match ses cases def_case info) = do
   ses_ts <- mapM checkSubExp ses
-  mapM_ (checkCase ses_ts) cases
-  checkCaseBody def_case
+  alternatives $
+    context "in body of last case" (checkCaseBody def_case)
+      : map (checkCase ses_ts) cases
   where
     checkVal t (Just v) = Prim (primValueType v) == t
     checkVal _ Nothing = True
@@ -1002,7 +1007,9 @@
           </> "cannot match pattern"
           </> indent 2 (ppTuple' $ map pretty vs)
       context ("in body of case " <> prettyTuple vs) $ checkCaseBody body
-    checkCaseBody = matchBranchType (matchReturns info)
+    checkCaseBody body = do
+      void $ checkBody body
+      matchBranchType (matchReturns info) body
 checkExp (Apply fname args rettype_annot _) = do
   (rettype_derived, paramtypes) <- lookupFun fname $ map fst args
   argflows <- mapM (checkArg . fst) args
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
@@ -40,6 +40,7 @@
     -- holes.)
     DynamicFun (Exp, StaticVal) StaticVal
   | IntrinsicSV
+  | HoleSV SrcLoc
   deriving (Show)
 
 -- | The type is Just if this is a polymorphic binding that must be
@@ -111,6 +112,8 @@
         replaceStaticValSizes globals orig_substs sv2
     IntrinsicSV ->
       IntrinsicSV
+    HoleSV loc ->
+      HoleSV loc
   where
     tv substs =
       identityMapper
@@ -219,6 +222,7 @@
     restrict' u (DynamicFun (e, sv1) sv2) =
       DynamicFun (e, restrict' u sv1) $ restrict' u sv2
     restrict' _ IntrinsicSV = IntrinsicSV
+    restrict' _ (HoleSV loc) = HoleSV loc
     restrict'' u (Binding t sv) = Binding t $ restrict' u sv
 
 -- | Defunctionalization monad.  The Reader environment tracks both
@@ -326,6 +330,8 @@
   sizesToRename sv1 <> sizesToRename sv2
 sizesToRename IntrinsicSV =
   mempty
+sizesToRename HoleSV {} =
+  mempty
 sizesToRename Dynamic {} =
   mempty
 sizesToRename (RecordSV fs) =
@@ -498,11 +504,13 @@
     IntrinsicSV -> do
       (pats, body, tp) <- etaExpand (typeOf e) e
       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty
+    HoleSV hole_loc ->
+      pure (Hole (Info t) hole_loc, HoleSV hole_loc)
     _ ->
       let tp = typeFromSV sv
        in pure (Var qn (Info tp) loc, sv)
 defuncExp (Hole (Info t) loc) =
-  pure (Hole (Info t) loc, IntrinsicSV)
+  pure (Hole (Info t) loc, HoleSV loc)
 defuncExp (Ascript e0 tydecl loc)
   | orderZero (typeOf e0) = do
       (e0', sv) <- defuncExp e0
@@ -933,8 +941,18 @@
       pure (apply_e, sv)
     -- Propagate the 'IntrinsicsSV' until we reach the outermost application,
     -- where we construct a dynamic static value with the appropriate type.
-    IntrinsicSV
-      | depth == 0 ->
+    IntrinsicSV -> intrinsicOrHole argtypes e' sv1
+    HoleSV _ -> intrinsicOrHole argtypes e' sv1
+    _ ->
+      error $
+        "Application of an expression\n"
+          ++ prettyString e1
+          ++ "\nthat is neither a static lambda "
+          ++ "nor a dynamic function, but has static value:\n"
+          ++ show sv1
+  where
+    intrinsicOrHole argtypes e' sv
+      | depth == 0 =
           -- If the intrinsic is fully applied, then we are done.
           -- Otherwise we need to eta-expand it and recursively
           -- defunctionalise. XXX: might it be better to simply
@@ -945,14 +963,7 @@
             else do
               (pats, body, tp) <- etaExpand (typeOf e') e'
               defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty
-      | otherwise -> pure (e', IntrinsicSV)
-    _ ->
-      error $
-        "Application of an expression\n"
-          ++ prettyString e1
-          ++ "\nthat is neither a static lambda "
-          ++ "nor a dynamic function, but has static value:\n"
-          ++ show sv1
+      | otherwise = pure (e', sv)
 defuncApply depth e@(Var qn (Info t) loc) = do
   let (argtypes, _) = unfoldFunType t
   sv <- lookupVar (toStruct t) (qualLeaf qn)
@@ -1129,6 +1140,8 @@
    in Scalar $ Sum $ M.insert name svs' $ M.fromList fields
 typeFromSV IntrinsicSV =
   error "Tried to get the type from the static value of an intrinsic."
+typeFromSV HoleSV {} =
+  error "Tried to get the type from the static value of a hole."
 
 -- | Construct the type for a fully-applied dynamic function from its
 -- static value and the original types of its arguments.
diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs
--- a/src/Futhark/Internalise/Defunctorise.hs
+++ b/src/Futhark/Internalise/Defunctorise.hs
@@ -198,21 +198,18 @@
         . localScope (const f_closure) -- Start afresh.
         . generating
         $ do
-          outer_substs <- scopeSubsts <$> askScope
           abs <- asks envAbs
-          let forward (k, v) = (lookupSubst k outer_substs, v)
-              p_substs' = M.fromList $ map forward $ M.toList p_substs
-              keep k _ = k `M.member` p_substs' || k `S.member` abs
+          let keep k _ = k `M.member` p_substs || k `S.member` abs
               abs_substs =
                 M.filterWithKey keep $
-                  M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs'
+                  M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs
                     <> scopeSubsts f_closure
                     <> scopeSubsts (modScope arg_mod)
           extendScope
             ( Scope
                 abs_substs
                 ( M.singleton (modParamName f_p) $
-                    substituteInMod p_substs' arg_mod
+                    substituteInMod p_substs arg_mod
                 )
             )
             $ do
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
@@ -349,11 +349,31 @@
 
       -- Some functions are magical (overloaded) and we handle that here.
       case () of
-        -- Overloaded functions never take array arguments (except
-        -- equality, but those cannot be existential), so we can safely
-        -- ignore the existential dimensions.
         ()
-          | Just internalise <- isOverloadedFunction qfname (map fst args) loc ->
+          -- Short-circuiting operators are magical.
+          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
+            baseString (qualLeaf qfname) == "&&",
+            [(x, _), (y, _)] <- args ->
+              internaliseExp desc $
+                E.AppExp
+                  (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)
+                  (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
+          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
+            baseString (qualLeaf qfname) == "||",
+            [(x, _), (y, _)] <- args ->
+              internaliseExp desc $
+                E.AppExp
+                  (E.If x (E.Literal (E.BoolValue True) mempty) y mempty)
+                  (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
+          -- Overloaded and intrinsic functions never take array
+          -- arguments (except equality, but those cannot be
+          -- existential), so we can safely ignore the existential
+          -- dimensions.
+          | Just internalise <- isOverloadedFunction qfname desc loc -> do
+              let prepareArg (arg, _) =
+                    (E.toStruct (E.typeOf arg),) <$> internaliseExp "arg" arg
+              internalise =<< mapM prepareArg args
+          | Just internalise <- isIntrinsicFunction qfname (map fst args) loc ->
               internalise desc
           | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
             Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do
@@ -1274,6 +1294,10 @@
   E.PrimType ->
   E.PrimType ->
   InternaliseM [I.SubExp]
+internaliseBinOp _ desc E.LogAnd x y E.Bool _ =
+  simpleBinOp desc I.LogAnd x y
+internaliseBinOp _ desc E.LogOr x y E.Bool _ =
+  simpleBinOp desc I.LogOr x y
 internaliseBinOp _ desc E.Plus x y (E.Signed t) _ =
   simpleBinOp desc (I.Add t I.OverflowWrap) x y
 internaliseBinOp _ desc E.Plus x y (E.Unsigned t) _ =
@@ -1459,82 +1483,30 @@
       rettype
       =<< bodyBind body
 
--- | Some operators and functions are overloaded or otherwise special
--- - we detect and treat them here.
+-- | Overloaded operators are treated here.
 isOverloadedFunction ::
   E.QualName VName ->
-  [E.Exp] ->
+  String ->
   SrcLoc ->
-  Maybe (String -> InternaliseM [SubExp])
-isOverloadedFunction qname args loc = do
+  Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
+isOverloadedFunction qname desc loc = do
   guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag
-  let handlers =
-        [ handleSign,
-          handleIntrinsicOps,
-          handleOps,
-          handleSOACs,
-          handleAccs,
-          handleAD,
-          handleRest
-        ]
-  msum [h args $ baseString $ qualLeaf qname | h <- handlers]
+  handle $ baseString $ qualLeaf qname
   where
-    handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x
-    handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x
-    handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x
-    handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x
-    handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x
-    handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x
-    handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x
-    handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x
-    handleSign _ _ = Nothing
-
-    handleIntrinsicOps [x] s
-      | Just unop <- find ((== s) . prettyString) allUnOps = Just $ \desc -> do
-          x' <- internaliseExp1 "x" x
-          fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'
-    handleIntrinsicOps [TupLit [x, y] _] s
-      | Just bop <- find ((== s) . prettyString) allBinOps = Just $ \desc -> do
-          x' <- internaliseExp1 "x" x
-          y' <- internaliseExp1 "y" y
-          fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'
-      | Just cmp <- find ((== s) . prettyString) allCmpOps = Just $ \desc -> do
-          x' <- internaliseExp1 "x" x
-          y' <- internaliseExp1 "y" y
-          fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'
-    handleIntrinsicOps [x] s
-      | Just conv <- find ((== s) . prettyString) allConvOps = Just $ \desc -> do
-          x' <- internaliseExp1 "x" x
-          fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'
-    handleIntrinsicOps _ _ = Nothing
-
-    -- Short-circuiting operators are magical.
-    handleOps [x, y] "&&" = Just $ \desc ->
-      internaliseExp desc $
-        E.AppExp
-          (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)
-          (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
-    handleOps [x, y] "||" = Just $ \desc ->
-      internaliseExp desc $
-        E.AppExp
-          (E.If x (E.Literal (E.BoolValue True) mempty) y mempty)
-          (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
     -- Handle equality and inequality specially, to treat the case of
     -- arrays.
-    handleOps [xe, ye] op
-      | Just cmp_f <- isEqlOp op = Just $ \desc -> do
-          xe' <- internaliseExp "x" xe
-          ye' <- internaliseExp "y" ye
-          rs <- zipWithM (doComparison desc) xe' ye'
-          cmp_f desc =<< letSubExp "eq" =<< eAll rs
+    handle op
+      | Just cmp_f <- isEqlOp op = Just $ \[(_, xe'), (_, ye')] -> do
+          rs <- zipWithM doComparison xe' ye'
+          cmp_f =<< letSubExp "eq" =<< eAll rs
       where
-        isEqlOp "!=" = Just $ \desc eq ->
+        isEqlOp "!=" = Just $ \eq ->
           letTupExp' desc $ I.BasicOp $ I.UnOp I.Not eq
-        isEqlOp "==" = Just $ \_ eq ->
+        isEqlOp "==" = Just $ \eq ->
           pure [eq]
         isEqlOp _ = Nothing
 
-        doComparison desc x y = do
+        doComparison x y = do
           x_t <- I.subExpType x
           y_t <- I.subExpType y
           case x_t of
@@ -1572,15 +1544,62 @@
 
               letSubExp "arrays_equal"
                 =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])
-    handleOps [x, y] name
+    handle name
       | Just bop <- find ((name ==) . prettyString) [minBound .. maxBound :: E.BinOp] =
-          Just $ \desc -> do
-            x' <- internaliseExp1 "x" x
-            y' <- internaliseExp1 "y" y
-            case (E.typeOf x, E.typeOf y) of
+          Just $ \[(x_t, [x']), (y_t, [y'])] ->
+            case (x_t, y_t) of
               (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) ->
                 internaliseBinOp loc desc bop x' y' t1 t2
               _ -> error "Futhark.Internalise.internaliseExp: non-primitive type in BinOp."
+    handle _ = Nothing
+
+-- | Handle intrinsic functions.  These are only allowed to be called
+-- in the prelude, and their internalisation may involve inspecting
+-- the AST.
+isIntrinsicFunction ::
+  E.QualName VName ->
+  [E.Exp] ->
+  SrcLoc ->
+  Maybe (String -> InternaliseM [SubExp])
+isIntrinsicFunction qname args loc = do
+  guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag
+  let handlers =
+        [ handleSign,
+          handleOps,
+          handleSOACs,
+          handleAccs,
+          handleAD,
+          handleRest
+        ]
+  msum [h args $ baseString $ qualLeaf qname | h <- handlers]
+  where
+    handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x
+    handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x
+    handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x
+    handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x
+    handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x
+    handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x
+    handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x
+    handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x
+    handleSign _ _ = Nothing
+
+    handleOps [x] s
+      | Just unop <- find ((== s) . prettyString) allUnOps = Just $ \desc -> do
+          x' <- internaliseExp1 "x" x
+          fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'
+    handleOps [TupLit [x, y] _] s
+      | Just bop <- find ((== s) . prettyString) allBinOps = Just $ \desc -> do
+          x' <- internaliseExp1 "x" x
+          y' <- internaliseExp1 "y" y
+          fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'
+      | Just cmp <- find ((== s) . prettyString) allCmpOps = Just $ \desc -> do
+          x' <- internaliseExp1 "x" x
+          y' <- internaliseExp1 "y" y
+          fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'
+    handleOps [x] s
+      | Just conv <- find ((== s) . prettyString) allConvOps = Just $ \desc -> do
+          x' <- internaliseExp1 "x" x
+          fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'
     handleOps _ _ = Nothing
 
     handleSOACs [TupLit [lam, arr] _] "map" = Just $ \desc -> do
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
@@ -211,7 +211,7 @@
             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' segthd_lvl [r_par, tseq_div_tpar] (t_par, t_par) $
+          segScatter2D (str_A ++ "_glb2loc") loc_sz_X 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)
@@ -406,7 +406,8 @@
             (height_A, width_B, rem_outer_dims)
             code2'
 
-        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt
+        let grid = KernelGrid (Count grid_size) (Count group_size)
+            level' = SegGroup SegNoVirt (Just grid)
             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_y, gridDim_y), (gid_x, gridDim_x)])
             kbody' = KernelBody () stms_seggroup ret_seggroup
         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'
@@ -576,7 +577,8 @@
             (height_A, width_B, rem_outer_dims)
             code2'
 
-        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt
+        let grid = KernelGrid (Count grid_size) (Count group_size)
+            level' = SegGroup SegNoVirt (Just grid)
             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)])
             kbody' = KernelBody () stms_seggroup ret_seggroup
         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'
@@ -770,7 +772,7 @@
 mkNewSegthdLvl tx ty grid_pexp = do
   grid_size <- letSubExp "grid_size" =<< toExp grid_pexp
   group_size <- letSubExp "group_size" =<< toExp (pe64 ty * pe64 tx)
-  let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))
+  let segthd_lvl = SegThreadInGroup (SegNoVirtFull (SegSeqDims []))
   pure (grid_size, group_size, segthd_lvl)
 
 mkGidsXYF :: Builder GPU (VName, VName, VName)
@@ -1084,7 +1086,7 @@
         let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev
         grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp
         group_size <- letSubExp "group_size_tile3d" =<< toExp (pe64 ty * pe64 tx)
-        let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))
+        let segthd_lvl = SegThreadInGroup (SegNoVirtFull (SegSeqDims []))
 
         count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz group_size
 
@@ -1285,7 +1287,8 @@
 
           pure $ map (RegTileReturns mempty regtile_ret_dims) epilogue_res'
         -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do
-        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt
+        let grid = KernelGrid (Count grid_size) (Count group_size)
+            level' = SegGroup SegNoVirt (Just grid)
             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_z, gridDim_z), (gid_y, gridDim_y), (gid_x, gridDim_x)])
             kbody' = KernelBody () stms_seggroup ret_seggroup
 
diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs
--- a/src/Futhark/Optimise/Fusion/GraphRep.hs
+++ b/src/Futhark/Optimise/Fusion/GraphRep.hs
@@ -39,10 +39,12 @@
 
     -- * Construction
     mkDepGraph,
+    mkDepGraphForFun,
     pprg,
   )
 where
 
+import Control.Monad.Reader
 import Data.Bifunctor (bimap)
 import Data.Foldable (foldlM)
 import Data.Graph.Inductive.Dot qualified as G
@@ -330,6 +332,12 @@
         makeAliasTable (bodyStms body),
         initialGraphConstruction
       ]
+
+-- | Make a dependency graph corresponding to a function.
+mkDepGraphForFun :: FunDef SOACS -> DepGraph
+mkDepGraphForFun f = runReader (mkDepGraph (funDefBody f)) scope
+  where
+    scope = scopeOfFParams (funDefParams f) <> scopeOf (bodyStms (funDefBody f))
 
 -- | Merges two contexts.
 mergedContext :: Ord b => a -> G.Context a b -> G.Context a b -> G.Context a b
diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs
--- a/src/Futhark/Optimise/GenRedOpt.hs
+++ b/src/Futhark/Optimise/GenRedOpt.hs
@@ -89,12 +89,9 @@
         Just stms_after -> pure $ Just $ stms_before <> stms_after
         Nothing -> pure $ Just $ stms_before <> oneStm ker_snd
 
-se1 :: SubExp
-se1 = intConst Int64 1
-
 genRed2Tile2d :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU))
 genRed2Tile2d env kerstm@(Let pat_ker aux (Op (SegOp (SegMap seg_thd seg_space kres_tps old_kbody))))
-  | (SegThread _ seg_group_size _novirt) <- seg_thd,
+  | SegThread _novirt _ <- seg_thd,
     -- novirt == SegNoVirtFull || novirt == SegNoVirt,
     KernelBody () kstms kres <- old_kbody,
     Just (css, r_ses) <- allGoodReturns kres,
@@ -161,11 +158,7 @@
       gid_flat_1 <- newVName "gid_flat"
       let space1 = SegSpace gid_flat_1 gid_dims_new
 
-      (grid_size, host_stms1) <- runBuilder $ do
-        let grid_pexp = foldl (\x d -> x * pe64 d) (pe64 se1) $ map snd gid_dims_new
-        dim_prod <- letSubExp "dim_prod" =<< toExp grid_pexp
-        letSubExp "grid_size" =<< ceilDiv dim_prod (unCount seg_group_size)
-      let level1 = SegThread (Count grid_size) seg_group_size (SegNoVirtFull (SegSeqDims [])) -- novirt ?
+      let level1 = SegThread (SegNoVirtFull (SegSeqDims [])) Nothing -- novirt ?
           kbody1 = KernelBody () ker1_stms [Returns ResultMaySimplify (Certs []) k1_res]
 
       -- is it OK here to use the "aux" from the parrent kernel?
@@ -177,7 +170,7 @@
       ker2_exp <- renameExp $ Op (SegOp (SegMap seg_thd seg_space kres_tps ker2_body))
       let ker2 = Let pat_ker aux ker2_exp
       pure $
-        Just (code1_tr_host <> host_stms1 <> oneStm ker1, ker2)
+        Just (code1_tr_host <> oneStm ker1, ker2)
   where
     isIndVarToParDim _ (Constant _) _ = False
     isIndVarToParDim variance (Var acc_ind) par_dim =
@@ -198,7 +191,6 @@
               (reverse acc_inds)
        in invar_dims ++ inner_dims
     --
-    ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y
     getAccLambda acc_tp =
       case acc_tp of
         (Acc tp_id _shp el_tps _) ->
diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs
--- a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs
+++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs
@@ -1108,10 +1108,7 @@
     collectHostOp (OtherOp op) = collectFree op
     collectHostOp GPUBody {} = pure ()
 
-    collectSegLevel (SegThread (Count num) (Count size) _) =
-      collectSubExp num >> collectSubExp size
-    collectSegLevel (SegGroup (Count num) (Count size) _) =
-      collectSubExp num >> collectSubExp size
+    collectSegLevel = mapM_ captureName . namesToList . freeIn
 
     collectSegSpace space =
       mapM_ collectSubExp (segSpaceDims space)
diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs
--- a/src/Futhark/Optimise/Simplify/Engine.hs
+++ b/src/Futhark/Optimise/Simplify/Engine.hs
@@ -452,7 +452,7 @@
   UT.UsageTable ->
   SimpleM rep (Stms (Wise rep), Stm (Wise rep))
 recSimplifyStm (Let pat (StmAux cs attrs (_, dec)) e) usage = do
-  ((e', e_hoisted), e_cs) <- collectCerts $ simplifyExp usage pat e
+  ((e', e_hoisted), e_cs) <- collectCerts $ simplifyExp (usage <> UT.usageInPat pat) pat e
   let aux' = StmAux (cs <> e_cs) attrs dec
   pure (e_hoisted, mkWiseStm (removePatWisdom pat) aux' e')
 
@@ -666,8 +666,8 @@
   all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn
 
 matchBlocker ::
-  (ASTRep rep, CanBeWise (Op rep), FreeIn a) =>
-  a ->
+  (ASTRep rep, CanBeWise (Op rep)) =>
+  [SubExp] ->
   MatchDec rt ->
   SimpleM rep (BlockPred (Wise rep))
 matchBlocker cond (MatchDec _ ifsort) = do
@@ -896,7 +896,7 @@
       Nothing ->
         pure (Nothing, mempty)
       Just (op_lam, nes) -> do
-        (op_lam', op_lam_stms) <- blockMigrated (simplifyLambda op_lam)
+        (op_lam', op_lam_stms) <- blockMigrated (simplifyLambda mempty op_lam)
         nes' <- simplify nes
         pure (Just (op_lam', nes'), op_lam_stms)
     (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs)
@@ -1063,11 +1063,12 @@
 
 simplifyLambda ::
   SimplifiableRep rep =>
+  Names ->
   Lambda (Wise rep) ->
   SimpleM rep (Lambda (Wise rep), Stms (Wise rep))
-simplifyLambda lam = do
+simplifyLambda extra_bound lam = do
   par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers
-  simplifyLambdaMaybeHoist par_blocker mempty lam
+  simplifyLambdaMaybeHoist (par_blocker `orIf` hasFree extra_bound) mempty lam
 
 simplifyLambdaNoHoisting ::
   SimplifiableRep rep =>
diff --git a/src/Futhark/Optimise/Simplify/Rules/Loop.hs b/src/Futhark/Optimise/Simplify/Rules/Loop.hs
--- a/src/Futhark/Optimise/Simplify/Rules/Loop.hs
+++ b/src/Futhark/Optimise/Simplify/Rules/Loop.hs
@@ -238,34 +238,6 @@
     notIndex _ = True
 simplifyLoopVariables _ _ _ _ = Skip
 
--- If a for-loop with no loop variables has a counter of type Int64,
--- and the bound is just a constant or sign-extended integer of
--- smaller type, then change the loop to iterate over the smaller type
--- instead.  We then move the sign extension inside the loop instead.
--- This addresses loops of the form @for i in x..<y@ in the source
--- language.
-narrowLoopType :: (BuilderOps rep) => TopDownRuleDoLoop rep
-narrowLoopType vtable pat aux (merge, ForLoop i Int64 n [], body)
-  | Just (n', it', cs) <- smallerType =
-      Simplify $ do
-        i' <- newVName $ baseString i
-        let form' = ForLoop i' it' n' []
-        body' <- insertStmsM . inScopeOf form' $ do
-          letBindNames [i] $ BasicOp $ ConvOp (SExt it' Int64) (Var i')
-          pure body
-        auxing aux $ certifying cs $ letBind pat $ DoLoop merge form' body'
-  where
-    smallerType
-      | Var n' <- n,
-        Just (ConvOp (SExt it' _) n'', cs) <- ST.lookupBasicOp n' vtable =
-          Just (n'', it', cs)
-      | Constant (IntValue (Int64Value n')) <- n,
-        toInteger n' <= toInteger (maxBound :: Int32) =
-          Just (intConst Int32 (toInteger n'), Int32, mempty)
-      | otherwise =
-          Nothing
-narrowLoopType _ _ _ _ = Skip
-
 unroll ::
   BuilderOps rep =>
   Integer ->
@@ -313,8 +285,7 @@
   [ RuleDoLoop hoistLoopInvariantMergeVariables,
     RuleDoLoop simplifyClosedFormLoop,
     RuleDoLoop simplifyKnownIterationLoop,
-    RuleDoLoop simplifyLoopVariables,
-    RuleDoLoop narrowLoopType
+    RuleDoLoop simplifyLoopVariables
   ]
 
 bottomUpRules :: BuilderOps rep => [BottomUpRule rep]
diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs
--- a/src/Futhark/Optimise/TileLoops.hs
+++ b/src/Futhark/Optimise/TileLoops.hs
@@ -126,13 +126,12 @@
         not $ null $ tiledInputs inputs,
         gtid_y : gtid_x : top_gtids_rev <- reverse gtids,
         kdim_y : kdim_x : top_kdims_rev <- reverse kdims,
-        (prestms', poststms') <-
+        Just (prestms', poststms') <-
           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms),
         used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res =
           Just . injectPrelude initial_space variance prestms' used
             <$> tileGeneric
               (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev)
-              initial_lvl
               res_ts
               (stmPat stm_to_tile)
               (gtid_x, gtid_y)
@@ -148,13 +147,12 @@
         inputs <- map (is1DTileable gtid variance) arrs,
         not $ null $ tiledInputs inputs,
         gtid `notNameIn` branch_variant,
-        (prestms', poststms') <-
+        Just (prestms', poststms') <-
           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms),
         used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res =
           Just . injectPrelude initial_space variance prestms' used
             <$> tileGeneric
               (tiling1d $ reverse top_space_rev)
-              initial_lvl
               res_ts
               (stmPat stm_to_tile)
               gtid
@@ -167,7 +165,7 @@
       -- Tiling inside for-loop.
       | DoLoop merge (ForLoop i it bound []) loopbody <- stmExp stm_to_tile,
         not $ any ((`nameIn` freeIn merge) . paramName . fst) merge,
-        (prestms', poststms') <-
+        Just (prestms', poststms') <-
           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = do
           let branch_variant' =
                 branch_variant
@@ -214,15 +212,22 @@
             descend (prestms <> oneStm stm_to_tile) poststms
 
 -- | Move statements from prelude to postlude if they are not used in
--- the tiled statement anyway.
+-- the tiled statement anyway.  Also, fail if the provided Stm uses
+-- anything from the resulting prelude whose size is not free in the
+-- prelude.
 preludeToPostlude ::
   VarianceTable ->
   Stms GPU ->
   Stm GPU ->
   Stms GPU ->
-  (Stms GPU, Stms GPU)
-preludeToPostlude variance prelude stm_to_tile postlude =
-  (prelude_used, prelude_not_used <> postlude)
+  Maybe (Stms GPU, Stms GPU)
+preludeToPostlude variance prelude stm_to_tile postlude = do
+  let prelude_sizes =
+        freeIn $ foldMap (patTypes . stmPat) prelude_used
+      prelude_bound =
+        namesFromList $ foldMap (patNames . stmPat) prelude_used
+  guard $ not $ prelude_sizes `namesIntersect` prelude_bound
+  Just (prelude_used, prelude_not_used <> postlude)
   where
     used_in_tiled = freeIn stm_to_tile
 
@@ -404,7 +409,7 @@
         mergeinit' <-
           fmap (map Var) $
             certifying (stmAuxCerts aux) $
-              tilingSegMap tiling "tiled_loopinit" (scalarLevel tiling) ResultPrivate $
+              tilingSegMap tiling "tiled_loopinit" ResultPrivate $
                 \in_bounds slice ->
                   fmap varsRes $
                     protectOutOfBounds "loopinit" in_bounds merge_ts $ do
@@ -428,8 +433,7 @@
 
         loopbody' <-
           localScope (scopeOfFParams mergeparams') . runBodyBuilder $
-            resultBody . map Var
-              <$> tiledBody private' privstms'
+            resultBody . map Var <$> tiledBody private' privstms'
         accs' <-
           letTupExp "tiled_inside_loop" $
             DoLoop merge' (ForLoop i it bound []) loopbody'
@@ -447,13 +451,12 @@
 doPrelude :: Tiling -> PrivStms -> Stms GPU -> [VName] -> Builder GPU [VName]
 doPrelude tiling privstms prestms prestms_live =
   -- Create a SegMap that takes care of the prelude for every thread.
-  tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $
-    \in_bounds slice -> do
-      ts <- mapM lookupType prestms_live
-      fmap varsRes . protectOutOfBounds "pre" in_bounds ts $ do
-        addPrivStms slice privstms
-        addStms prestms
-        pure $ varsRes prestms_live
+  tilingSegMap tiling "prelude" ResultPrivate $ \in_bounds slice -> do
+    ts <- mapM lookupType prestms_live
+    fmap varsRes . protectOutOfBounds "pre" in_bounds ts $ do
+      addPrivStms slice privstms
+      addStms prestms
+      pure $ varsRes prestms_live
 
 liveSet :: FreeIn a => Stms GPU -> a -> Names
 liveSet stms after =
@@ -578,7 +581,6 @@
 data Tiling = Tiling
   { tilingSegMap ::
       String ->
-      SegLevel ->
       ResultManifest ->
       (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result) ->
       Builder GPU [VName],
@@ -604,13 +606,7 @@
   }
 
 type DoTiling gtids kdims =
-  SegLevel -> gtids -> kdims -> SubExp -> Builder GPU Tiling
-
-scalarLevel :: Tiling -> SegLevel
-scalarLevel tiling =
-  SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt
-  where
-    lvl = tilingLevel tiling
+  gtids -> kdims -> SubExp -> Builder GPU Tiling
 
 protectOutOfBounds ::
   String ->
@@ -643,7 +639,7 @@
   [Type] ->
   Builder GPU [VName]
 postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts =
-  tilingSegMap tiling "thread_res" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do
+  tilingSegMap tiling "thread_res" ResultPrivate $ \in_bounds slice -> do
     -- Read our per-thread result from the tiled loop.
     forM_ (zip (patNames pat) accs') $ \(us, everyone) -> do
       everyone_t <- lookupType everyone
@@ -664,7 +660,6 @@
 
 tileGeneric ::
   DoTiling gtids kdims ->
-  SegLevel ->
   [Type] ->
   Pat Type ->
   gtids ->
@@ -675,8 +670,8 @@
   Stms GPU ->
   Result ->
   TileM (Stms GPU, Tiling, TiledBody)
-tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form inputs poststms poststms_res = do
-  (tiling, tiling_stms) <- runBuilder $ doTiling initial_lvl gtids kdims w
+tileGeneric doTiling res_ts pat gtids kdims w form inputs poststms poststms_res = do
+  (tiling, tiling_stms) <- runBuilder $ doTiling gtids kdims w
 
   pure (tiling_stms, tiling, tiledBody tiling)
   where
@@ -690,7 +685,7 @@
 
       -- We don't use a Replicate here, because we want to enforce a
       -- scalar memory space.
-      mergeinits <- tilingSegMap tiling "mergeinit" (scalarLevel tiling) ResultPrivate $ \in_bounds slice ->
+      mergeinits <- tilingSegMap tiling "mergeinit" ResultPrivate $ \in_bounds slice ->
         -- Constant neutral elements (a common case) do not need protection from OOB.
         if freeIn red_nes == mempty
           then pure $ subExpsRes red_nes
@@ -736,10 +731,9 @@
 
 mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude
 mkReadPreludeValues prestms_live_arrs prestms_live slice =
-  fmap mconcat $
-    forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do
-      arr_t <- lookupType arr
-      letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice
+  fmap mconcat . forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do
+    arr_t <- lookupType arr
+    letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice
 
 tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Builder GPU KernelResult
 tileReturns dims_on_top dims arr = do
@@ -776,16 +770,15 @@
   SubExp ->
   VName ->
   VName ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
+  KernelGrid ->
   TileKind ->
   PrivStms ->
   SubExp ->
   [InputArray] ->
   Builder GPU [InputTile]
-readTile1D tile_size gid gtid num_groups group_size kind privstms tile_id inputs =
+readTile1D tile_size gid gtid (KernelGrid _num_groups group_size) kind privstms tile_id inputs =
   fmap (inputsToTiles inputs)
-    . segMap1D "full_tile" lvl ResultNoSimplify
+    . segMap1D "full_tile" lvl ResultNoSimplify tile_size
     $ \ltid -> do
       j <-
         letSubExp "j"
@@ -813,18 +806,17 @@
           TileFull ->
             mapM readTileElem arrs
   where
-    lvl = SegThread num_groups group_size SegNoVirt
+    lvl = SegThreadInGroup SegNoVirt
 
 processTile1D ::
   VName ->
   VName ->
   SubExp ->
   SubExp ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
+  KernelGrid ->
   ProcessTileArgs ->
   Builder GPU [VName]
-processTile1D gid gtid kdim tile_size num_groups group_size tile_args = do
+processTile1D gid gtid kdim tile_size (KernelGrid _num_groups group_size) tile_args = do
   let red_comm = processComm tile_args
       privstms = processPrivStms tile_args
       map_lam = processMapLam tile_args
@@ -833,7 +825,7 @@
       tile_id = processTileId tile_args
       accs = processAcc tile_args
 
-  segMap1D "acc" lvl ResultPrivate $ \ltid -> do
+  segMap1D "acc" lvl ResultPrivate (unCount group_size) $ \ltid -> do
     reconstructGtids1D group_size gtid gid ltid
     addPrivStms [DimFix $ Var ltid] privstms
 
@@ -857,18 +849,17 @@
           (eBody [pure $ Op $ OtherOp $ Screma tile_size tiles' form'])
           (resultBodyM thread_accs)
   where
-    lvl = SegThread num_groups group_size SegNoVirt
+    lvl = SegThreadInGroup SegNoVirt
 
 processResidualTile1D ::
   VName ->
   VName ->
   SubExp ->
   SubExp ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
+  KernelGrid ->
   ResidualTileArgs ->
   Builder GPU [VName]
-processResidualTile1D gid gtid kdim tile_size num_groups group_size args = do
+processResidualTile1D gid gtid kdim tile_size grid args = do
   -- The number of residual elements that are not covered by
   -- the whole tiles.
   residual_input <-
@@ -899,8 +890,7 @@
           tile_size
           gid
           gtid
-          num_groups
-          group_size
+          grid
           TilePartial
           privstms
           num_whole_tiles
@@ -921,61 +911,52 @@
       let tile_args =
             ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles
       resultBody . map Var
-        <$> processTile1D gid gtid kdim residual_input num_groups group_size tile_args
+        <$> processTile1D gid gtid kdim residual_input grid tile_args
 
 tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp
-tiling1d dims_on_top initial_lvl gtid kdim w = do
+tiling1d dims_on_top gtid kdim w = do
   gid <- newVName "gid"
   gid_flat <- newVName "gid_flat"
 
-  (lvl, space) <-
-    if null dims_on_top
-      then
-        pure
-          ( SegGroup (segNumGroups initial_lvl) (segGroupSize initial_lvl) $ segVirt initial_lvl,
-            SegSpace gid_flat [(gid, unCount $ segNumGroups initial_lvl)]
-          )
-      else do
-        group_size <-
-          letSubExp "computed_group_size" $
-            BasicOp $
-              BinOp (SMin Int64) (unCount (segGroupSize initial_lvl)) kdim
+  tile_size_key <- nameFromString . prettyString <$> newVName "tile_size"
+  tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeGroup
+  let group_size = tile_size
 
-        -- How many groups we need to exhaust the innermost dimension.
-        ldim <-
-          letSubExp "ldim" $
-            BasicOp $
-              BinOp (SDivUp Int64 Unsafe) kdim group_size
+  (grid, space) <- do
+    -- How many groups we need to exhaust the innermost dimension.
+    ldim <-
+      letSubExp "ldim" . BasicOp $
+        BinOp (SDivUp Int64 Unsafe) kdim group_size
 
-        num_groups <-
-          letSubExp "computed_num_groups"
-            =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top)
+    num_groups <-
+      letSubExp "computed_num_groups"
+        =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top)
 
-        pure
-          ( SegGroup (Count num_groups) (Count group_size) SegNoVirt,
-            SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]
-          )
-  let tile_size = unCount $ segGroupSize lvl
+    pure
+      ( KernelGrid (Count num_groups) (Count group_size),
+        SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]
+      )
+  let tiling_lvl = SegThreadInGroup SegNoVirt
 
   pure
     Tiling
-      { tilingSegMap = \desc lvl' manifest f -> segMap1D desc lvl' manifest $ \ltid -> do
+      { tilingSegMap = \desc manifest f -> segMap1D desc tiling_lvl manifest tile_size $ \ltid -> do
           letBindNames [gtid]
             =<< toExp (le64 gid * pe64 tile_size + le64 ltid)
           f (untyped $ le64 gtid .<. pe64 kdim) [DimFix $ Var ltid],
         tilingReadTile =
-          readTile1D tile_size gid gtid (segNumGroups lvl) (segGroupSize lvl),
+          readTile1D tile_size gid gtid grid,
         tilingProcessTile =
-          processTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl),
+          processTile1D gid gtid kdim tile_size grid,
         tilingProcessResidualTile =
-          processResidualTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl),
+          processResidualTile1D gid gtid kdim tile_size grid,
         tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)],
         tilingTileShape = Shape [tile_size],
         tilingNumWholeTiles =
           letSubExp "num_whole_tiles" $
             BasicOp $
               BinOp (SQuot Int64 Unsafe) w tile_size,
-        tilingLevel = lvl,
+        tilingLevel = SegGroup SegNoVirt (Just grid),
         tilingSpace = space
       }
 
@@ -1015,18 +996,16 @@
   (VName, VName) ->
   (VName, VName) ->
   SubExp ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
   TileKind ->
   PrivStms ->
   SubExp ->
   [InputArray] ->
   Builder GPU [InputTile]
-readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id inputs =
+readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size kind privstms tile_id inputs =
   fmap (inputsToTiles inputs)
     . segMap2D
       "full_tile"
-      (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))
+      (SegThread (SegNoVirtFull (SegSeqDims [])) Nothing)
       ResultNoSimplify
       (tile_size, tile_size)
     $ \(ltid_x, ltid_y) -> do
@@ -1088,11 +1067,9 @@
   (VName, VName) ->
   (SubExp, SubExp) ->
   SubExp ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
   ProcessTileArgs ->
   Builder GPU [VName]
-processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size num_groups group_size tile_args = do
+processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size tile_args = do
   let privstms = processPrivStms tile_args
       red_comm = processComm tile_args
       red_lam = processRedLam tile_args
@@ -1106,7 +1083,7 @@
 
   segMap2D
     "acc"
-    (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))
+    (SegThreadInGroup (SegNoVirtFull (SegSeqDims [])))
     ResultPrivate
     (tile_size, tile_size)
     $ \(ltid_x, ltid_y) -> do
@@ -1146,82 +1123,69 @@
   (VName, VName) ->
   (SubExp, SubExp) ->
   SubExp ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
   ResidualTileArgs ->
   Builder GPU [VName]
-processResidualTile2D
-  gids
-  gtids
-  kdims
-  tile_size
-  num_groups
-  group_size
-  args = do
-    -- The number of residual elements that are not covered by
-    -- the whole tiles.
-    residual_input <-
-      letSubExp "residual_input" $
-        BasicOp $
-          BinOp (SRem Int64 Unsafe) w tile_size
+processResidualTile2D gids gtids kdims tile_size args = do
+  -- The number of residual elements that are not covered by
+  -- the whole tiles.
+  residual_input <-
+    letSubExp "residual_input" $
+      BasicOp $
+        BinOp (SRem Int64 Unsafe) w tile_size
 
-    letTupExp "acc_after_residual"
-      =<< eIf
-        (toExp $ pe64 residual_input .==. 0)
-        (resultBodyM $ map Var accs)
-        (nonemptyTile residual_input)
-    where
-      privstms = residualPrivStms args
-      red_comm = residualComm args
-      red_lam = residualRedLam args
-      map_lam = residualMapLam args
-      accs = residualAcc args
-      inputs = residualInput args
-      num_whole_tiles = residualNumWholeTiles args
-      w = residualInputSize args
+  letTupExp "acc_after_residual"
+    =<< eIf
+      (toExp $ pe64 residual_input .==. 0)
+      (resultBodyM $ map Var accs)
+      (nonemptyTile residual_input)
+  where
+    privstms = residualPrivStms args
+    red_comm = residualComm args
+    red_lam = residualRedLam args
+    map_lam = residualMapLam args
+    accs = residualAcc args
+    inputs = residualInput args
+    num_whole_tiles = residualNumWholeTiles args
+    w = residualInputSize args
 
-      nonemptyTile residual_input = renameBody <=< runBodyBuilder $ do
-        -- Collectively construct a tile.  Threads that are out-of-bounds
-        -- provide a blank dummy value.
-        full_tile <-
-          readTile2D
-            kdims
-            gtids
-            gids
-            tile_size
-            num_groups
-            group_size
-            TilePartial
-            privstms
-            num_whole_tiles
-            inputs
+    nonemptyTile residual_input = renameBody <=< runBodyBuilder $ do
+      -- Collectively construct a tile.  Threads that are out-of-bounds
+      -- provide a blank dummy value.
+      full_tile <-
+        readTile2D
+          kdims
+          gtids
+          gids
+          tile_size
+          TilePartial
+          privstms
+          num_whole_tiles
+          inputs
 
-        let slice =
-              DimSlice (intConst Int64 0) residual_input (intConst Int64 1)
-        tiles <- forM full_tile $ \case
-          InputTiled perm tile' ->
-            InputTiled perm
-              <$> letExp "partial_tile" (BasicOp $ Index tile' (Slice [slice, slice]))
-          InputUntiled arr ->
-            pure $ InputUntiled arr
+      let slice =
+            DimSlice (intConst Int64 0) residual_input (intConst Int64 1)
+      tiles <- forM full_tile $ \case
+        InputTiled perm tile' ->
+          InputTiled perm
+            <$> letExp "partial_tile" (BasicOp $ Index tile' (Slice [slice, slice]))
+        InputUntiled arr ->
+          pure $ InputUntiled arr
 
-        let tile_args =
-              ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles
+      let tile_args =
+            ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles
 
-        -- Now each thread performs a traversal of the tile and
-        -- updates its accumulator.
-        resultBody . map Var
-          <$> processTile2D
-            gids
-            gtids
-            kdims
-            tile_size
-            num_groups
-            group_size
-            tile_args
+      -- Now each thread performs a traversal of the tile and
+      -- updates its accumulator.
+      resultBody . map Var
+        <$> processTile2D
+          gids
+          gtids
+          kdims
+          tile_size
+          tile_args
 
 tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)
-tiling2d dims_on_top _initial_lvl (gtid_x, gtid_y) (kdim_x, kdim_y) w = do
+tiling2d dims_on_top (gtid_x, gtid_y) (kdim_x, kdim_y) w = do
   gid_x <- newVName "gid_x"
   gid_y <- newVName "gid_y"
 
@@ -1246,24 +1210,26 @@
         (num_groups_y : map snd dims_on_top)
 
   gid_flat <- newVName "gid_flat"
-  let lvl = SegGroup (Count num_groups) (Count group_size) (SegNoVirtFull (SegSeqDims []))
+  let grid = KernelGrid (Count num_groups) (Count group_size)
+      lvl = SegGroup (SegNoVirtFull (SegSeqDims [])) (Just grid)
       space =
         SegSpace gid_flat $
           dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)]
+      tiling_lvl = SegThreadInGroup SegNoVirt
 
   pure
     Tiling
-      { tilingSegMap = \desc lvl' manifest f ->
-          segMap2D desc lvl' manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do
+      { tilingSegMap = \desc manifest f ->
+          segMap2D desc tiling_lvl manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do
             reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y)
             f
               ( untyped $
                   le64 gtid_x .<. pe64 kdim_x .&&. le64 gtid_y .<. pe64 kdim_y
               )
               [DimFix $ Var ltid_x, DimFix $ Var ltid_y],
-        tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size (segNumGroups lvl) (segGroupSize lvl),
-        tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl),
-        tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl),
+        tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size,
+        tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,
+        tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,
         tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)],
         tilingTileShape = Shape [tile_size, tile_size],
         tilingNumWholeTiles =
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
@@ -85,12 +85,13 @@
   String ->
   SegLevel ->
   ResultManifest ->
+  SubExp -> -- dim_x
   (VName -> Builder GPU Result) ->
   Builder GPU [VName]
-segMap1D desc lvl manifest f = do
+segMap1D desc lvl manifest w f = do
   ltid <- newVName "ltid"
   ltid_flat <- newVName "ltid_flat"
-  let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)]
+  let space = SegSpace ltid_flat [(ltid, w)]
 
   ((ts, res), stms) <- localScope (scopeOfSegSpace space) . runBuilder $ do
     res <- f ltid
@@ -164,12 +165,11 @@
   String ->
   SubExp ->
   VName ->
-  SegLevel -> -- lvl
   [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 lvl seq_dims (dim_x, dim_y) f = do
+segScatter2D desc arr_size 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"
@@ -178,10 +178,8 @@
   let seq_space = zip seq_is seq_dims
 
   let segspace = SegSpace ltid_flat $ seq_space ++ [(ltid_y, dim_y), (ltid_x, dim_x)]
-      lvl' =
-        SegThread
-          (segNumGroups lvl)
-          (segGroupSize lvl)
+      lvl =
+        SegThreadInGroup
           (SegNoVirtFull (SegSeqDims [0 .. length seq_dims - 1]))
 
   ((t_v, res_v, res_i), stms) <- runBuilder $ do
@@ -194,7 +192,7 @@
   let ret = WriteReturns mempty (Shape [arr_size]) updt_arr [(Slice [DimFix res_i], res_v)]
   let body = KernelBody () stms [ret]
 
-  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl' segspace [t_v] body
+  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [t_v] 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/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs
--- a/src/Futhark/Pass/ExpandAllocations.hs
+++ b/src/Futhark/Pass/ExpandAllocations.hs
@@ -110,33 +110,33 @@
 
 transformExp :: Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem)
 transformExp (Op (Inner (SegOp (SegMap lvl space ts kbody)))) = do
-  (alloc_stms, (_, kbody')) <- transformScanRed lvl space [] kbody
+  (alloc_stms, (lvl', _, kbody')) <- transformScanRed lvl space [] kbody
   pure
     ( alloc_stms,
-      Op $ Inner $ SegOp $ SegMap lvl space ts kbody'
+      Op $ Inner $ SegOp $ SegMap lvl' space ts kbody'
     )
 transformExp (Op (Inner (SegOp (SegRed lvl space reds ts kbody)))) = do
-  (alloc_stms, (lams, kbody')) <-
+  (alloc_stms, (lvl', lams, kbody')) <-
     transformScanRed lvl space (map segBinOpLambda reds) kbody
   let reds' = zipWith (\red lam -> red {segBinOpLambda = lam}) reds lams
   pure
     ( alloc_stms,
-      Op $ Inner $ SegOp $ SegRed lvl space reds' ts kbody'
+      Op $ Inner $ SegOp $ SegRed lvl' space reds' ts kbody'
     )
 transformExp (Op (Inner (SegOp (SegScan lvl space scans ts kbody)))) = do
-  (alloc_stms, (lams, kbody')) <-
+  (alloc_stms, (lvl', lams, kbody')) <-
     transformScanRed lvl space (map segBinOpLambda scans) kbody
   let scans' = zipWith (\red lam -> red {segBinOpLambda = lam}) scans lams
   pure
     ( alloc_stms,
-      Op $ Inner $ SegOp $ SegScan lvl space scans' ts kbody'
+      Op $ Inner $ SegOp $ SegScan lvl' space scans' ts kbody'
     )
 transformExp (Op (Inner (SegOp (SegHist lvl space ops ts kbody)))) = do
-  (alloc_stms, (lams', kbody')) <- transformScanRed lvl space lams kbody
+  (alloc_stms, (lvl', lams', kbody')) <- transformScanRed lvl space lams kbody
   let ops' = zipWith onOp ops lams'
   pure
     ( alloc_stms,
-      Op $ Inner $ SegOp $ SegHist lvl space ops' ts kbody'
+      Op $ Inner $ SegOp $ SegHist lvl' space ops' ts kbody'
     )
   where
     lams = map histOp ops
@@ -156,7 +156,7 @@
       let -- XXX: fake a SegLevel, which we don't have here.  We will not
           -- use it for anything, as we will not allow irregular
           -- allocations inside the update function.
-          lvl = SegThread (Count $ intConst Int64 0) (Count $ intConst Int64 0) SegNoVirt
+          lvl = SegThread SegNoVirt Nothing
           (op_lam', lam_allocs) =
             extractLambdaAllocations (lvl, [0]) bound_outside mempty op_lam
           variantAlloc (_, Var v, _) = v `notNameIn` bound_outside
@@ -186,12 +186,33 @@
 transformExp e =
   pure (mempty, e)
 
+ensureGridKnown :: SegLevel -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)
+ensureGridKnown lvl =
+  case lvl of
+    SegThread _ (Just grid) -> pure (mempty, lvl, grid)
+    SegGroup _ (Just grid) -> pure (mempty, lvl, grid)
+    SegThread virt Nothing -> mkGrid (SegThread virt)
+    SegGroup virt Nothing -> mkGrid (SegGroup virt)
+    SegThreadInGroup {} -> error "ensureGridKnown: SegThreadInGroup"
+  where
+    mkGrid f = do
+      (grid, stms) <-
+        runBuilder $
+          KernelGrid
+            <$> (Count <$> getSize "num_groups" SizeNumGroups)
+            <*> (Count <$> getSize "group_size" SizeGroup)
+      pure (stms, f $ Just grid, grid)
+
+    getSize desc size_class = do
+      size_key <- nameFromString . prettyString <$> newVName desc
+      letSubExp desc $ Op $ Inner $ SizeOp $ GetSize size_key size_class
+
 transformScanRed ::
   SegLevel ->
   SegSpace ->
   [Lambda GPUMem] ->
   KernelBody GPUMem ->
-  ExpandM (Stms GPUMem, ([Lambda GPUMem], KernelBody GPUMem))
+  ExpandM (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem))
 transformScanRed lvl space ops kbody = do
   bound_outside <- asks $ namesFromList . M.keys
   let user = (lvl, [le64 $ segFlat space])
@@ -221,10 +242,14 @@
     _ ->
       pure ()
 
-  allocsForBody variant_allocs invariant_allocs lvl space kbody' $ \alloc_stms kbody'' -> do
-    ops'' <- forM ops' $ \op' ->
-      localScope (scopeOf op') $ offsetMemoryInLambda op'
-    pure (alloc_stms, (ops'', kbody''))
+  if null variant_allocs && null invariant_allocs
+    then pure (mempty, (lvl, ops, kbody))
+    else do
+      (lvl_stms, lvl', grid) <- ensureGridKnown lvl
+      allocsForBody variant_allocs invariant_allocs grid space kbody' $ \alloc_stms kbody'' -> do
+        ops'' <- forM ops' $ \op' ->
+          localScope (scopeOf op') $ offsetMemoryInLambda op'
+        pure (lvl_stms <> alloc_stms, (lvl', ops'', kbody''))
   where
     bound_in_kernel =
       namesFromList (M.keys $ scopeOfSegSpace space)
@@ -236,15 +261,15 @@
 allocsForBody ::
   Extraction ->
   Extraction ->
-  SegLevel ->
+  KernelGrid ->
   SegSpace ->
   KernelBody GPUMem ->
   (Stms GPUMem -> KernelBody GPUMem -> OffsetM b) ->
   ExpandM b
-allocsForBody variant_allocs invariant_allocs lvl space kbody' m = do
+allocsForBody variant_allocs invariant_allocs grid space kbody' m = do
   (alloc_offsets, alloc_stms) <-
     memoryRequirements
-      lvl
+      grid
       space
       (kernelBodyStms kbody')
       variant_allocs
@@ -258,26 +283,26 @@
       m alloc_stms kbody''
 
 memoryRequirements ::
-  SegLevel ->
+  KernelGrid ->
   SegSpace ->
   Stms GPUMem ->
   Extraction ->
   Extraction ->
   ExpandM (RebaseMap, Stms GPUMem)
-memoryRequirements lvl space kstms variant_allocs invariant_allocs = do
+memoryRequirements grid space kstms variant_allocs invariant_allocs = do
   (num_threads, num_threads_stms) <-
     runBuilder . letSubExp "num_threads" . BasicOp $
       BinOp
         (Mul Int64 OverflowUndef)
-        (unCount $ segNumGroups lvl)
-        (unCount $ segGroupSize lvl)
+        (unCount $ gridNumGroups grid)
+        (unCount $ gridGroupSize grid)
 
   (invariant_alloc_stms, invariant_alloc_offsets) <-
     inScopeOf num_threads_stms $
       expandedInvariantAllocations
         num_threads
-        (segNumGroups lvl)
-        (segGroupSize lvl)
+        (gridNumGroups grid)
+        (gridGroupSize grid)
         invariant_allocs
 
   (variant_alloc_stms, variant_alloc_offsets) <-
@@ -437,7 +462,7 @@
 
     untouched d = DimSlice 0 d 1
 
-    newBase user@(SegThread {}, _) (old_shape, _) =
+    newBaseThread user (old_shape, _) =
       let (users_shape, user_ids) = getNumUsers user
           num_dims = length old_shape
           perm = [num_dims .. num_dims + shapeRank users_shape - 1] ++ [0 .. num_dims - 1]
@@ -448,7 +473,10 @@
               Slice $
                 map DimFix user_ids ++ map untouched old_shape
        in offset_ixfun
-    newBase user@(SegGroup {}, _) (old_shape, _) =
+
+    newBase user@(SegThreadInGroup {}, _) = newBaseThread user
+    newBase user@(SegThread {}, _) = newBaseThread user
+    newBase user@(SegGroup {}, _) = \(old_shape, _) ->
       let (users_shape, user_ids) = getNumUsers user
           root_ixfun = IxFun.iota $ map pe64 (shapeDims users_shape) ++ old_shape
           offset_ixfun =
@@ -467,6 +495,8 @@
   where
     getNumUsers (SegThread {}, [gtid]) = (Shape [num_threads], [gtid])
     getNumUsers (SegThread {}, [gid, ltid]) = (Shape [num_groups, group_size], [gid, ltid])
+    getNumUsers (SegThreadInGroup {}, [gtid]) = (Shape [num_threads], [gtid])
+    getNumUsers (SegThreadInGroup {}, [gid, ltid]) = (Shape [num_groups, group_size], [gid, ltid])
     getNumUsers (SegGroup {}, [gid]) = (Shape [num_groups], [gid])
     getNumUsers user = error $ "getNumUsers: unhandled " ++ show user
 
diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs
--- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs
+++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs
@@ -20,32 +20,48 @@
   opIsConst (SizeOp GetSizeMax {}) = True
   opIsConst _ = False
 
-allocAtLevel :: SegLevel -> AllocM fromrep trep a -> AllocM fromrep trep a
+allocAtLevel :: SegLevel -> AllocM GPU GPUMem a -> AllocM GPU GPUMem a
 allocAtLevel lvl = local $ \env ->
   env
     { allocSpace = space,
-      aggressiveReuse = True
+      aggressiveReuse = True,
+      allocInOp = handleHostOp (Just lvl)
     }
   where
     space = case lvl of
-      SegThread {} -> DefaultSpace
       SegGroup {} -> Space "local"
+      SegThread {} -> DefaultSpace
+      SegThreadInGroup {} -> DefaultSpace
 
 handleSegOp ::
+  Maybe SegLevel ->
   SegOp SegLevel GPU ->
   AllocM GPU GPUMem (SegOp SegLevel GPUMem)
-handleSegOp op = do
+handleSegOp outer_lvl op = do
   num_threads <-
-    letSubExp "num_threads" $
-      BasicOp $
-        BinOp
-          (Mul Int64 OverflowUndef)
-          (unCount (segNumGroups lvl))
-          (unCount (segGroupSize lvl))
-  allocAtLevel lvl $ mapSegOpM (mapper num_threads) op
+    letSubExp "num_threads"
+      =<< case maybe_grid of
+        Just grid ->
+          pure . BasicOp $
+            BinOp
+              (Mul Int64 OverflowUndef)
+              (unCount (gridNumGroups grid))
+              (unCount (gridGroupSize grid))
+        Nothing ->
+          foldBinOp
+            (Mul Int64 OverflowUndef)
+            (intConst Int64 1)
+            (segSpaceDims $ segSpace op)
+  allocAtLevel (segLevel op) $ mapSegOpM (mapper num_threads) op
   where
+    maybe_grid =
+      case (outer_lvl, segLevel op) of
+        (Just (SegThread _ (Just grid)), _) -> Just grid
+        (Just (SegGroup _ (Just grid)), _) -> Just grid
+        (_, SegThread _ (Just grid)) -> Just grid
+        (_, SegGroup _ (Just grid)) -> Just grid
+        _ -> Nothing
     scope = scopeOfSegSpace $ segSpace op
-    lvl = segLevel op
     mapper num_threads =
       identitySegOpMapper
         { mapOnSegOpBody =
@@ -56,20 +72,22 @@
         }
     f = case segLevel op of
       SegThread {} -> inThread
+      SegThreadInGroup {} -> inThread
       SegGroup {} -> inGroup
     inThread env = env {envExpHints = inThreadExpHints}
     inGroup env = env {envExpHints = inGroupExpHints}
 
 handleHostOp ::
+  Maybe SegLevel ->
   HostOp GPU (SOAC GPU) ->
   AllocM GPU GPUMem (MemOp (HostOp GPUMem ()))
-handleHostOp (SizeOp op) =
+handleHostOp _ (SizeOp op) =
   pure $ Inner $ SizeOp op
-handleHostOp (OtherOp op) =
+handleHostOp _ (OtherOp op) =
   error $ "Cannot allocate memory in SOAC: " ++ prettyString op
-handleHostOp (SegOp op) =
-  Inner . SegOp <$> handleSegOp op
-handleHostOp (GPUBody ts (Body _ stms res)) =
+handleHostOp outer_lvl (SegOp op) =
+  Inner . SegOp <$> handleSegOp outer_lvl op
+handleHostOp _ (GPUBody ts (Body _ stms res)) =
   fmap (Inner . GPUBody ts) . buildBody_ . allocInStms stms $ pure res
 
 kernelExpHints :: Exp GPUMem -> AllocM GPU GPUMem [ExpHint]
@@ -79,9 +97,9 @@
       dims' = rearrangeShape perm dims
       ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv
   pure [Hint ixfun DefaultSpace]
-kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread {} space ts body)))) =
+kernelExpHints (Op (Inner (SegOp (SegMap lvl@(SegThread _ _) space ts body)))) =
   zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body
-kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread {} space reds ts body)))) =
+kernelExpHints (Op (Inner (SegOp (SegRed lvl@(SegThread _ _) space reds ts body)))) =
   (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res
   where
     num_reds = segBinOpResults reds
@@ -166,11 +184,11 @@
 
 -- | The pass from 'GPU' to 'GPUMem'.
 explicitAllocations :: Pass GPU GPUMem
-explicitAllocations = explicitAllocationsGeneric handleHostOp kernelExpHints
+explicitAllocations = explicitAllocationsGeneric (handleHostOp Nothing) kernelExpHints
 
 -- | Convert some 'GPU' stms to 'GPUMem'.
 explicitAllocationsInStms ::
   (MonadFreshNames m, HasScope GPUMem m) =>
   Stms GPU ->
   m (Stms GPUMem)
-explicitAllocationsInStms = explicitAllocationsInStmsGeneric handleHostOp kernelExpHints
+explicitAllocationsInStms = explicitAllocationsInStmsGeneric (handleHostOp Nothing) kernelExpHints
diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
--- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
+++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
@@ -21,6 +21,7 @@
 import Control.Monad.Writer
 import Futhark.Analysis.PrimExp
 import Futhark.IR
+import Futhark.IR.GPU.Op (SegVirt (..))
 import Futhark.IR.Prop.Aliases
 import Futhark.IR.SegOp
 import Futhark.MonadFreshNames
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
@@ -39,6 +39,7 @@
 import Data.Map qualified as M
 import Data.Maybe
 import Futhark.IR
+import Futhark.IR.GPU.Op (SegVirt (..))
 import Futhark.IR.SOACS (SOACS)
 import Futhark.IR.SOACS qualified as SOACS
 import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest)
@@ -375,7 +376,7 @@
         Nothing -> addStmToAcc stm acc
         Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmAux stm) w lam arrs) acc'
 maybeDistributeStm stm@(Let pat aux (DoLoop merge form@ForLoop {} body)) acc
-  | all (`notNameIn` freeIn pat) (patNames pat),
+  | all (`notNameIn` freeIn (patTypes pat)) (patNames pat),
     bodyContainsParallelism body =
       distributeSingleStm acc stm >>= \case
         Just (kernels, res, nest, acc')
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
@@ -60,12 +60,9 @@
   let body = lambdaBody lam
 
   group_size <- newVName "computed_group_size"
-  let intra_lvl = SegThread (Count num_groups) (Count $ Var group_size) SegNoVirt
-
   (wss_min, wss_avail, log, kbody) <-
-    lift $
-      localScope (scopeOfLParams $ lambdaParams lam) $
-        intraGroupParalleliseBody intra_lvl body
+    lift . localScope (scopeOfLParams $ lambdaParams lam) $
+      intraGroupParalleliseBody body
 
   outside_scope <- lift askScope
   -- outside_scope may also contain the inputs, even though those are
@@ -116,12 +113,10 @@
 
   let nested_pat = loopNestingPat first_nest
       rts = map (length ispace `stripArray`) $ patTypes nested_pat
-      lvl = SegGroup (Count num_groups) (Count $ Var group_size) SegNoVirt
+      grid = KernelGrid (Count num_groups) (Count $ Var group_size)
+      lvl = SegGroup SegNoVirt (Just grid)
       kstm =
-        Let nested_pat aux $
-          Op $
-            SegOp $
-              SegMap lvl kspace rts kbody'
+        Let nested_pat aux $ Op $ SegOp $ SegMap lvl kspace rts kbody'
 
   let intra_min_par = intra_avail_par
   pure
@@ -184,21 +179,21 @@
         accAvailPar = S.singleton ws
       }
 
-intraGroupBody :: SegLevel -> Body SOACS -> IntraGroupM (Body GPU)
-intraGroupBody lvl body = do
-  stms <- collectStms_ $ intraGroupStms lvl $ bodyStms body
+intraGroupBody :: Body SOACS -> IntraGroupM (Body GPU)
+intraGroupBody body = do
+  stms <- collectStms_ $ intraGroupStms $ bodyStms body
   pure $ mkBody stms $ bodyResult body
 
-intraGroupStm :: SegLevel -> Stm SOACS -> IntraGroupM ()
-intraGroupStm lvl stm@(Let pat aux e) = do
+intraGroupStm :: Stm SOACS -> IntraGroupM ()
+intraGroupStm stm@(Let pat aux e) = do
   scope <- askScope
-  let lvl' = SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt
+  let lvl = SegThread SegNoVirt Nothing
 
   case e of
     DoLoop merge form loopbody ->
       localScope (scopeOf form') $
         localScope (scopeOfFParams $ map fst merge) $ do
-          loopbody' <- intraGroupBody lvl loopbody
+          loopbody' <- intraGroupBody loopbody
           certifying (stmAuxCerts aux) $
             letBind pat $
               DoLoop merge form' loopbody'
@@ -207,13 +202,13 @@
           ForLoop i it bound inps -> ForLoop i it bound inps
           WhileLoop cond -> WhileLoop cond
     Match cond cases defbody ifdec -> do
-      cases' <- mapM (traverse $ intraGroupBody lvl) cases
-      defbody' <- intraGroupBody lvl defbody
+      cases' <- mapM (traverse intraGroupBody) cases
+      defbody' <- intraGroupBody defbody
       certifying (stmAuxCerts aux) . letBind pat $
         Match cond cases' defbody' ifdec
     Op soac
       | "sequential_outer" `inAttrs` stmAuxAttrs aux ->
-          intraGroupStms lvl . fmap (certify (stmAuxCerts aux))
+          intraGroupStms . fmap (certify (stmAuxCerts aux))
             =<< runBuilder_ (FOT.transformSOAC pat soac)
     Op (Screma w arrs form)
       | Just lam <- isMapSOAC form -> do
@@ -229,7 +224,7 @@
                     distOnInnerMap =
                       distributeMap,
                     distOnTopLevelStms =
-                      liftInner . collectStms_ . intraGroupStms lvl,
+                      liftInner . collectStms_ . intraGroupStms,
                     distSegLevel = \minw _ _ -> do
                       lift $ parallelMin minw
                       pure lvl,
@@ -252,7 +247,7 @@
           let scanfun' = soacsLambdaToGPU scanfun
               mapfun' = soacsLambdaToGPU mapfun
           certifying (stmAuxCerts aux) $
-            addStms =<< segScan lvl' pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] []
+            addStms =<< segScan lvl pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] []
           parallelMin [w]
     Op (Screma w arrs form)
       | Just (reds, map_lam) <- isRedomapSOAC form,
@@ -260,7 +255,7 @@
           let red_lam' = soacsLambdaToGPU red_lam
               map_lam' = soacsLambdaToGPU map_lam
           certifying (stmAuxCerts aux) $
-            addStms =<< segRed lvl' pat mempty w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []
+            addStms =<< segRed lvl pat mempty w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []
           parallelMin [w]
     Op (Hist w arrs ops bucket_fun) -> do
       ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do
@@ -270,7 +265,7 @@
 
       let bucket_fun' = soacsLambdaToGPU bucket_fun
       certifying (stmAuxCerts aux) $
-        addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs
+        addStms =<< segHist lvl pat w [] [] ops' bucket_fun' arrs
       parallelMin [w]
     Op (Stream w arrs accs lam)
       | chunk_size_param : _ <- lambdaParams lam -> do
@@ -281,7 +276,7 @@
               replace se = se
               replaceSets (IntraAcc x y log) =
                 IntraAcc (S.map (map replace) x) (S.map (map replace) y) log
-          censor replaceSets $ intraGroupStms lvl stream_stms
+          censor replaceSets $ intraGroupStms stream_stms
     Op (Scatter w ivs lam dests) -> do
       write_i <- newVName "write_i"
       space <- mkSegSpace [(write_i, w)]
@@ -308,23 +303,22 @@
       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
+        letBind pat $ Op $ SegOp $ SegMap lvl space ts body
 
       parallelMin [w]
     _ ->
       addStm $ soacsStmToGPU stm
 
-intraGroupStms :: SegLevel -> Stms SOACS -> IntraGroupM ()
-intraGroupStms lvl = mapM_ (intraGroupStm lvl)
+intraGroupStms :: Stms SOACS -> IntraGroupM ()
+intraGroupStms = mapM_ intraGroupStm
 
 intraGroupParalleliseBody ::
   (MonadFreshNames m, HasScope GPU m) =>
-  SegLevel ->
   Body SOACS ->
   m ([[SubExp]], [[SubExp]], Log, KernelBody GPU)
-intraGroupParalleliseBody lvl body = do
+intraGroupParalleliseBody body = do
   (IntraAcc min_ws avail_ws log, kstms) <-
-    runIntraGroupM $ intraGroupStms lvl $ bodyStms body
+    runIntraGroupM $ intraGroupStms $ bodyStms body
   pure
     ( S.toList min_ws,
       S.toList avail_ws,
diff --git a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs
--- a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs
+++ b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs
@@ -74,7 +74,9 @@
             (SDivUp Int64 Unsafe)
             (eSubExp w)
             (eSubExp =<< asIntS Int64 group_size)
-      pure $ SegThread (Count usable_groups) (Count group_size) SegNoVirt
+      let grid = KernelGrid (Count usable_groups) (Count group_size)
+      pure $ SegThread SegNoVirt (Just grid)
     NoRecommendation v -> do
       (num_groups, _) <- numberOfGroups desc w group_size
-      pure $ SegThread (Count num_groups) (Count group_size) v
+      let grid = KernelGrid (Count num_groups) (Count group_size)
+      pure $ SegThread v (Just grid)
diff --git a/src/Futhark/Pass/ExtractKernels/ToGPU.hs b/src/Futhark/Pass/ExtractKernels/ToGPU.hs
--- a/src/Futhark/Pass/ExtractKernels/ToGPU.hs
+++ b/src/Futhark/Pass/ExtractKernels/ToGPU.hs
@@ -34,10 +34,12 @@
   String ->
   m SegLevel
 segThread desc =
-  SegThread
-    <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups)
-    <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup)
-    <*> pure SegVirt
+  SegThread SegVirt <$> (Just <$> kernelGrid)
+  where
+    kernelGrid =
+      KernelGrid
+        <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups)
+        <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup)
 
 injectSOACS ::
   ( Monad m,
diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs
--- a/src/Futhark/Util.hs
+++ b/src/Futhark/Util.hs
@@ -43,6 +43,7 @@
     zEncodeString,
     atMostChars,
     invertMap,
+    cartesian,
     traverseFold,
     fixPoint,
   )
@@ -57,7 +58,7 @@
 import Data.ByteString.Base16 qualified as Base16
 import Data.Char
 import Data.Either
-import Data.Foldable (fold)
+import Data.Foldable (fold, toList)
 import Data.Function ((&))
 import Data.List (foldl', genericDrop, genericSplitAt, sortBy)
 import Data.List.NonEmpty qualified as NE
@@ -411,6 +412,13 @@
   M.toList m
     & fmap (swap . first S.singleton)
     & foldr (uncurry $ M.insertWith (<>)) mempty
+
+-- | Compute the cartesian product of two foldable collections, using the given
+-- combinator function.
+cartesian :: (Monoid m, Foldable t) => (a -> a -> m) -> t a -> t a -> m
+cartesian f xs ys =
+  [(x, y) | x <- toList xs, y <- toList ys]
+    & foldMap (uncurry f)
 
 -- | Applicatively fold a traversable.
 traverseFold :: (Monoid m, Traversable t, Applicative f) => (a -> f m) -> t a -> f m
diff --git a/src/Futhark/Util/Pretty.hs b/src/Futhark/Util/Pretty.hs
--- a/src/Futhark/Util/Pretty.hs
+++ b/src/Futhark/Util/Pretty.hs
@@ -35,6 +35,7 @@
     stack,
     parensIf,
     ppTuple',
+    ppTupleLines',
 
     -- * Operators
     (</>),
@@ -113,15 +114,16 @@
 ppTuple' :: [Doc a] -> Doc a
 ppTuple' ets = braces $ commasep $ map align ets
 
+ppTupleLines' :: [Doc a] -> Doc a
+ppTupleLines' ets = braces $ commastack $ map align ets
+
 -- | Prettyprint a list enclosed in curly braces.
 prettyTuple :: Pretty a => [a] -> Text
 prettyTuple = docText . ppTuple' . map pretty
 
 -- | Like 'prettyTuple', but put a linebreak after every element.
 prettyTupleLines :: Pretty a => [a] -> Text
-prettyTupleLines = docText . ppTupleLines'
-  where
-    ppTupleLines' = braces . hsep . punctuate comma . map (align . pretty)
+prettyTupleLines = docText . ppTupleLines' . map pretty
 
 -- | The document @'apply' ds@ separates @ds@ with commas and encloses them with
 -- parentheses.
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
@@ -1033,7 +1033,7 @@
     onModule (Module (Env terms types _)) =
       Module $ Env (replaceM onTerm terms) (replaceM onType types) mempty
     onModule (ModuleFun f) =
-      ModuleFun $ \m -> onModule <$> f (substituteInModule (M.mapMaybe maybeHead rev_substs) m)
+      ModuleFun $ \m -> onModule <$> f (substituteInModule substs m)
     onTerm (TermValue t v) = TermValue t v
     onTerm (TermPoly t v) = TermPoly t v
     onTerm (TermModule m) = TermModule $ onModule m
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
@@ -253,8 +253,7 @@
       _ -> hasArrayLit e
 prettyAppExp _ (LetFun fname (tparams, params, retdecl, rettype, e) body _) =
   "let"
-    <+> prettyName fname
-    <+> hsep (map pretty tparams ++ map pretty params)
+    <+> hsep (prettyName fname : map pretty tparams ++ map pretty params)
       <> retdecl'
     <+> equals
     </> indent 2 (pretty e)
@@ -468,8 +467,7 @@
 instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where
   pretty (TypeBind name l params te rt _ _) =
     "type" <> pretty l
-      <+> prettyName name
-      <+> hsep (map pretty params)
+      <+> hsep (prettyName name : map pretty params)
       <+> equals
       <+> maybe (pretty te) pretty (unAnnot rt)
 
@@ -502,9 +500,9 @@
 instance (Eq vn, IsName vn, Annot f) => Pretty (SpecBase f vn) where
   pretty (TypeAbbrSpec tpsig) = pretty tpsig
   pretty (TypeSpec l name ps _ _) =
-    "type" <> pretty l <+> prettyName name <+> hsep (map pretty ps)
+    "type" <> pretty l <+> hsep (prettyName name : map pretty ps)
   pretty (ValSpec name tparams vtype _ _ _) =
-    "val" <+> prettyName name <+> hsep (map pretty tparams) <> colon <+> pretty vtype
+    "val" <+> hsep (prettyName name : map pretty tparams) <> colon <+> pretty vtype
   pretty (ModSpec name sig _ _) =
     "module" <+> prettyName name <> colon <+> pretty sig
   pretty (IncludeSpec e _) =
@@ -531,11 +529,11 @@
 
 instance (Eq vn, IsName vn, Annot f) => Pretty (ModBindBase f vn) where
   pretty (ModBind name ps sig e _ _) =
-    "module" <+> prettyName name <+> hsep (map pretty ps) <+> sig' <> " =" <+> pretty e
+    "module" <+> hsep (prettyName name : map pretty ps) <> sig' <> " =" <+> pretty e
     where
       sig' = case sig of
         Nothing -> mempty
-        Just (s, _) -> colon <+> pretty s <> " "
+        Just (s, _) -> " " <> colon <+> pretty s <> " "
 
 ppBinOp :: IsName v => QualName v -> Doc a
 ppBinOp bop =
diff --git a/src/Language/Futhark/TypeChecker/Match.hs b/src/Language/Futhark/TypeChecker/Match.hs
--- a/src/Language/Futhark/TypeChecker/Match.hs
+++ b/src/Language/Futhark/TypeChecker/Match.hs
@@ -8,6 +8,7 @@
   )
 where
 
+import Data.List qualified as L
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Futhark.Util (maybeHead, nubOrd)
@@ -23,16 +24,16 @@
   deriving (Eq, Ord, Show)
 
 -- | A representation of the essentials of a pattern.
-data Match
-  = MatchWild StructType
-  | MatchConstr Constr [Match] StructType
+data Match t
+  = MatchWild t
+  | MatchConstr Constr [Match t] t
   deriving (Eq, Ord, Show)
 
-matchType :: Match -> StructType
+matchType :: Match StructType -> StructType
 matchType (MatchWild t) = t
 matchType (MatchConstr _ _ t) = t
 
-pprMatch :: Int -> Match -> Doc a
+pprMatch :: Int -> Match t -> Doc a
 pprMatch _ MatchWild {} = "_"
 pprMatch _ (MatchConstr (ConstrLit l) _ _) = pretty l
 pprMatch p (MatchConstr (Constr c) ps _) =
@@ -45,10 +46,10 @@
   where
     ppField name t = pretty (nameToString name) <> equals <> pprMatch (-1) t
 
-instance Pretty Match where
+instance Pretty (Match t) where
   pretty = pprMatch (-1)
 
-patternToMatch :: Pat -> Match
+patternToMatch :: Pat -> Match StructType
 patternToMatch (Id _ (Info t) _) = MatchWild $ toStruct t
 patternToMatch (Wildcard (Info t) _) = MatchWild $ toStruct t
 patternToMatch (PatParens p _) = patternToMatch p
@@ -67,29 +68,35 @@
 patternToMatch (PatConstr c (Info t) args _) =
   MatchConstr (Constr c) (map patternToMatch args) $ toStruct t
 
-isConstr :: Match -> Maybe Name
+isConstr :: Match t -> Maybe Name
 isConstr (MatchConstr (Constr c) _ _) = Just c
 isConstr _ = Nothing
 
-complete :: [Match] -> Bool
+isBool :: Match t -> Maybe Bool
+isBool (MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) _ _) = Just b
+isBool _ = Nothing
+
+complete :: [Match StructType] -> Bool
 complete xs
   | Just x <- maybeHead xs,
     Scalar (Sum all_cs) <- matchType x,
     Just xs_cs <- mapM isConstr xs =
       all (`elem` xs_cs) (M.keys all_cs)
   | otherwise =
-      (any (isBool True) xs && any (isBool False) xs)
+      all (`elem` fromMaybe [] (mapM isBool xs)) [True, False]
         || all isRecord xs
         || all isTuple xs
   where
-    isBool b1 (MatchConstr (ConstrLit (PatLitPrim (BoolValue b2))) _ _) = b1 == b2
-    isBool _ _ = False
     isRecord (MatchConstr ConstrRecord {} _ _) = True
     isRecord _ = False
     isTuple (MatchConstr ConstrTuple _ _) = True
     isTuple _ = False
 
-specialise :: [StructType] -> Match -> [[Match]] -> [[Match]]
+specialise ::
+  [StructType] ->
+  Match StructType ->
+  [[Match StructType]] ->
+  [[Match StructType]]
 specialise ats c1 = go
   where
     go ((c2 : row) : ps)
@@ -109,14 +116,14 @@
     match _ _ =
       Nothing
 
-defaultMat :: [[Match]] -> [[Match]]
+defaultMat :: [[Match t]] -> [[Match t]]
 defaultMat = mapMaybe onRow
   where
     onRow (MatchConstr {} : _) = Nothing
     onRow (MatchWild {} : ps) = Just ps
     onRow [] = Nothing -- Should not happen.
 
-findUnmatched :: [[Match]] -> Int -> [[Match]]
+findUnmatched :: [[Match StructType]] -> Int -> [[Match ()]]
 findUnmatched pmat n
   | ((p : _) : _) <- pmat,
     Just heads <- mapM maybeHead pmat =
@@ -133,40 +140,44 @@
           pmat' = specialise ats c pmat
       u <- findUnmatched pmat' (a_k + n - 1)
       pure $ case c of
-        MatchConstr c' _ t ->
+        MatchConstr c' _ _ ->
           let (r, p) = splitAt a_k u
-           in MatchConstr c' r t : p
-        MatchWild t ->
-          MatchWild t : u
+           in MatchConstr c' r () : p
+        MatchWild _ ->
+          MatchWild () : u
 
     incompleteCase pt cs = do
       u <- findUnmatched (defaultMat pmat) (n - 1)
       if null cs
-        then pure $ MatchWild pt : u
+        then pure $ MatchWild () : u
         else case pt of
           Scalar (Sum all_cs) -> do
             -- Figure out which constructors are missing.
             let sigma = mapMaybe isConstr cs
                 notCovered (k, _) = k `notElem` sigma
             (cname, ts) <- filter notCovered $ M.toList all_cs
-            pure $ MatchConstr (Constr cname) (map MatchWild ts) pt : u
-          _ ->
-            -- This is where we could have enumerated missing match
-            -- values (e.g. for booleans), rather than just emitting a
-            -- wildcard.
-            pure $ MatchWild pt : u
-
--- If we get here, then the number of columns must be zero.
-findUnmatched [] _ = [[]]
+            pure $ MatchConstr (Constr cname) (map (const (MatchWild ())) ts) () : u
+          Scalar (Prim Bool) -> do
+            -- Figure out which constants are missing.
+            let sigma = mapMaybe isBool cs
+            b <- filter (`notElem` sigma) [True, False]
+            pure $ MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) [] () : u
+          _ -> do
+            -- FIXME: this is wrong in the unlikely case where someone
+            -- is pattern-matching every single possible number for
+            -- some numeric type.  It should be handled more like Bool
+            -- above.
+            pure $ MatchWild () : u
+findUnmatched [] n = [replicate n $ MatchWild ()]
 findUnmatched _ _ = []
 
 {-# NOINLINE unmatched #-}
 
 -- | Find the unmatched cases.
-unmatched :: [Pat] -> [Match]
+unmatched :: [Pat] -> [Match ()]
 unmatched orig_ps =
   -- The algorithm may find duplicate example, which we filter away
   -- here.
   nubOrd $
     mapMaybe maybeHead $
-      findUnmatched (map ((: []) . patternToMatch) orig_ps) 1
+      findUnmatched (map (L.singleton . patternToMatch) orig_ps) 1
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
@@ -375,8 +375,7 @@
         <+> hsep (map pretty ps)
 ppTypeAbbr _ name (l, ps, t) =
   "type" <> pretty l
-    <+> pretty name
-    <+> hsep (map pretty ps)
+    <+> hsep (pretty name : map pretty ps)
     <+> equals
     <+> nest 2 (align (pretty t))
 
@@ -455,11 +454,19 @@
       (ModFun (FunSig mod_abs mod_pmod mod_mod))
       (ModFun (FunSig sig_abs sig_pmod sig_mod))
       loc = do
+        -- We need to use different substitutions when matching
+        -- parameter and body signatures - this is because the
+        -- concrete parameter must be *at least as* general as the
+        -- ascripted parameter, while the concrete body must be *at
+        -- most as* general as the ascripted body.
         abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc
+        p_abs_substs <- resolveAbsTypes sig_abs sig_pmod mod_abs loc
         let abs_subst_to_type =
               old_abs_subst_to_type <> M.map (substFromAbbr . snd) abs_substs
+            p_abs_subst_to_type =
+              old_abs_subst_to_type <> M.map (substFromAbbr . snd) p_abs_substs
             abs_name_substs = M.map (qualLeaf . fst) abs_substs
-        pmod_substs <- matchMods abs_subst_to_type quals mod_pmod sig_pmod loc
+        pmod_substs <- matchMods p_abs_subst_to_type quals sig_pmod mod_pmod loc
         mod_substs <- matchMTys' abs_subst_to_type quals mod_mod sig_mod loc
         pure (pmod_substs <> mod_substs <> abs_name_substs)
 
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
@@ -895,10 +895,9 @@
            in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1
         _ -> pure ()
 
-      occurs <- (dflow `seqOccurrences`) <$> consumeArg argloc argtype' (diet tp1')
-
-      checkIfConsumable loc $ S.map AliasBound $ allConsumed occurs
-      occur occurs
+      arg_consumed <- consumedByArg argloc argtype' (diet tp1')
+      checkIfConsumable loc $ mconcat arg_consumed
+      occur $ dflow `seqOccurrences` map (`consumption` argloc) arg_consumed
 
       -- Unification ignores uniqueness in higher-order arguments, so
       -- we check for that here.
@@ -1011,12 +1010,12 @@
 maskAliases t FuncDiet {} = t
 maskAliases _ _ = error "Invalid arguments passed to maskAliases."
 
-consumeArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Occurrence]
-consumeArg loc (Scalar (Record ets)) (RecordDiet ds) =
-  concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds)
-consumeArg loc (Scalar (Sum ets)) (SumDiet ds) =
-  concat <$> traverse (uncurry $ consumeArg loc) (concat $ M.elems $ M.intersectionWith zip ets ds)
-consumeArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _)
+consumedByArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Aliasing]
+consumedByArg loc (Scalar (Record ets)) (RecordDiet ds) =
+  mconcat . M.elems <$> traverse (uncurry $ consumedByArg loc) (M.intersectionWith (,) ets ds)
+consumedByArg loc (Scalar (Sum ets)) (SumDiet ds) =
+  mconcat <$> traverse (uncurry $ consumedByArg loc) (concat $ M.elems $ M.intersectionWith zip ets ds)
+consumedByArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _)
   | not $ contravariantArg t1 d =
       typeError loc mempty . withIndexLink "consuming-argument" $
         "Non-consuming higher-order parameter passed consuming argument."
@@ -1031,8 +1030,8 @@
       contravariantArg tp dp && contravariantArg tr dr
     contravariantArg _ _ =
       True
-consumeArg loc at Consume = pure [consumption (aliases at) loc]
-consumeArg loc at _ = pure [observation (aliases at) loc]
+consumedByArg _ at Consume = pure [aliases at]
+consumedByArg _ _ _ = pure []
 
 -- | Type-check a single expression in isolation.  This expression may
 -- turn out to be polymorphic, in which case the list of type
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
@@ -69,7 +69,6 @@
     -- * Errors
     useAfterConsume,
     unusedSize,
-    notConsumable,
     uniqueReturnAliased,
     returnAliased,
     badLetWithValue,
@@ -970,14 +969,15 @@
 checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM ()
 checkIfConsumable loc als = do
   vtable <- asks $ scopeVtable . termScope
-  let consumable v = case M.lookup v vtable of
+  let boundAlias (AliasBound v) = Just v
+      boundAlias (AliasFree _) = Nothing
+      consumable v = case M.lookup v vtable of
         Just (BoundV Local _ t)
           | Scalar Arrow {} <- t -> False
           | otherwise -> True
         Just (BoundV l _ _) -> l == Local
-        _ -> True
-  -- The sort ensures that AliasBound vars are shown before AliasFree.
-  case map aliasVar $ sort $ filter (not . consumable . aliasVar) $ S.toList als of
+        _ -> False -- Implies name from module.
+  case sort $ filter (not . consumable) $ mapMaybe boundAlias $ S.toList als of
     v : _ -> notConsumable loc =<< describeVar loc v
     [] -> pure ()