diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,39 @@
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
+## [0.25.10]
+
+### Added
+
+* Faster non-commutative reductions in the GPU backends. Work by
+  Anders Holst and Christian Påbøl Jacobsen.
+
+### Fixed
+
+* Interpreter crash for certain complicated size expressions involving
+  internal bindings (#2053).
+
+* Incorrect type checking of `let` binding with explicit size
+  quantification, where size appears in type of body (#2048).
+
+* GPU code generation for non-commutative non-segmented reductions
+  with array operands (#2051).
+
+* Histogram with non-vectorised reduction operators (#2056). (But it
+  is probably not a good idea to write such programs.)
+
+* Futhark's LSP server should work better with Eglot.
+
+* Incorrect copy removal inside histograms could cause compiler error
+  (#2058).
+
+* CUDA backend now correctly queries for available shared memory,
+  which affects performance (hopefully positively).
+
+* `futhark literate` now switches to the directory containing the
+  `.fut` file before executing its contents. This fixes accessing
+  files through relative paths.
+
 ## [0.25.9]
 
 ### Added
diff --git a/docs/c-api.rst b/docs/c-api.rst
--- a/docs/c-api.rst
+++ b/docs/c-api.rst
@@ -141,9 +141,7 @@
    cache was hit succesfully, but you can enable logging to see what
    happens.
 
-   The lifespan of ``fname`` must exceed the lifespan of the
-   configuration object.  Pass ``NULL`` to disable caching (this is
-   the default).
+   Pass ``NULL`` to disable caching (this is the default).
 
 Context
 -------
diff --git a/docs/language-reference.rst b/docs/language-reference.rst
--- a/docs/language-reference.rst
+++ b/docs/language-reference.rst
@@ -576,6 +576,9 @@
   ``t``.  To pass a single array-typed parameter, enclose it in
   parens.
 
+* The bodies of ``let``, ``if``, and ``loop`` extend as far to the
+  right as possible.
+
 * The following table describes the precedence and associativity of
   infix operators in both expressions and type expressions.  All
   operators in the same row have the same precedence.  The rows are
@@ -836,10 +839,10 @@
 
     Company any two values of numeric type for equality.
 
-  ```symbol```
+  ```qualname```
 
-    Use ``symbol``, which may be any non-operator function name, as an
-    infix operator.
+    Use ``qualname``, which may be any non-operator function name, as
+    an infix operator.
 
 ``x && y``
 ..........
@@ -1499,14 +1502,14 @@
    mod_param: "(" `name` ":" `mod_type_exp` ")"
    mod_type_bind: "module" "type" `name` "=" `mod_type_exp`
 
-Futhark supports an ML-style higher-order module system.  *Modules*
-can contain types, functions, and other modules and module types.
-*Module types* are used to classify the contents of modules, and
-*parametric modules* are used to abstract over modules (essentially
-module-level functions).  In Standard ML, modules, module types and
-parametric modules are called structs, signatures, and functors,
-respectively.  Module names exist in the same name space as values,
-but module types are their own name space.
+Futhark supports an ML-style higher-order module system. *Modules* can
+contain types, functions, and other modules and module types. *Module
+types* are used to classify the contents of modules, and *parametric
+modules* are used to abstract over modules (essentially module-level
+functions). In Standard ML, modules, module types and parametric
+modules are called *structs*, *signatures*, and *functors*,
+respectively. Module names exist in the same name space as values, but
+module types are their own name space.
 
 Module bindings
 ~~~~~~~~~~~~~~~
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -1,6 +1,6 @@
 cabal-version: 2.4
 name:           futhark
-version:        0.25.9
+version:        0.25.10
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -313,9 +313,6 @@
       Futhark.Optimise.Fusion.TryFusion
       Futhark.Optimise.GenRedOpt
       Futhark.Optimise.HistAccs
-      Futhark.Optimise.InPlaceLowering
-      Futhark.Optimise.InPlaceLowering.LowerIntoStm
-      Futhark.Optimise.InPlaceLowering.SubstituteIndices
       Futhark.Optimise.InliningDeadFun
       Futhark.Optimise.MemoryBlockMerging
       Futhark.Optimise.MemoryBlockMerging.GreedyColoring
diff --git a/rts/c/backends/cuda.h b/rts/c/backends/cuda.h
--- a/rts/c/backends/cuda.h
+++ b/rts/c/backends/cuda.h
@@ -289,6 +289,8 @@
   size_t max_threshold;
   size_t max_local_memory;
   size_t max_bespoke;
+  size_t max_registers;
+  size_t max_cache;
 
   size_t lockstep_width;
 
@@ -429,6 +431,10 @@
   int arch_set = 0, num_extra_opts;
   struct futhark_context_config *cfg = ctx->cfg;
 
+  char** macro_names;
+  int64_t* macro_vals;
+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);
+
   // nvrtc cannot handle multiple -arch options.  Hence, if one of the
   // extra_opts is -arch, we have to be careful not to do our usual
   // automatic generation.
@@ -441,7 +447,7 @@
     }
   }
 
-  size_t i = 0, n_opts_alloc = 20 + num_extra_opts + cfg->num_tuning_params;
+  size_t i = 0, n_opts_alloc = 20 + num_macros + num_extra_opts + cfg->num_tuning_params;
   char **opts = (char**) malloc(n_opts_alloc * sizeof(char *));
   if (!arch_set) {
     opts[i++] = strdup("-arch");
@@ -457,6 +463,17 @@
   opts[i++] = msgprintf("-D%s=%d",
                         "max_group_size",
                         (int)ctx->max_group_size);
+  opts[i++] = msgprintf("-D%s=%d",
+                        "max_local_memory",
+                        (int)ctx->max_local_memory);
+  opts[i++] = msgprintf("-D%s=%d",
+                        "max_registers",
+                        (int)ctx->max_registers);
+
+  for (int j = 0; j < num_macros; j++) {
+    opts[i++] = msgprintf("-D%s=%zu", macro_names[j], macro_vals[j]);
+  }
+
   for (int j = 0; j < cfg->num_tuning_params; j++) {
     opts[i++] = msgprintf("-D%s=%zu", cfg->tuning_param_vars[j],
                           cfg->tuning_params[j]);
@@ -485,6 +502,9 @@
   opts[i++] = msgprintf("-DTR_TILE_DIM=%d", TR_TILE_DIM);
   opts[i++] = msgprintf("-DTR_ELEMS_PER_THREAD=%d", TR_ELEMS_PER_THREAD);
 
+  free(macro_names);
+  free(macro_vals);
+
   *n_opts = i;
   *opts_out = opts;
 }
@@ -798,12 +818,19 @@
 
   free_list_init(&ctx->gpu_free_list);
 
-  ctx->max_local_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);
+  // MAX_SHARED_MEMORY_PER_BLOCK gives bogus numbers (48KiB); probably
+  // for backwards compatibility.  Add _OPTIN and you seem to get the
+  // right number.
+  ctx->max_local_memory =
+    device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) -
+    device_query(ctx->dev, RESERVED_SHARED_MEMORY_PER_BLOCK);
   ctx->max_group_size = device_query(ctx->dev, MAX_THREADS_PER_BLOCK);
   ctx->max_grid_size = device_query(ctx->dev, MAX_GRID_DIM_X);
   ctx->max_tile_size = sqrt(ctx->max_group_size);
   ctx->max_threshold = 0;
   ctx->max_bespoke = 0;
+  ctx->max_registers = device_query(ctx->dev, MAX_REGISTERS_PER_BLOCK);
+  ctx->max_cache = device_query(ctx->dev, L2_CACHE_SIZE);
   ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);
   CUDA_SUCCEED_FATAL(cuStreamCreate(&ctx->stream, CU_STREAM_DEFAULT));
   cuda_size_setup(ctx);
@@ -853,6 +880,10 @@
     fprintf(ctx->log, "Creating kernel %s.\n", name);
   }
   CUDA_SUCCEED_FATAL(cuModuleGetFunction(kernel, ctx->module, name));
+  // Unless the below is set, the kernel is limited to 48KiB of memory.
+  CUDA_SUCCEED_FATAL(cuFuncSetAttribute(*kernel,
+                                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                        ctx->max_local_memory));
 }
 
 static void gpu_free_kernel(struct futhark_context *ctx,
diff --git a/rts/c/backends/hip.h b/rts/c/backends/hip.h
--- a/rts/c/backends/hip.h
+++ b/rts/c/backends/hip.h
@@ -263,6 +263,8 @@
   size_t max_threshold;
   size_t max_local_memory;
   size_t max_bespoke;
+  size_t max_registers;
+  size_t max_cache;
 
   size_t lockstep_width;
 
@@ -452,6 +454,10 @@
   int arch_set = 0, num_extra_opts;
   struct futhark_context_config *cfg = ctx->cfg;
 
+  char** macro_names;
+  int64_t* macro_vals;
+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);
+
   for (num_extra_opts = 0; extra_opts[num_extra_opts] != NULL; num_extra_opts++) {
     if (strstr(extra_opts[num_extra_opts], "--gpu-architecture")
         == extra_opts[num_extra_opts]) {
@@ -459,7 +465,7 @@
     }
   }
 
-  size_t i = 0, n_opts_alloc = 20 + num_extra_opts + cfg->num_tuning_params;
+  size_t i = 0, n_opts_alloc = 20 + num_macros + num_extra_opts + cfg->num_tuning_params;
   char **opts = (char**) malloc(n_opts_alloc * sizeof(char *));
   if (!arch_set) {
     hipDeviceProp_t props;
@@ -473,6 +479,17 @@
   opts[i++] = msgprintf("-D%s=%d",
                         "max_group_size",
                         (int)ctx->max_group_size);
+  opts[i++] = msgprintf("-D%s=%d",
+                        "max_local_memory",
+                        (int)ctx->max_local_memory);
+  opts[i++] = msgprintf("-D%s=%d",
+                        "max_registers",
+                        (int)ctx->max_registers);
+
+  for (int j = 0; j < num_macros; j++) {
+    opts[i++] = msgprintf("-D%s=%zu", macro_names[j], macro_vals[j]);
+  }
+
   for (int j = 0; j < cfg->num_tuning_params; j++) {
     opts[i++] = msgprintf("-D%s=%zu", cfg->tuning_param_vars[j],
                           cfg->tuning_params[j]);
@@ -488,6 +505,9 @@
   opts[i++] = msgprintf("-DTR_TILE_DIM=%d", TR_TILE_DIM);
   opts[i++] = msgprintf("-DTR_ELEMS_PER_THREAD=%d", TR_ELEMS_PER_THREAD);
 
+  free(macro_names);
+  free(macro_vals);
+
   *n_opts = i;
   *opts_out = opts;
 }
@@ -659,6 +679,8 @@
   ctx->max_tile_size = sqrt(ctx->max_group_size);
   ctx->max_threshold = 0;
   ctx->max_bespoke = 0;
+  ctx->max_registers = device_query(ctx->dev, hipDeviceAttributeMaxRegistersPerBlock);
+  ctx->max_cache = device_query(ctx->dev, hipDeviceAttributeL2CacheSize);
   // FIXME: in principle we should query hipDeviceAttributeWarpSize
   // from the device, which will provide 64 on AMD GPUs.
   // Unfortunately, we currently do nasty implicit intra-warp
diff --git a/rts/c/backends/opencl.h b/rts/c/backends/opencl.h
--- a/rts/c/backends/opencl.h
+++ b/rts/c/backends/opencl.h
@@ -528,6 +528,8 @@
   size_t max_tile_size;
   size_t max_threshold;
   size_t max_local_memory;
+  size_t max_registers;
+  size_t max_cache;
 
   size_t lockstep_width;
 
@@ -579,10 +581,18 @@
     compile_opts_size += strlen(ctx->cfg->tuning_param_names[i]) + 20;
   }
 
+  char** macro_names;
+  int64_t* macro_vals;
+  int num_macros = gpu_macros(ctx, &macro_names, &macro_vals);
+
   for (int i = 0; extra_build_opts[i] != NULL; i++) {
     compile_opts_size += strlen(extra_build_opts[i] + 1);
   }
 
+  for (int i = 0; i < num_macros; i++) {
+    compile_opts_size += strlen(macro_names[i]) + 1 + 20;
+  }
+
   char *compile_opts = (char*) malloc(compile_opts_size);
 
   int w = snprintf(compile_opts, compile_opts_size,
@@ -594,6 +604,16 @@
                 "max_group_size",
                 (int)ctx->max_group_size);
 
+  w += snprintf(compile_opts+w, compile_opts_size-w,
+                "-D%s=%d ",
+                "max_local_memory",
+                (int)ctx->max_local_memory);
+
+  w += snprintf(compile_opts+w, compile_opts_size-w,
+                "-D%s=%d ",
+                "max_registers",
+                (int)ctx->max_registers);
+
   for (int i = 0; i < ctx->cfg->num_tuning_params; i++) {
     w += snprintf(compile_opts+w, compile_opts_size-w,
                   "-D%s=%d ",
@@ -606,6 +626,11 @@
                   "%s ", extra_build_opts[i]);
   }
 
+  for (int i = 0; i < num_macros; i++) {
+    w += snprintf(compile_opts+w, compile_opts_size-w,
+                  "-D%s=%zu ", macro_names[i], macro_vals[i]);
+  }
+
   w += snprintf(compile_opts+w, compile_opts_size-w,
                 "-DTR_BLOCK_DIM=%d -DTR_TILE_DIM=%d -DTR_ELEMS_PER_THREAD=%d ",
                 TR_BLOCK_DIM, TR_TILE_DIM, TR_ELEMS_PER_THREAD);
@@ -616,6 +641,9 @@
     w += snprintf(compile_opts+w, compile_opts_size-w, "-DEMULATE_F16 ");
   }
 
+  free(macro_names);
+  free(macro_vals);
+
   return compile_opts;
 }
 
@@ -784,6 +812,20 @@
     }
     ctx->cfg->default_tile_size = max_tile_size;
   }
+
+
+  cl_ulong cache_size;
+  OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_GLOBAL_MEM_CACHE_SIZE,
+                                       sizeof(cache_size), &cache_size, NULL));
+
+  if (cache_size == 0) {
+    // Some code assumes nonzero cache.
+    cache_size = 1024*1024;
+  }
+
+  ctx->max_cache = cache_size;
+
+  ctx->max_registers = 1<<16; // I cannot find a way to query for this.
 
   ctx->max_group_size = max_group_size;
   ctx->max_tile_size = max_tile_size; // No limit.
diff --git a/rts/c/context.h b/rts/c/context.h
--- a/rts/c/context.h
+++ b/rts/c/context.h
@@ -130,8 +130,8 @@
   ctx->profiling_paused = 0;
   ctx->error = NULL;
   ctx->log = stderr;
+  set_tuning_params(ctx);
   if (backend_context_setup(ctx) == 0) {
-    set_tuning_params(ctx);
     setup_program(ctx);
     init_constants(ctx);
     (void)futhark_context_clear_caches(ctx);
diff --git a/rts/c/gpu_prototypes.h b/rts/c/gpu_prototypes.h
--- a/rts/c/gpu_prototypes.h
+++ b/rts/c/gpu_prototypes.h
@@ -5,6 +5,9 @@
 #define TR_TILE_DIM (TR_BLOCK_DIM*2)
 #define TR_ELEMS_PER_THREAD 8
 
+// Must be defined by the user.
+static int gpu_macros(struct futhark_context *ctx, char*** names, int64_t** values);
+
 struct builtin_kernels* init_builtin_kernels(struct futhark_context* ctx);
 void free_builtin_kernels(struct futhark_context* ctx, struct builtin_kernels* kernels);
 static int gpu_free_all(struct futhark_context *ctx);
diff --git a/rts/python/opencl.py b/rts/python/opencl.py
--- a/rts/python/opencl.py
+++ b/rts/python/opencl.py
@@ -102,6 +102,16 @@
     return sizes
 
 
+def to_c_str_rep(x):
+    if type(x) is bool or type(x) is np.bool_:
+        if x:
+            return "true"
+        else:
+            return "false"
+    else:
+        return str(x)
+
+
 def initialise_opencl_object(
     self,
     program_src="",
@@ -119,6 +129,7 @@
     required_types=[],
     all_sizes={},
     user_sizes={},
+    constants=[],
 ):
     if command_queue is None:
         self.ctx = get_prefered_context(
@@ -283,6 +294,10 @@
                 v,
             )
             for (s, v) in self.sizes.items()
+        ]
+
+        build_options += [
+            "-D{}={}".format(s, to_c_str_rep(f())) for (s, f) in constants
         ]
 
         if self.platform.name == "Oclgrind":
diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs
--- a/src/Futhark/Actions.hs
+++ b/src/Futhark/Actions.hs
@@ -166,7 +166,7 @@
   Action
     { actionName = "Compile imperative kernels",
       actionDescription = "Translate program into imperative IL with kernels and write it on standard output.",
-      actionProcedure = liftIO . putStrLn . prettyString . snd <=< ImpGenGPU.compileProgOpenCL
+      actionProcedure = liftIO . putStrLn . prettyString . snd <=< ImpGenGPU.compileProgHIP
     }
 
 -- | Convert the program to CPU multicore ImpCode and print it to stdout.
diff --git a/src/Futhark/CLI/Defs.hs b/src/Futhark/CLI/Defs.hs
--- a/src/Futhark/CLI/Defs.hs
+++ b/src/Futhark/CLI/Defs.hs
@@ -33,7 +33,7 @@
     defsInDec (LocalDec d _) = defsInDec d
     defsInDec (OpenDec me _) = defsInModExp me
     defsInDec (ModDec mb) = defsInModExp $ modExp mb
-    defsInDec SigDec {} = mempty
+    defsInDec ModTypeDec {} = mempty
     defsInDec ImportDec {} = mempty
 
     defsInModExp ModVar {} = mempty
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
@@ -35,7 +35,6 @@
 import Futhark.Optimise.DoubleBuffer
 import Futhark.Optimise.Fusion
 import Futhark.Optimise.HistAccs
-import Futhark.Optimise.InPlaceLowering
 import Futhark.Optimise.InliningDeadFun
 import Futhark.Optimise.MemoryBlockMerging qualified as MemoryBlockMerging
 import Futhark.Optimise.ReduceDeviceSyncs (reduceDeviceSyncs)
@@ -330,23 +329,6 @@
     long = [passLongOption pass]
     pass = Seq.explicitAllocations
 
-iplOption :: String -> FutharkOption
-iplOption short =
-  passOption (passDescription pass) (UntypedPass perform) short long
-  where
-    perform (GPU prog) config =
-      GPU
-        <$> runPipeline (onePass inPlaceLoweringGPU) config prog
-    perform (Seq prog) config =
-      Seq
-        <$> runPipeline (onePass inPlaceLoweringSeq) config prog
-    perform s _ =
-      externalErrorS $
-        "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s
-
-    long = [passLongOption pass]
-    pass = inPlaceLoweringSeq
-
 cseOption :: String -> FutharkOption
 cseOption short =
   passOption (passDescription pass) (UntypedPass perform) short long
@@ -616,7 +598,6 @@
     kernelsPassOption reduceDeviceSyncs [],
     typedPassOption soacsProg GPU extractKernels [],
     typedPassOption soacsProg MC extractMulticore [],
-    iplOption [],
     allocateOption "a",
     kernelsMemPassOption doubleBufferGPU [],
     mcMemPassOption doubleBufferMC [],
diff --git a/src/Futhark/CLI/LSP.hs b/src/Futhark/CLI/LSP.hs
--- a/src/Futhark/CLI/LSP.hs
+++ b/src/Futhark/CLI/LSP.hs
@@ -39,7 +39,7 @@
 syncOptions :: TextDocumentSyncOptions
 syncOptions =
   TextDocumentSyncOptions
-    { _openClose = Just True,
+    { _openClose = Just False,
       _change = Just TextDocumentSyncKind_Incremental,
       _willSave = Just False,
       _willSaveWaitUntil = Just False,
diff --git a/src/Futhark/CLI/Literate.hs b/src/Futhark/CLI/Literate.hs
--- a/src/Futhark/CLI/Literate.hs
+++ b/src/Futhark/CLI/Literate.hs
@@ -45,7 +45,9 @@
   ( copyFile,
     createDirectoryIfMissing,
     doesFileExist,
+    getCurrentDirectory,
     removePathForcibly,
+    setCurrentDirectory,
   )
 import System.Environment (getExecutablePath)
 import System.Exit
@@ -675,18 +677,15 @@
 
 data Env = Env
   { envImgDir :: FilePath,
-    -- | Image dir relative to program.
-    envRelImgDir :: FilePath,
     envOpts :: Options,
     envServer :: ScriptServer,
     envHash :: T.Text
   }
 
-newFileWorker :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM (FilePath, FilePath)
-newFileWorker env (fname_desired, template) m = do
+newFile :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM FilePath
+newFile env (fname_desired, template) m = do
   let fname_base = fromMaybe (T.unpack (envHash env) <> "-" <> template) fname_desired
       fname = envImgDir env </> fname_base
-      fname_rel = envRelImgDir env </> fname_base
   exists <- liftIO $ doesFileExist fname
   liftIO $ createDirectoryIfMissing True $ envImgDir env
   when (exists && scriptVerbose (envOpts env) > 0) $
@@ -698,14 +697,11 @@
         "Generating new file: " <> T.pack fname
     m fname
   modify $ \s -> s {stateFiles = S.insert fname $ stateFiles s}
-  pure (fname, fname_rel)
-
-newFile :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM FilePath
-newFile env f m = snd <$> newFileWorker env f m
+  pure fname
 
 newFileContents :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM T.Text
 newFileContents env f m =
-  liftIO . T.readFile . fst =<< newFileWorker env f m
+  liftIO . T.readFile =<< newFile env f m
 
 processDirective :: Env -> Directive -> ScriptM T.Text
 processDirective env (DirectiveBrief d) =
@@ -1165,8 +1161,8 @@
           system futhark ["hash", prog] mempty
 
       let mdfile = fromMaybe (prog `replaceExtension` "md") $ scriptOutput opts
-          imgdir_rel = dropExtension (takeFileName mdfile) <> "-img"
-          imgdir = takeDirectory mdfile </> imgdir_rel
+          prog_dir = takeDirectory prog
+          imgdir = dropExtension (takeFileName mdfile) <> "-img"
           run_options = scriptExtraOptions opts
           onLine "call" l = T.putStrLn l
           onLine _ _ = pure ()
@@ -1178,16 +1174,22 @@
                     else const . const $ pure ()
               }
 
+      orig_dir <- getCurrentDirectory
+
       withScriptServer cfg $ \server -> do
         let env =
               Env
                 { envServer = server,
                   envOpts = opts,
                   envHash = proghash,
-                  envImgDir = imgdir,
-                  envRelImgDir = imgdir_rel
+                  envImgDir = imgdir
                 }
+
+        when (scriptVerbose opts > 0) $ do
+          T.hPutStrLn stderr $ "Executing from " <> T.pack prog_dir
+        setCurrentDirectory prog_dir
+
         (failure, md) <- processScript env script
-        T.writeFile mdfile md
+        T.writeFile (orig_dir </> mdfile) md
         when (failure == Failure) exitFailure
     _ -> Nothing
diff --git a/src/Futhark/CLI/Misc.hs b/src/Futhark/CLI/Misc.hs
--- a/src/Futhark/CLI/Misc.hs
+++ b/src/Futhark/CLI/Misc.hs
@@ -107,7 +107,9 @@
       [ "You're welcome!",
         "Tell all your friends about Futhark!",
         "Likewise!",
-        "And thank you in return for trying the language!"
+        "And thank you in return for trying the language!",
+        "It's our pleasure!",
+        "Have fun with Futhark!"
       ]
 
 -- | @futhark tokens@
diff --git a/src/Futhark/CodeGen/Backends/CCUDA.hs b/src/Futhark/CodeGen/Backends/CCUDA.hs
--- a/src/Futhark/CodeGen/Backends/CCUDA.hs
+++ b/src/Futhark/CodeGen/Backends/CCUDA.hs
@@ -29,13 +29,15 @@
 
 mkBoilerplate ::
   T.Text ->
+  [(Name, KernelConstExp)] ->
   M.Map Name KernelSafety ->
   [PrimType] ->
   [FailureMsg] ->
   GC.CompilerM OpenCL () ()
-mkBoilerplate cuda_program kernels types failures = do
+mkBoilerplate cuda_program macros kernels types failures = do
   generateGPUBoilerplate
     cuda_program
+    macros
     backendsCudaH
     (M.keys kernels)
     types
@@ -67,7 +69,7 @@
                           fprintf(stderr, "%s: %s\n", optarg, strerror(errno));
                           exit(1);
                         }
-                        exit(1);}|]
+                        exit(0);}|]
            },
          Option
            { optionLongName = "load-cuda",
@@ -113,7 +115,7 @@
 compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts)
 compileProg version prog = do
   ( ws,
-    Program cuda_code cuda_prelude kernels types params failures prog'
+    Program cuda_code cuda_prelude macros kernels types params failures prog'
     ) <-
     ImpGen.compileProg prog
   (ws,)
@@ -122,7 +124,7 @@
       version
       params
       operations
-      (mkBoilerplate (cuda_prelude <> cuda_code) kernels types failures)
+      (mkBoilerplate (cuda_prelude <> cuda_code) macros kernels types failures)
       cuda_includes
       (Space "device", [Space "device", DefaultSpace])
       cliOptions
diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs
--- a/src/Futhark/CodeGen/Backends/COpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/COpenCL.hs
@@ -77,13 +77,15 @@
 
 mkBoilerplate ::
   T.Text ->
+  [(Name, KernelConstExp)] ->
   M.Map Name KernelSafety ->
   [PrimType] ->
   [FailureMsg] ->
   GC.CompilerM OpenCL () ()
-mkBoilerplate opencl_program kernels types failures = do
+mkBoilerplate opencl_program macros kernels types failures = do
   generateGPUBoilerplate
     opencl_program
+    macros
     backendsOpenclH
     (M.keys kernels)
     types
@@ -187,7 +189,7 @@
 compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts)
 compileProg version prog = do
   ( ws,
-    Program opencl_code opencl_prelude kernels types params failures prog'
+    Program opencl_code opencl_prelude macros kernels types params failures prog'
     ) <-
     ImpGen.compileProg prog
   (ws,)
@@ -196,7 +198,7 @@
       version
       params
       operations
-      (mkBoilerplate (opencl_prelude <> opencl_code) kernels types failures)
+      (mkBoilerplate (opencl_prelude <> opencl_code) macros kernels types failures)
       opencl_includes
       (Space "device", [Space "device", DefaultSpace])
       cliOptions
diff --git a/src/Futhark/CodeGen/Backends/GPU.hs b/src/Futhark/CodeGen/Backends/GPU.hs
--- a/src/Futhark/CodeGen/Backends/GPU.hs
+++ b/src/Futhark/CodeGen/Backends/GPU.hs
@@ -12,12 +12,13 @@
 where
 
 import Control.Monad
+import Control.Monad.Identity
 import Data.Bifunctor (bimap)
 import Data.Map qualified as M
 import Data.Text qualified as T
 import Futhark.CodeGen.Backends.GenericC qualified as GC
 import Futhark.CodeGen.Backends.GenericC.Options
-import Futhark.CodeGen.Backends.GenericC.Pretty (idText)
+import Futhark.CodeGen.Backends.GenericC.Pretty (expText, idText)
 import Futhark.CodeGen.Backends.SimpleRep (primStorageType, toStorage)
 import Futhark.CodeGen.ImpCode.OpenCL
 import Futhark.CodeGen.RTS.C (gpuH, gpuPrototypesH)
@@ -66,9 +67,12 @@
           ([C.cinit|&ctx->global_failure_args|], [C.cinit|sizeof(ctx->global_failure_args)|])
         ]
 
+getParamByKey :: Name -> C.Exp
+getParamByKey key = [C.cexp|*ctx->tuning_params.$id:key|]
+
 kernelConstToExp :: KernelConst -> C.Exp
-kernelConstToExp (SizeConst key) =
-  [C.cexp|*ctx->tuning_params.$id:key|]
+kernelConstToExp (SizeConst key _) =
+  getParamByKey key
 kernelConstToExp (SizeMaxConst size_class) =
   [C.cexp|ctx->$id:field|]
   where
@@ -133,13 +137,11 @@
         )
 
 callKernel :: GC.OpCompiler OpenCL ()
-callKernel (GetSize v key) = do
-  let e = kernelConstToExp $ SizeConst key
-  GC.stm [C.cstm|$id:v = $exp:e;|]
+callKernel (GetSize v key) =
+  GC.stm [C.cstm|$id:v = $exp:(getParamByKey key);|]
 callKernel (CmpSizeLe v key x) = do
-  let e = kernelConstToExp $ SizeConst key
   x' <- GC.compileExp x
-  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]
+  GC.stm [C.cstm|$id:v = $exp:(getParamByKey key) <= $exp:x';|]
   -- Output size information if logging is enabled.  The autotuner
   -- depends on the format of this output, so use caution if changing
   -- it.
@@ -381,16 +383,20 @@
                   return strdup("Unknown error.  This is a compiler bug.");
                 }|]
 
+compileConstExp :: KernelConstExp -> C.Exp
+compileConstExp e = runIdentity $ GC.compilePrimExp (pure . kernelConstToExp) e
+
 -- | Called after most code has been generated to generate the bulk of
 -- the boilerplate.
 generateGPUBoilerplate ::
   T.Text ->
+  [(Name, KernelConstExp)] ->
   T.Text ->
   [Name] ->
   [PrimType] ->
   [FailureMsg] ->
   GC.CompilerM OpenCL () ()
-generateGPUBoilerplate gpu_program backendH kernels types failures = do
+generateGPUBoilerplate gpu_program macros backendH kernels types failures = do
   createKernels kernels
   let gpu_program_fragments =
         -- Some C compilers limit the size of literal strings, so
@@ -402,6 +408,13 @@
         | FloatType Float64 `elem` types = [C.cexp|1|]
         | otherwise = [C.cexp|0|]
       max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures
+
+      setMacro i (name, e) =
+        [C.cstm|{names[$int:i] = $string:(nameToString name);
+                 values[$int:i] = $esc:e';}|]
+        where
+          e' = T.unpack $ expText $ compileConstExp e
+
   mapM_
     GC.earlyDecl
     [C.cunit|static const int max_failure_args = $int:max_failure_args;
@@ -410,6 +423,17 @@
              $esc:(T.unpack gpuPrototypesH)
              $esc:(T.unpack backendH)
              $esc:(T.unpack gpuH)
+             static int gpu_macros(struct futhark_context *ctx, char*** names_out, typename int64_t** values_out) {
+               int num_macros = $int:(length macros);
+               char** names = malloc(num_macros * sizeof(char*));
+               typename int64_t* values = malloc(num_macros * sizeof(int64_t));
+
+               $stms:(zipWith setMacro [(0::Int)..] macros)
+
+               *names_out = names;
+               *values_out = values;
+               return num_macros;
+             }
             |]
   GC.earlyDecl $ failureMsgFunction failures
 
diff --git a/src/Futhark/CodeGen/Backends/GenericC/Code.hs b/src/Futhark/CodeGen/Backends/GenericC/Code.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/Code.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/Code.hs
@@ -4,7 +4,6 @@
 module Futhark.CodeGen.Backends.GenericC.Code
   ( compilePrimExp,
     compileExp,
-    compileExpToName,
     compileCode,
     compileDest,
     compileArg,
@@ -44,15 +43,6 @@
       onPart (ErrorVal (FloatType Float64) x) = ("%f",) <$> compileExp x
   (formatstrs, formatargs) <- mapAndUnzipM onPart parts
   pure (mconcat formatstrs, formatargs)
-
-compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName
-compileExpToName _ _ (LeafExp v _) =
-  pure v
-compileExpToName desc t e = do
-  desc' <- newVName desc
-  e' <- compileExp e
-  decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|]
-  pure desc'
 
 compileExp :: Exp -> CompilerM op s C.Exp
 compileExp = compilePrimExp $ \v -> pure [C.cexp|$id:v|]
diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs
--- a/src/Futhark/CodeGen/Backends/GenericPython.hs
+++ b/src/Futhark/CodeGen/Backends/GenericPython.hs
@@ -14,7 +14,6 @@
     compileCode,
     compilePrimValue,
     compilePrimType,
-    compilePrimTypeExt,
     compilePrimToNp,
     compilePrimToExtNp,
     fromStorage,
@@ -1026,24 +1025,6 @@
     FloatType Float64 -> "ct.c_double"
     Imp.Bool -> "ct.c_bool"
     Unit -> "ct.c_bool"
-
--- | The ctypes type corresponding to a 'PrimType', taking sign into account.
-compilePrimTypeExt :: PrimType -> Imp.Signedness -> String
-compilePrimTypeExt t ept =
-  case (t, ept) of
-    (IntType Int8, Imp.Unsigned) -> "ct.c_uint8"
-    (IntType Int16, Imp.Unsigned) -> "ct.c_uint16"
-    (IntType Int32, Imp.Unsigned) -> "ct.c_uint32"
-    (IntType Int64, Imp.Unsigned) -> "ct.c_uint64"
-    (IntType Int8, _) -> "ct.c_int8"
-    (IntType Int16, _) -> "ct.c_int16"
-    (IntType Int32, _) -> "ct.c_int32"
-    (IntType Int64, _) -> "ct.c_int64"
-    (FloatType Float16, _) -> "ct.c_uint16"
-    (FloatType Float32, _) -> "ct.c_float"
-    (FloatType Float64, _) -> "ct.c_double"
-    (Imp.Bool, _) -> "ct.c_bool"
-    (Unit, _) -> "ct.c_byte"
 
 -- | The Numpy type corresponding to a 'PrimType'.
 compilePrimToNp :: Imp.PrimType -> String
diff --git a/src/Futhark/CodeGen/Backends/HIP.hs b/src/Futhark/CodeGen/Backends/HIP.hs
--- a/src/Futhark/CodeGen/Backends/HIP.hs
+++ b/src/Futhark/CodeGen/Backends/HIP.hs
@@ -29,13 +29,15 @@
 
 mkBoilerplate ::
   T.Text ->
+  [(Name, KernelConstExp)] ->
   M.Map Name KernelSafety ->
   [PrimType] ->
   [FailureMsg] ->
   GC.CompilerM OpenCL () ()
-mkBoilerplate hip_program kernels types failures = do
+mkBoilerplate hip_program macros kernels types failures = do
   generateGPUBoilerplate
     hip_program
+    macros
     backendsHipH
     (M.keys kernels)
     types
@@ -95,7 +97,7 @@
 compileProg :: (MonadFreshNames m) => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts)
 compileProg version prog = do
   ( ws,
-    Program hip_code hip_prelude kernels types params failures prog'
+    Program hip_code hip_prelude macros kernels types params failures prog'
     ) <-
     ImpGen.compileProg prog
   (ws,)
@@ -104,7 +106,7 @@
       version
       params
       operations
-      (mkBoilerplate (hip_prelude <> hip_code) kernels types failures)
+      (mkBoilerplate (hip_prelude <> hip_code) macros kernels types failures)
       hip_includes
       (Space "device", [Space "device", DefaultSpace])
       cliOptions
diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
--- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
@@ -33,6 +33,7 @@
     Imp.Program
       opencl_code
       opencl_prelude
+      macros
       kernels
       types
       sizes
@@ -89,7 +90,7 @@
             "default_threshold=default_threshold",
             "sizes=sizes"
           ]
-          [Escape $ openClInit types assign sizes failures]
+          [Escape $ openClInit macros types assign sizes failures]
       options =
         [ Option
             { optionLongName = "platform",
@@ -202,9 +203,12 @@
 asLong :: PyExp -> PyExp
 asLong x = simpleCall "np.int64" [x]
 
+getParamByKey :: Name -> PyExp
+getParamByKey key = Index (Var "self.sizes") (IdxExp $ String $ prettyText key)
+
 kernelConstToExp :: Imp.KernelConst -> PyExp
-kernelConstToExp (Imp.SizeConst key) =
-  Index (Var "self.sizes") (IdxExp $ String $ prettyText key)
+kernelConstToExp (Imp.SizeConst key _) =
+  getParamByKey key
 kernelConstToExp (Imp.SizeMaxConst size_class) =
   Var $ "self.max_" <> prettyString size_class
 
@@ -215,13 +219,11 @@
 callKernel :: OpCompiler Imp.OpenCL ()
 callKernel (Imp.GetSize v key) = do
   v' <- compileVar v
-  stm $ Assign v' $ kernelConstToExp $ Imp.SizeConst key
+  stm $ Assign v' $ getParamByKey key
 callKernel (Imp.CmpSizeLe v key x) = do
   v' <- compileVar v
   x' <- compileExp x
-  stm $
-    Assign v' $
-      BinOp "<=" (kernelConstToExp (Imp.SizeConst key)) x'
+  stm $ Assign v' $ BinOp "<=" (getParamByKey key) x'
 callKernel (Imp.GetSizeMax v size_class) = do
   v' <- compileVar v
   stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class
diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs
@@ -15,6 +15,8 @@
   ( ErrorMsg (..),
     ErrorMsgPart (..),
     FailureMsg (..),
+    KernelConst (..),
+    KernelConstExp,
     ParamMap,
     PrimType (..),
     errorMsgArgTypes,
@@ -28,15 +30,28 @@
 errorMsgNumArgs :: ErrorMsg a -> Int
 errorMsgNumArgs = length . errorMsgArgTypes
 
+getParamByKey :: Name -> PyExp
+getParamByKey key = Index (Var "self.sizes") (IdxExp $ String $ prettyText key)
+
+kernelConstToExp :: KernelConst -> PyExp
+kernelConstToExp (SizeConst key _) =
+  getParamByKey key
+kernelConstToExp (SizeMaxConst size_class) =
+  Var $ "self.max_" <> prettyString size_class
+
+compileConstExp :: KernelConstExp -> PyExp
+compileConstExp e = runIdentity $ Py.compilePrimExp (pure . kernelConstToExp) e
+
 -- | Python code (as a string) that calls the
 -- @initiatialize_opencl_object@ procedure.  Should be put in the
 -- class constructor.
-openClInit :: [PrimType] -> String -> ParamMap -> [FailureMsg] -> T.Text
-openClInit types assign sizes failures =
+openClInit :: [(Name, KernelConstExp)] -> [PrimType] -> String -> ParamMap -> [FailureMsg] -> T.Text
+openClInit constants types assign sizes failures =
   [text|
 size_heuristics=$size_heuristics
 self.global_failure_args_max = $max_num_args
 self.failure_msgs=$failure_msgs
+constants = $constants'
 program = initialise_opencl_object(self,
                                    program_src=fut_opencl_src,
                                    build_options=build_options,
@@ -52,7 +67,8 @@
                                    size_heuristics=size_heuristics,
                                    required_types=$types',
                                    user_sizes=sizes,
-                                   all_sizes=$sizes')
+                                   all_sizes=$sizes',
+                                   constants=constants)
 $assign'
 |]
   where
@@ -62,6 +78,12 @@
     sizes' = prettyText $ sizeClassesToPython sizes
     max_num_args = prettyText $ foldl max 0 $ map (errorMsgNumArgs . failureError) failures
     failure_msgs = prettyText $ List $ map formatFailure failures
+    onConstant (name, e) =
+      Tuple
+        [ String (nameToText name),
+          Lambda "" (compileConstExp e)
+        ]
+    constants' = prettyText $ List $ map onConstant constants
 
 formatFailure :: FailureMsg -> PyExp
 formatFailure (FailureMsg (ErrorMsg parts) backtrace) =
diff --git a/src/Futhark/CodeGen/ImpCode/GPU.hs b/src/Futhark/CodeGen/ImpCode/GPU.hs
--- a/src/Futhark/CodeGen/ImpCode/GPU.hs
+++ b/src/Futhark/CodeGen/ImpCode/GPU.hs
@@ -34,7 +34,7 @@
 
 -- | A run-time constant related to kernels.
 data KernelConst
-  = SizeConst Name
+  = SizeConst Name SizeClass
   | SizeMaxConst SizeClass
   deriving (Eq, Ord, Show)
 
@@ -62,15 +62,14 @@
     -- | A short descriptive and _unique_ name - should be
     -- alphanumeric and without spaces.
     kernelName :: Name,
-    -- | If true, this kernel does not need to check
-    -- whether we are in a failing state, as it can cope.
-    -- Intuitively, it means that the kernel does not
-    -- depend on any non-scalar parameters to make control
-    -- flow decisions.  Replication, transpose, and copy
+    -- | If true, this kernel does not need to check whether we are in
+    -- a failing state, as it can cope. Intuitively, it means that the
+    -- kernel does not depend on any non-scalar parameters to make
+    -- control flow decisions. Replication, transpose, and copy
     -- kernels are examples of this.
     kernelFailureTolerant :: Bool,
     -- | If true, multi-versioning branches will consider this kernel
-    -- when considering the local memory requirements.  Set this to
+    -- when considering the local memory requirements. Set this to
     -- false for kernels that do their own checking.
     kernelCheckLocalMemory :: Bool
   }
@@ -86,11 +85,13 @@
   deriving (Eq, Ord, Show)
 
 instance Pretty KernelConst where
-  pretty (SizeConst key) = "get_size" <> parens (pretty key)
-  pretty (SizeMaxConst size_class) = "get_max_size" <> parens (pretty size_class)
+  pretty (SizeConst key size_class) =
+    "get_size" <> parens (commasep [pretty key, pretty size_class])
+  pretty (SizeMaxConst size_class) =
+    "get_max_size" <> parens (pretty size_class)
 
 instance FreeIn KernelConst where
-  freeIn' (SizeConst _) = mempty
+  freeIn' SizeConst {} = mempty
   freeIn' (SizeMaxConst _) = mempty
 
 instance Pretty KernelUse where
diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs
--- a/src/Futhark/CodeGen/ImpCode/OpenCL.hs
+++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs
@@ -18,6 +18,7 @@
     FailureMsg (..),
     GroupDim,
     KernelConst (..),
+    KernelConstExp,
     module Futhark.CodeGen.ImpCode,
     module Futhark.IR.GPU.Sizes,
   )
@@ -26,7 +27,7 @@
 import Data.Map qualified as M
 import Data.Text qualified as T
 import Futhark.CodeGen.ImpCode
-import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..))
+import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..), KernelConstExp)
 import Futhark.IR.GPU.Sizes
 import Futhark.Util.Pretty
 
@@ -35,6 +36,9 @@
   { openClProgram :: T.Text,
     -- | Must be prepended to the program.
     openClPrelude :: T.Text,
+    -- | Definitions to be passed as macro definitions to the kernel
+    -- compiler.
+    openClMacroDefs :: [(Name, KernelConstExp)],
     openClKernelNames :: M.Map KernelName KernelSafety,
     -- | So we can detect whether the device is capable.
     openClUsedTypes :: [PrimType],
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
@@ -38,7 +38,6 @@
     hasFunction,
     collect,
     collect',
-    comment,
     VarEntry (..),
     ArrayEntry (..),
 
@@ -76,6 +75,7 @@
     caseMatch,
 
     -- * Constructing code.
+    newVName,
     dLParams,
     dFParams,
     addLoopVar,
@@ -391,13 +391,6 @@
   modify $ \s -> s {stateCode = prev_code}
   pure (x, new_code)
 
--- | Execute a code generation action, wrapping the generated code
--- within a 'Imp.Comment' with the given description.
-comment :: T.Text -> ImpM rep r op () -> ImpM rep r op ()
-comment desc m = do
-  code <- collect m
-  emit $ Imp.Comment desc code
-
 -- | Emit some generated imperative code.
 emit :: Imp.Code op -> ImpM rep r op ()
 emit code = modify $ \s -> s {stateCode = stateCode s <> code}
@@ -870,7 +863,7 @@
   copyDWIM (patElemName pe) [] se []
   case op of
     OpaqueNil -> pure ()
-    OpaqueTrace s -> comment ("Trace: " <> s) $ do
+    OpaqueTrace s -> sComment ("Trace: " <> s) $ do
       se_t <- subExpType se
       case se_t of
         Prim t -> tracePrim s t se
@@ -937,7 +930,7 @@
       dPrimV "x" . TPrimExp $
         BinOpExp (Add it OverflowUndef) e' $
           BinOpExp (Mul it OverflowUndef) i' s'
-    copyDWIM (patElemName pe) [DimFix i] (Var (tvVar x)) []
+    copyDWIMFix (patElemName pe) [i] (Var (tvVar x)) []
 defCompileBasicOp (Pat [pe]) (Manifest _ src) =
   copyDWIM (patElemName pe) [] (Var src) []
 defCompileBasicOp (Pat [pe]) (Concat i (x :| ys) _) = do
@@ -967,7 +960,7 @@
       copy t dest_mem static_src
   | otherwise =
       forM_ (zip [0 ..] es) $ \(i, e) ->
-        copyDWIM (patElemName pe) [DimFix $ fromInteger i] e []
+        copyDWIMFix (patElemName pe) [fromInteger i] e []
   where
     isLiteral (Constant v) = Just v
     isLiteral _ = Nothing
@@ -1665,6 +1658,8 @@
   body' <- collect body
   emit $ Imp.While cond body'
 
+-- | Execute a code generation action, wrapping the generated code
+-- within a 'Imp.Comment' with the given description.
 sComment :: T.Text -> ImpM rep r op () -> ImpM rep r op ()
 sComment s code = do
   code' <- collect code
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
@@ -27,7 +27,7 @@
 import Futhark.Error
 import Futhark.IR.GPUMem
 import Futhark.MonadFreshNames
-import Futhark.Util.IntegralExp (divUp, rem)
+import Futhark.Util.IntegralExp (divUp, nextMul)
 import Prelude hiding (quot, rem)
 
 callKernelOperations :: Operations GPUMem HostEnv Imp.HostOp
@@ -134,7 +134,7 @@
   sKernelThread "gpuseq" tid (defKernelAttrs one one) $
     compileStms (freeIn res) stms $
       forM_ (zip pes res) $ \(pe, SubExpRes _ se) ->
-        copyDWIM (patElemName pe) [DimFix 0] se []
+        copyDWIMFix (patElemName pe) [0] se []
 opCompiler pat e =
   compilerBugS $
     "opCompiler: Invalid pattern\n  "
@@ -170,14 +170,13 @@
 -- otherwise protected by their own multi-versioning branches deeper
 -- down.  Currently the compiler will not generate multi-versioning
 -- that makes this a problem, but it might in the future.
-checkLocalMemoryReqs :: Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))
-checkLocalMemoryReqs code = do
-  scope <- askScope
+checkLocalMemoryReqs :: (VName -> Bool) -> Imp.HostCode -> CallKernelGen (Maybe (Imp.TExp Bool))
+checkLocalMemoryReqs in_scope code = do
   let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getGPU code
 
   -- If any of the sizes involve a variable that is not known at this
   -- point, then we cannot check the requirements.
-  if any (`M.notMember` scope) (namesToList $ freeIn alloc_sizes)
+  if not $ all in_scope $ namesToList $ freeIn alloc_sizes
     then pure Nothing
     else do
       local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32
@@ -200,7 +199,7 @@
     -- These allocations will actually be padded to an 8-byte aligned
     -- size, so we should take that into account when checking whether
     -- they fit.
-    alignedSize x = x + ((8 - (x `rem` 8)) `rem` 8)
+    alignedSize x = nextMul x 8
 
 withAcc ::
   Pat LetDecMem ->
@@ -250,9 +249,10 @@
 -- always be safe (and what would we do if none of the branches would
 -- work?).
 expCompiler dest (Match cond (first_case : cases) defbranch sort@(MatchDec _ MatchEquiv)) = do
+  scope <- askScope
   tcode <- collect $ compileBody dest $ caseBody first_case
   fcode <- collect $ expCompiler dest $ Match cond cases defbranch sort
-  check <- checkLocalMemoryReqs tcode
+  check <- checkLocalMemoryReqs (`M.member` scope) tcode
   let matches = caseMatch cond (casePat first_case)
   emit $ case check of
     Nothing -> fcode
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
@@ -29,6 +29,9 @@
     fenceForArrays,
     updateAcc,
     genZeroes,
+    isPrimParam,
+    kernelConstToExp,
+    getChunkSize,
 
     -- * Host-level bulk operations
     sReplicate,
@@ -52,7 +55,6 @@
 import Futhark.Error
 import Futhark.IR.GPUMem
 import Futhark.IR.Mem.LMAD qualified as LMAD
-import Futhark.MonadFreshNames
 import Futhark.Transform.Rename
 import Futhark.Util (dropLast, nubOrd, splitFromEnd)
 import Futhark.Util.IntegralExp (divUp, quot, rem)
@@ -257,6 +259,44 @@
         . entryArrayLoc
         =<< lookupArray arr
 
+isPrimParam :: (Typed p) => Param p -> Bool
+isPrimParam = primType . paramType
+
+kernelConstToExp :: Imp.KernelConstExp -> CallKernelGen Imp.Exp
+kernelConstToExp = traverse f
+  where
+    f (Imp.SizeMaxConst c) = do
+      v <- dPrim (prettyString c) int64
+      sOp $ Imp.GetSizeMax (tvVar v) c
+      pure $ tvVar v
+    f (Imp.SizeConst k c) = do
+      v <- dPrim (nameToString k) int64
+      sOp $ Imp.GetSize (tvVar v) k c
+      pure $ tvVar v
+
+-- | Given available register and a list of parameter types, compute
+-- the largest available chunk size given the parameters for which we
+-- want chunking and the available resources. Used in
+-- 'SegScan.SinglePass.compileSegScan', and 'SegRed.compileSegRed'
+-- (with primitive non-commutative operators only).
+getChunkSize :: [Type] -> Imp.KernelConstExp
+getChunkSize types = do
+  let max_group_size = Imp.SizeMaxConst SizeGroup
+      max_group_mem = Imp.SizeMaxConst SizeLocalMemory
+      max_group_reg = Imp.SizeMaxConst SizeRegisters
+      k_mem = le64 max_group_mem `quot` le64 max_group_size
+      k_reg = le64 max_group_reg `quot` le64 max_group_size
+      types' = map elemType $ filter primType types
+      sizes = map primByteSize types'
+
+      sum_sizes = sum sizes
+      sum_sizes' = sum (map (sMax64 4 . primByteSize) types') `quot` 4
+      max_size = maximum sizes
+
+      mem_constraint = max k_mem sum_sizes `quot` max_size
+      reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes')
+  untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint
+
 inBlockScan ::
   KernelConstants ->
   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->
@@ -275,7 +315,7 @@
         splitAt (length actual_params `div` 2) actual_params
       y_to_x =
         forM_ (zip x_params y_params) $ \(x, y) ->
-          when (primType (paramType x)) $
+          when (isPrimParam x) $
             copyDWIM (paramName x) [] (Var (paramName y)) []
 
   -- Set initial y values
@@ -342,23 +382,21 @@
     array_scan = not $ all primType $ lambdaReturnType scan_lam
 
     readInitial p arr
-      | primType $ paramType p =
-          copyDWIM (paramName p) [] (Var arr) [DimFix ltid]
+      | isPrimParam p =
+          copyDWIMFix (paramName p) [] (Var arr) [ltid]
       | otherwise =
-          copyDWIM (paramName p) [] (Var arr) [DimFix gtid]
+          copyDWIMFix (paramName p) [] (Var arr) [gtid]
 
     readParam behind p arr
-      | primType $ paramType p =
-          copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind]
+      | isPrimParam p =
+          copyDWIMFix (paramName p) [] (Var arr) [ltid - behind]
       | otherwise =
-          copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size]
+          copyDWIMFix (paramName p) [] (Var arr) [gtid - behind + arrs_full_size]
 
-    writeResult x y arr
-      | primType $ paramType x = do
-          copyDWIM arr [DimFix ltid] (Var $ paramName x) []
-          copyDWIM (paramName y) [] (Var $ paramName x) []
-      | otherwise =
-          copyDWIM (paramName y) [] (Var $ paramName x) []
+    writeResult x y arr = do
+      when (isPrimParam x) $
+        copyDWIMFix arr [ltid] (Var $ paramName x) []
+      copyDWIM (paramName y) [] (Var $ paramName x) []
 
 groupScan ::
   Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) ->
@@ -421,16 +459,16 @@
       group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants
 
       writeBlockResult p arr
-        | primType $ paramType p =
-            copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) []
+        | isPrimParam p =
+            copyDWIMFix arr [sExt64 block_id] (Var $ paramName p) []
         | otherwise =
-            copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) []
+            copyDWIMFix arr [group_offset + sExt64 block_id] (Var $ paramName p) []
 
       readPrevBlockResult p arr
-        | primType $ paramType p =
-            copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1]
+        | isPrimParam p =
+            copyDWIMFix (paramName p) [] (Var arr) [sExt64 block_id - 1]
         | otherwise =
-            copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1]
+            copyDWIMFix (paramName p) [] (Var arr) [group_offset + sExt64 block_id - 1]
 
   doInBlockScan seg_flag ltid_in_bounds lam
   barrier
@@ -440,8 +478,8 @@
     sComment "save correct values for first block" $
       sWhen is_first_block $
         forM_ (zip x_params arrs) $ \(x, arr) ->
-          unless (primType $ paramType x) $
-            copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []
+          unless (isPrimParam x) $
+            copyDWIMFix arr [arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) []
 
     barrier
 
@@ -457,7 +495,7 @@
         flag_true <- seg_flag
         Just $ \from to ->
           flag_true (from * block_size + block_size - 1) (to * block_size + block_size - 1)
-  comment
+  sComment
     "scan the first block, after which offset 'i' contains carry-in for block 'i+1'"
     $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam
 
@@ -467,12 +505,12 @@
     sComment "move correct values for first block back a block" $
       sWhen is_first_block $
         forM_ (zip x_params arrs) $ \(x, arr) ->
-          unless (primType $ paramType x) $
-            copyDWIM
+          unless (isPrimParam x) $
+            copyDWIMFix
               arr
-              [DimFix $ arrs_full_size + group_offset + ltid]
+              [arrs_full_size + group_offset + ltid]
               (Var arr)
-              [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid]
+              [arrs_full_size + group_offset + sExt64 block_size + ltid]
 
     barrier
 
@@ -498,8 +536,8 @@
 
       write_final_result =
         forM_ (zip x_params arrs) $ \(p, arr) ->
-          when (primType $ paramType p) $
-            copyDWIM arr [DimFix ltid] (Var $ paramName p) []
+          when (isPrimParam p) $
+            copyDWIMFix arr [ltid] (Var $ paramName p) []
 
   sComment "carry-in for every block except the first" $
     localOps threadOperations $ do
@@ -512,9 +550,9 @@
   sComment "restore correct values for first block" $
     sWhen (is_first_block .&&. ltid_in_bounds) $
       forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) ->
-        if primType (paramType y)
-          then copyDWIM arr [DimFix ltid] (Var $ paramName y) []
-          else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid]
+        if isPrimParam y
+          then copyDWIMFix arr [ltid] (Var $ paramName y) []
+          else copyDWIMFix (paramName x) [] (Var arr) [arrs_full_size + group_offset + ltid]
 
   barrier
 
@@ -550,17 +588,13 @@
         let i = local_tid + tvExp offset
         copyDWIMFix (paramName param) [] (Var arr) [sExt64 i]
 
-      writeReduceOpResult param arr
-        | Prim _ <- paramType param =
-            copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []
-        | otherwise =
-            pure ()
+      writeReduceOpResult param arr =
+        when (isPrimParam param) $
+          copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []
 
-      writeArrayOpResult param arr
-        | Prim _ <- paramType param =
-            pure ()
-        | otherwise =
-            copyDWIMFix arr [0] (Var $ paramName param) []
+      writeArrayOpResult param arr =
+        unless (isPrimParam param) $
+          copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) []
 
   let (reduce_acc_params, reduce_arr_params) =
         splitAt (length arrs) $ lambdaParams lam
@@ -570,17 +604,17 @@
 
   offset <-- (0 :: Imp.TExp Int32)
 
-  comment "participating threads read initial accumulator" $
+  sComment "participating threads read initial accumulator" $
     localOps threadOperations . sWhen (local_tid .<. w) $
       zipWithM_ readReduceArgument reduce_acc_params arrs
 
   let do_reduce = localOps threadOperations $ do
-        comment "read array element" $
+        sComment "read array element" $
           zipWithM_ readReduceArgument reduce_arr_params arrs
-        comment "apply reduction operation" $
+        sComment "apply reduction operation" $
           compileBody' reduce_acc_params $
             lambdaBody lam
-        comment "write result of operation" $
+        sComment "write result of operation" $
           zipWithM_ writeReduceOpResult reduce_acc_params arrs
       in_wave_reduce = everythingVolatile do_reduce
 
@@ -622,10 +656,11 @@
   cross_wave_reductions
   errorsync
 
-  sComment "Copy array-typed operands to result array" $ do
-    sWhen (local_tid .==. 0) $
-      localOps threadOperations $
-        zipWithM_ writeArrayOpResult reduce_acc_params arrs
+  unless (all isPrimParam reduce_acc_params) $
+    sComment "Copy array-typed operands to result array" $
+      sWhen (local_tid .==. 0) $
+        localOps threadOperations $
+          zipWithM_ writeArrayOpResult reduce_acc_params arrs
 
 compileThreadOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp
 compileThreadOp pat (Alloc size space) =
@@ -900,10 +935,10 @@
   let onLeaf name _ = lookupConstExp name
       lookupConstExp name =
         constExp =<< hasExp =<< M.lookup name vtable
-      constExp (Op (Inner (SizeOp (GetSize key _)))) =
-        Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32
-      constExp (Op (Inner (SizeOp (GetSizeMax size_class)))) =
-        Just $ LeafExp (Imp.SizeMaxConst size_class) int32
+      constExp (Op (Inner (SizeOp (GetSize key c)))) =
+        Just $ LeafExp (Imp.SizeConst (keyWithEntryPoint fname key) c) int32
+      constExp (Op (Inner (SizeOp (GetSizeMax c)))) =
+        Just $ LeafExp (Imp.SizeMaxConst c) int32
       constExp e = primExpFromExp lookupConstExp e
   pure $ replaceInPrimExpM onLeaf size
   where
@@ -1092,7 +1127,12 @@
     -- | Number of groups.
     kAttrNumGroups :: Count NumGroups SubExp,
     -- | Group size.
-    kAttrGroupSize :: Count GroupSize SubExp
+    kAttrGroupSize :: Count GroupSize SubExp,
+    -- | Variables that are specially in scope inside the kernel.
+    -- Operationally, these will be available at kernel compile time
+    -- (which happens at run-time, with access to machine-specific
+    -- information).
+    kAttrConstExps :: M.Map VName Imp.KernelConstExp
   }
 
 -- | The default kernel attributes.
@@ -1105,7 +1145,8 @@
     { kAttrFailureTolerant = False,
       kAttrCheckLocalMemory = True,
       kAttrNumGroups = num_groups,
-      kAttrGroupSize = group_size
+      kAttrGroupSize = group_size,
+      kAttrConstExps = mempty
     }
 
 getSize :: String -> SizeClass -> CallKernelGen (TV Int64)
@@ -1170,12 +1211,12 @@
 sKernelOp attrs constants ops name m = do
   HostEnv atomics _ locks <- askEnv
   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m
-  uses <- computeKernelUses body mempty
+  uses <- computeKernelUses body $ M.keys $ kAttrConstExps attrs
   group_size <- onGroupSize $ kernelGroupSize constants
   emit . Imp.Op . Imp.CallKernel $
     Imp.Kernel
       { Imp.kernelBody = body,
-        Imp.kernelUses = uses,
+        Imp.kernelUses = uses <> map constToUse (M.toList (kAttrConstExps attrs)),
         Imp.kernelNumGroups = [untyped $ kernelNumGroups constants],
         Imp.kernelGroupSize = [group_size],
         Imp.kernelName = name,
@@ -1192,6 +1233,8 @@
         case x of
           Just (LeafExp kc _) -> Right kc
           _ -> Left $ untyped e
+
+    constToUse (v, e) = Imp.ConstUse v e
 
 sKernelFailureTolerant ::
   Bool ->
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
@@ -27,7 +27,6 @@
 import Futhark.Error
 import Futhark.IR.GPUMem
 import Futhark.IR.Mem.LMAD qualified as LMAD
-import Futhark.MonadFreshNames
 import Futhark.Transform.Rename
 import Futhark.Util (chunks, mapAccumLM, takeLast)
 import Futhark.Util.IntegralExp (divUp, rem)
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
@@ -48,8 +48,8 @@
 import Futhark.Construct (fullSliceNum)
 import Futhark.IR.GPUMem
 import Futhark.IR.Mem.LMAD qualified as LMAD
-import Futhark.MonadFreshNames
 import Futhark.Pass.ExplicitAllocations ()
+import Futhark.Transform.Substitute
 import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast)
 import Futhark.Util.IntegralExp (divUp, quot, rem)
 import Prelude hiding (quot, rem)
@@ -242,17 +242,9 @@
           t64 $
             r64 hist_T / hist_C_max
 
-  -- Querying L2 cache size is not reliable.  Instead we provide a
-  -- tunable knob with a hopefully sane default.
-  let hist_L2_def = 4 * 1024 * 1024
   hist_L2 <- dPrim "L2_size" int32
-  entry <- askFunction
   -- Equivalent to F_L2*L2 in paper.
-  sOp
-    $ Imp.GetSize
-      (tvVar hist_L2)
-      (keyWithEntryPoint entry $ nameFromString (prettyString (tvVar hist_L2)))
-    $ Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def
+  sOp $ Imp.GetSizeMax (tvVar hist_L2) Imp.SizeCache
 
   let hist_L2_ln_sz = 16 * 4 -- L2 cache line size approximation
   hist_RACE_exp <-
@@ -1073,7 +1065,8 @@
   KernelBody GPUMem ->
   CallKernelGen ()
 compileSegHist (Pat pes) lvl space ops kbody = do
-  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl
+  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = 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
@@ -1180,8 +1173,11 @@
                   ++ zip bucket_ids (shapeDims (histShape op))
                   ++ zip vector_ids (shapeDims $ histOpShape op)
                   ++ [(subhistogram_id, Var $ tvVar num_histos)]
+            -- The operator may have references to the old flat thread
+            -- ID, which we must update to point at the new one.
+            subst = M.singleton (segFlat space) flat_gtid
 
-        let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty
+        let segred_op = SegBinOp Commutative (substituteNames subst $ histOp op) (histNeutral op) mempty
         compileSegRed' (Pat red_pes) grid segred_space [segred_op] $ \red_cont ->
           red_cont . flip map subhistos $ \subhisto ->
             ( Var subhisto,
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
@@ -41,845 +41,1035 @@
 --   cases, we would allocate a /whole/ group per segment with the
 --   large strategy, but at most 50% of the threads in the group would
 --   have any element to read, which becomes highly inefficient.
-module Futhark.CodeGen.ImpGen.GPU.SegRed
-  ( compileSegRed,
-    compileSegRed',
-    DoSegBody,
-  )
-where
-
-import Control.Monad
-import Data.List (genericLength, zip7)
-import Data.Maybe
-import Futhark.CodeGen.ImpCode.GPU qualified as Imp
-import Futhark.CodeGen.ImpGen
-import Futhark.CodeGen.ImpGen.GPU.Base
-import Futhark.Error
-import Futhark.IR.GPUMem
-import Futhark.IR.Mem.LMAD qualified as LMAD
-import Futhark.Transform.Rename
-import Futhark.Util (chunks)
-import Futhark.Util.IntegralExp (divUp, quot, rem)
-import Prelude hiding (quot, rem)
-
--- | The maximum number of operators we support in a single SegRed.
--- This limit arises out of the static allocation of counters.
-maxNumOps :: Int32
-maxNumOps = 10
-
--- | Code generation for the body of the SegRed, taking a continuation
--- for saving the results of the body.  The results should be
--- represented as a pairing of a t'SubExp' along with a list of
--- indexes into that t'SubExp' for reading the result.
-type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen ()
-
--- | Compile 'SegRed' instance to host-level code with calls to
--- various kernels.
-compileSegRed ::
-  Pat LetDecMem ->
-  SegLevel ->
-  SegSpace ->
-  [SegBinOp GPUMem] ->
-  KernelBody GPUMem ->
-  CallKernelGen ()
-compileSegRed pat lvl space reds body = do
-  emit $ Imp.DebugPrint "\n# SegRed" Nothing
-  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
-
-      sComment "save map-out results" $ do
-        let map_arrs = drop (segBinOpResults reds) $ patElems pat
-        zipWithM_ (compileThreadResult space) map_arrs map_res
-
-      red_cont $ map ((,[]) . kernelResultSubExp) red_res
-  emit $ Imp.DebugPrint "" Nothing
-
--- | Like 'compileSegRed', but where the body is a monadic action.
-compileSegRed' ::
-  Pat LetDecMem ->
-  KernelGrid ->
-  SegSpace ->
-  [SegBinOp GPUMem] ->
-  DoSegBody ->
-  CallKernelGen ()
-compileSegRed' pat grid space reds body
-  | genericLength reds > maxNumOps =
-      compilerLimitationS $
-        "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."
-  | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space =
-      nonsegmentedReduction pat num_groups group_size space reds body
-  | otherwise = do
-      let group_size' = pe64 $ unCount group_size
-          segment_size = pe64 $ last $ segSpaceDims space
-          use_small_segments = segment_size * 2 .<. group_size'
-      sIf
-        use_small_segments
-        (smallSegmentsReduction pat num_groups group_size space reds body)
-        (largeSegmentsReduction pat num_groups group_size space reds body)
-  where
-    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
--- those arrays inside the kernel), while array-typed arguments go in
--- global memory.  Allocations for the former have already been
--- performed.  This policy is baked into how the allocations are done
--- in ExplicitAllocations.
-intermediateArrays ::
-  Count GroupSize SubExp ->
-  SubExp ->
-  SegBinOp GPUMem ->
-  InKernelGen [VName]
-intermediateArrays (Count group_size) num_threads (SegBinOp _ red_op nes _) = do
-  let red_op_params = lambdaParams red_op
-      (red_acc_params, _) = splitAt (length nes) red_op_params
-  forM red_acc_params $ \p ->
-    case paramDec p of
-      MemArray pt shape _ (ArrayIn mem _) -> do
-        let shape' = Shape [num_threads] <> shape
-        sArray "red_arr" pt shape' mem $
-          LMAD.iota 0 (map pe64 $ shapeDims shape')
-      _ -> do
-        let pt = elemType $ paramType p
-            shape = Shape [group_size]
-        sAllocArray "red_arr" pt shape $ Space "local"
-
--- | Arrays for storing group results.
---
--- The group-result arrays have an extra dimension because they are
--- also used for keeping vectorised accumulators for first-stage
--- reduction, if necessary.  If necessary, this dimension has size
--- group_size, and otherwise 1.  When actually storing group results,
--- the first index is set to 0.
-groupResultArrays ::
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
-  [SegBinOp GPUMem] ->
-  CallKernelGen [[VName]]
-groupResultArrays (Count virt_num_groups) (Count group_size) reds =
-  forM reds $ \(SegBinOp _ lam _ shape) ->
-    forM (lambdaReturnType lam) $ \t -> do
-      let pt = elemType t
-          extra_dim
-            | primType t, shapeRank shape == 0 = intConst Int64 1
-            | otherwise = group_size
-          full_shape = Shape [extra_dim, virt_num_groups] <> shape <> arrayShape t
-          -- Move the groupsize dimension last to ensure coalesced
-          -- memory access.
-          perm = [1 .. shapeRank full_shape - 1] ++ [0]
-      sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm
-
-nonsegmentedReduction ::
-  Pat LetDecMem ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
-  SegSpace ->
-  [SegBinOp GPUMem] ->
-  DoSegBody ->
-  CallKernelGen ()
-nonsegmentedReduction segred_pat num_groups group_size space reds body = do
-  let (gtids, dims) = unzip $ unSegSpace space
-      dims' = map pe64 dims
-      num_groups' = fmap pe64 num_groups
-      group_size' = fmap pe64 group_size
-      global_tid = Imp.le64 $ segFlat space
-      w = last dims'
-
-  counter <- genZeroes "counters" $ fromIntegral maxNumOps
-
-  reds_group_res_arrs <- groupResultArrays num_groups group_size reds
-
-  num_threads <-
-    dPrimV "num_threads" $
-      unCount num_groups' * unCount group_size'
-
-  sKernelThread "segred_nonseg" (segFlat space) (defKernelAttrs num_groups group_size) $ do
-    constants <- kernelConstants <$> askEnv
-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
-    reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds
-
-    -- Since this is the nonsegmented case, all outer segment IDs must
-    -- necessarily be 0.
-    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64)
-
-    let num_elements = Imp.elements w
-        elems_per_thread =
-          num_elements
-            `divUp` Imp.elements (sExt64 (kernelNumThreads constants))
-
-    slugs <-
-      mapM (segBinOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $
-        zip3 reds reds_arrs reds_group_res_arrs
-    reds_op_renamed <-
-      reductionStageOne
-        constants
-        (zip gtids dims')
-        num_elements
-        global_tid
-        elems_per_thread
-        (tvExp num_threads)
-        slugs
-        body
-
-    let segred_pes =
-          chunks (map (length . segBinOpNeutral) reds) $
-            patElems segred_pat
-    forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $
-      \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do
-        let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op
-        reductionStageTwo
-          constants
-          pes
-          (kernelGroupId constants)
-          0
-          [0]
-          0
-          (sExt64 $ kernelNumGroups constants)
-          slug
-          red_x_params
-          red_y_params
-          red_op_renamed
-          nes
-          1
-          counter
-          (fromInteger i)
-          sync_arr
-          group_res_arrs
-          red_arrs
-
-smallSegmentsReduction ::
-  Pat LetDecMem ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
-  SegSpace ->
-  [SegBinOp GPUMem] ->
-  DoSegBody ->
-  CallKernelGen ()
-smallSegmentsReduction (Pat segred_pes) num_groups group_size space reds body = do
-  let (gtids, dims) = unzip $ unSegSpace space
-      dims' = map pe64 dims
-      segment_size = last dims'
-
-  -- Careful to avoid division by zero now.
-  segment_size_nonzero <-
-    dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size
-
-  let num_groups' = fmap pe64 num_groups
-      group_size' = fmap pe64 group_size
-  num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size'
-  let num_segments = product $ init dims'
-      segments_per_group = unCount group_size' `quot` segment_size_nonzero
-      required_groups = sExt32 $ num_segments `divUp` segments_per_group
-
-  emit $ Imp.DebugPrint "# SegRed-small" Nothing
-  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments
-  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size
-  emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group
-  emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups
-
-  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_groups group_size) $ do
-    constants <- kernelConstants <$> askEnv
-    reds_arrs <- mapM (intermediateArrays group_size (Var $ tvVar num_threads)) reds
-
-    -- We probably do not have enough actual workgroups to cover the
-    -- entire iteration space.  Some groups thus have to perform double
-    -- duty; we put an outer loop to accomplish this.
-    virtualiseGroups SegVirt required_groups $ \group_id' -> do
-      -- Compute the 'n' input indices.  The outer 'n-1' correspond to
-      -- the segment ID, and are computed from the group id.  The inner
-      -- is computed from the local thread id, and may be out-of-bounds.
-      let ltid = sExt64 $ kernelLocalThreadId constants
-          segment_index =
-            (ltid `quot` segment_size_nonzero)
-              + (sExt64 group_id' * sExt64 segments_per_group)
-          index_within_segment = ltid `rem` segment_size
-
-      dIndexSpace (zip (init gtids) (init dims')) segment_index
-      dPrimV_ (last gtids) index_within_segment
-
-      let out_of_bounds =
-            forM_ (zip reds reds_arrs) $ \(SegBinOp _ _ nes _, red_arrs) ->
-              forM_ (zip red_arrs nes) $ \(arr, ne) ->
-                copyDWIMFix arr [ltid] ne []
-
-          in_bounds =
-            body $ \red_res ->
-              sComment "save results to be reduced" $ do
-                let red_dests = map (,[ltid]) (concat reds_arrs)
-                forM_ (zip red_dests red_res) $ \((d, d_is), (res, res_is)) ->
-                  copyDWIMFix d d_is res res_is
-
-      sComment "apply map function if in bounds" $
-        sIf
-          ( segment_size
-              .>. 0
-              .&&. isActive (init $ zip gtids dims)
-              .&&. ltid
-              .<. segment_size
-              * segments_per_group
-          )
-          in_bounds
-          out_of_bounds
-
-      sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.
-      let crossesSegment from to =
-            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)
-      sWhen (segment_size .>. 0) $
-        sComment "perform segmented scan to imitate reduction" $
-          forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) ->
-            groupScan
-              (Just crossesSegment)
-              (sExt64 $ tvExp num_threads)
-              (segment_size * segments_per_group)
-              red_op
-              red_arrs
-
-      sOp $ Imp.Barrier Imp.FenceLocal
-
-      sComment "save final values of segments"
-        $ sWhen
-          ( sExt64 group_id'
-              * segments_per_group
-              + sExt64 ltid
-                .<. num_segments
-                .&&. ltid
-                .<. segments_per_group
-          )
-        $ forM_ (zip segred_pes (concat reds_arrs))
-        $ \(pe, arr) -> do
-          -- Figure out which segment result this thread should write...
-          let flat_segment_index =
-                sExt64 group_id' * segments_per_group + sExt64 ltid
-              gtids' =
-                unflattenIndex (init dims') flat_segment_index
-          copyDWIMFix
-            (patElemName pe)
-            gtids'
-            (Var arr)
-            [(ltid + 1) * segment_size_nonzero - 1]
-
-      -- Finally another barrier, because we will be writing to the
-      -- local memory array first thing in the next iteration.
-      sOp $ Imp.Barrier Imp.FenceLocal
-
-largeSegmentsReduction ::
-  Pat LetDecMem ->
-  Count NumGroups SubExp ->
-  Count GroupSize SubExp ->
-  SegSpace ->
-  [SegBinOp GPUMem] ->
-  DoSegBody ->
-  CallKernelGen ()
-largeSegmentsReduction segred_pat num_groups group_size space reds body = do
-  let (gtids, dims) = unzip $ unSegSpace space
-      dims' = map pe64 dims
-      num_segments = product $ init dims'
-      segment_size = last dims'
-      num_groups' = fmap pe64 num_groups
-      group_size' = fmap pe64 group_size
-
-  (groups_per_segment, elems_per_thread) <-
-    groupsPerSegmentAndElementsPerThread
-      segment_size
-      num_segments
-      num_groups'
-      group_size'
-  virt_num_groups <-
-    dPrimV "virt_num_groups" $
-      groups_per_segment * num_segments
-
-  num_threads <-
-    dPrimV "num_threads" $
-      unCount num_groups' * unCount group_size'
-
-  threads_per_segment <-
-    dPrimV "threads_per_segment" $
-      groups_per_segment * unCount group_size'
-
-  emit $ Imp.DebugPrint "# SegRed-large" Nothing
-  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments
-  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size
-  emit $ Imp.DebugPrint "virt_num_groups" $ Just $ untyped $ tvExp virt_num_groups
-  emit $ Imp.DebugPrint "num_groups" $ Just $ untyped $ Imp.unCount num_groups'
-  emit $ Imp.DebugPrint "group_size" $ Just $ untyped $ Imp.unCount group_size'
-  emit $ Imp.DebugPrint "elems_per_thread" $ Just $ untyped $ Imp.unCount elems_per_thread
-  emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment
-
-  reds_group_res_arrs <- groupResultArrays (Count (tvSize virt_num_groups)) group_size reds
-
-  -- In principle we should have a counter for every segment.  Since
-  -- the number of segments is a dynamic quantity, we would have to
-  -- allocate and zero out an array here, which is expensive.
-  -- However, we exploit the fact that the number of segments being
-  -- reduced at any point in time is limited by the number of
-  -- workgroups. If we bound the number of workgroups, we can get away
-  -- with using that many counters.  FIXME: Is this limit checked
-  -- anywhere?  There are other places in the compiler that will fail
-  -- if the group count exceeds the maximum group size, which is at
-  -- most 1024 anyway.
-  let num_counters = fromIntegral maxNumOps * 1024
-  counter <- genZeroes "counters" num_counters
-
-  sKernelThread "segred_large" (segFlat space) (defKernelAttrs num_groups group_size) $ do
-    constants <- kernelConstants <$> askEnv
-    reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds
-    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
-
-    -- We probably do not have enough actual workgroups to cover the
-    -- entire iteration space.  Some groups thus have to perform double
-    -- duty; we put an outer loop to accomplish this.
-    virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do
-      let segment_gtids = init gtids
-          w = last dims
-          local_tid = kernelLocalThreadId constants
-
-      flat_segment_id <-
-        dPrimVE "flat_segment_id" $
-          group_id `quot` sExt32 groups_per_segment
-
-      global_tid <-
-        dPrimVE "global_tid" $
-          (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid)
-            `rem` (sExt64 (unCount group_size') * groups_per_segment)
-
-      let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment
-      dIndexSpace (zip segment_gtids (init dims')) $ sExt64 flat_segment_id
-      dPrim_ (last gtids) int64
-      let num_elements = Imp.elements $ pe64 w
-
-      slugs <-
-        mapM (segBinOpSlug local_tid group_id) $
-          zip3 reds reds_arrs reds_group_res_arrs
-      reds_op_renamed <-
-        reductionStageOne
-          constants
-          (zip gtids dims')
-          num_elements
-          global_tid
-          elems_per_thread
-          (tvExp threads_per_segment)
-          slugs
-          body
-
-      let segred_pes =
-            chunks (map (length . segBinOpNeutral) reds) $
-              patElems segred_pat
-
-          multiple_groups_per_segment =
-            forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $
-              \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do
-                let (red_x_params, red_y_params) =
-                      splitAt (length nes) $ lambdaParams red_op
-                reductionStageTwo
-                  constants
-                  pes
-                  group_id
-                  flat_segment_id
-                  (map Imp.le64 segment_gtids)
-                  (sExt64 first_group_for_segment)
-                  groups_per_segment
-                  slug
-                  red_x_params
-                  red_y_params
-                  red_op_renamed
-                  nes
-                  (fromIntegral num_counters)
-                  counter
-                  (fromInteger i)
-                  sync_arr
-                  group_res_arrs
-                  red_arrs
-
-          one_group_per_segment =
-            comment "first thread in group saves final result to memory" $
-              forM_ (zip slugs segred_pes) $ \(slug, pes) ->
-                sWhen (local_tid .==. 0) $
-                  forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->
-                    copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is
-
-      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment
-
--- Careful to avoid division by zero here.  We have at least one group
--- per segment.
-groupsPerSegmentAndElementsPerThread ::
-  Imp.TExp Int64 ->
-  Imp.TExp Int64 ->
-  Count NumGroups (Imp.TExp Int64) ->
-  Count GroupSize (Imp.TExp Int64) ->
-  CallKernelGen
-    ( Imp.TExp Int64,
-      Imp.Count Imp.Elements (Imp.TExp Int64)
-    )
-groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do
-  groups_per_segment <-
-    dPrimVE "groups_per_segment" $
-      unCount num_groups_hint `divUp` sMax64 1 num_segments
-  elements_per_thread <-
-    dPrimVE "elements_per_thread" $
-      segment_size `divUp` (unCount group_size * groups_per_segment)
-  pure (groups_per_segment, Imp.elements elements_per_thread)
-
--- | A SegBinOp with auxiliary information.
-data SegBinOpSlug = SegBinOpSlug
-  { slugOp :: SegBinOp GPUMem,
-    -- | The arrays used for computing the intra-group reduction
-    -- (either local or global memory).
-    slugArrs :: [VName],
-    -- | Places to store accumulator in stage 1 reduction.
-    slugAccs :: [(VName, [Imp.TExp Int64])]
-  }
-
-slugBody :: SegBinOpSlug -> Body GPUMem
-slugBody = lambdaBody . segBinOpLambda . slugOp
-
-slugParams :: SegBinOpSlug -> [LParam GPUMem]
-slugParams = lambdaParams . segBinOpLambda . slugOp
-
-slugNeutral :: SegBinOpSlug -> [SubExp]
-slugNeutral = segBinOpNeutral . slugOp
-
-slugShape :: SegBinOpSlug -> Shape
-slugShape = segBinOpShape . slugOp
-
-slugsComm :: [SegBinOpSlug] -> Commutativity
-slugsComm = mconcat . map (segBinOpComm . slugOp)
-
-accParams, nextParams :: SegBinOpSlug -> [LParam GPUMem]
-accParams slug = take (length (slugNeutral slug)) $ slugParams slug
-nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug
-
-segBinOpSlug :: Imp.TExp Int32 -> Imp.TExp Int32 -> (SegBinOp GPUMem, [VName], [VName]) -> InKernelGen SegBinOpSlug
-segBinOpSlug local_tid group_id (op, group_res_arrs, param_arrs) =
-  SegBinOpSlug op group_res_arrs
-    <$> zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs
-  where
-    mkAcc p param_arr
-      | Prim t <- paramType p,
-        shapeRank (segBinOpShape op) == 0 = do
-          acc <- dPrim (baseString (paramName p) <> "_acc") t
-          pure (tvVar acc, [])
-      | otherwise =
-          pure (param_arr, [sExt64 local_tid, sExt64 group_id])
-
-computeThreadChunkSize ::
-  Commutativity ->
-  Imp.TExp Int64 ->
-  Imp.TExp Int64 ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  TV Int64 ->
-  ImpM rep r op ()
-computeThreadChunkSize Commutative threads_per_segment thread_index elements_per_thread num_elements chunk_var =
-  chunk_var
-    <-- sMin64
-      (Imp.unCount elements_per_thread)
-      ((Imp.unCount num_elements - thread_index) `divUp` threads_per_segment)
-computeThreadChunkSize Noncommutative _ thread_index elements_per_thread num_elements chunk_var = do
-  starting_point <-
-    dPrimV "starting_point" $
-      thread_index * Imp.unCount elements_per_thread
-  remaining_elements <-
-    dPrimV "remaining_elements" $
-      Imp.unCount num_elements - tvExp starting_point
-
-  let no_remaining_elements = tvExp remaining_elements .<=. 0
-      beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point
-
-  sIf
-    (no_remaining_elements .||. beyond_bounds)
-    (chunk_var <-- 0)
-    ( sIf
-        is_last_thread
-        (chunk_var <-- Imp.unCount last_thread_elements)
-        (chunk_var <-- Imp.unCount elements_per_thread)
-    )
-  where
-    last_thread_elements =
-      num_elements - Imp.elements thread_index * elements_per_thread
-    is_last_thread =
-      Imp.unCount num_elements
-        .<. (thread_index + 1)
-        * Imp.unCount elements_per_thread
-
-reductionStageZero ::
-  KernelConstants ->
-  [(VName, Imp.TExp Int64)] ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  Imp.TExp Int64 ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  Imp.TExp Int64 ->
-  [SegBinOpSlug] ->
-  DoSegBody ->
-  InKernelGen ([Lambda GPUMem], InKernelGen ())
-reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do
-  let (gtids, _dims) = unzip ispace
-      gtid = mkTV (last gtids) int64
-      local_tid = sExt64 $ kernelLocalThreadId constants
-
-  -- Figure out how many elements this thread should process.
-  chunk_size <- dPrim "chunk_size" int64
-  computeThreadChunkSize
-    (slugsComm slugs)
-    threads_per_segment
-    (sExt64 global_tid)
-    elems_per_thread
-    num_elements
-    chunk_size
-
-  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs
-
-  sComment "neutral-initialise the accumulators" $
-    forM_ slugs $ \slug ->
-      forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->
-        sLoopNest (slugShape slug) $ \vec_is ->
-          copyDWIMFix acc (acc_is ++ vec_is) ne []
-
-  slugs_op_renamed <- mapM (renameLambda . segBinOpLambda . slugOp) slugs
-
-  let doTheReduction =
-        forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) ->
-          sLoopNest (slugShape slug) $ \vec_is -> do
-            comment "to reduce current chunk, first store our result in memory" $ do
-              forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->
-                copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)
-
-              forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) ->
-                when (primType $ paramType p) $
-                  copyDWIMFix arr [local_tid] (Var $ paramName p) []
-
-            sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.
-            groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug)
-
-            sOp $ Imp.Barrier Imp.FenceLocal
-
-            sComment "first thread saves the result in accumulator" $
-              sWhen (local_tid .==. 0) $
-                forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) ->
-                  copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) []
-
-  -- If this is a non-commutative reduction, each thread must run the
-  -- loop the same number of iterations, because we will be performing
-  -- a group-wide reduction in there.
-  let comm = slugsComm slugs
-      (bound, check_bounds) =
-        case comm of
-          Commutative -> (tvExp chunk_size, id)
-          Noncommutative ->
-            ( Imp.unCount elems_per_thread,
-              sWhen (tvExp gtid .<. Imp.unCount num_elements)
-            )
-
-  sFor "i" bound $ \i -> do
-    gtid
-      <-- case comm of
-        Commutative ->
-          global_tid + threads_per_segment * i
-        Noncommutative ->
-          let index_in_segment = global_tid `quot` kernelGroupSize constants
-           in sExt64 local_tid
-                + (index_in_segment * Imp.unCount elems_per_thread + i)
-                  * kernelGroupSize constants
-
-    check_bounds $
-      sComment "apply map function" $
-        body $ \all_red_res -> do
-          let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res
-
-          forM_ (zip slugs slugs_res) $ \(slug, red_res) ->
-            sLoopNest (slugShape slug) $ \vec_is -> do
-              sComment "load accumulator" $
-                forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->
-                  copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)
-              sComment "load new values" $
-                forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) ->
-                  copyDWIMFix (paramName p) [] res (res_is ++ vec_is)
-              sComment "apply reduction operator"
-                $ compileStms mempty (bodyStms $ slugBody slug)
-                $ sComment "store in accumulator"
-                $ forM_
-                  ( zip
-                      (slugAccs slug)
-                      (map resSubExp $ bodyResult $ slugBody slug)
-                  )
-                $ \((acc, acc_is), se) ->
-                  copyDWIMFix acc (acc_is ++ vec_is) se []
-
-    case comm of
-      Noncommutative -> do
-        doTheReduction
-        sComment "first thread keeps accumulator; others reset to neutral element" $ do
-          let reset_to_neutral =
-                forM_ slugs $ \slug ->
-                  forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->
-                    sLoopNest (slugShape slug) $ \vec_is ->
-                      copyDWIMFix acc (acc_is ++ vec_is) ne []
-          sUnless (local_tid .==. 0) reset_to_neutral
-      _ -> pure ()
-  sOp $ Imp.ErrorSync Imp.FenceLocal
-  pure (slugs_op_renamed, doTheReduction)
-
-reductionStageOne ::
-  KernelConstants ->
-  [(VName, Imp.TExp Int64)] ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  Imp.TExp Int64 ->
-  Imp.Count Imp.Elements (Imp.TExp Int64) ->
-  Imp.TExp Int64 ->
-  [SegBinOpSlug] ->
-  DoSegBody ->
-  InKernelGen [Lambda GPUMem]
-reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do
-  (slugs_op_renamed, doTheReduction) <-
-    reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body
-
-  case slugsComm slugs of
-    Noncommutative -> pure ()
-    Commutative -> doTheReduction
-
-  pure slugs_op_renamed
-
-reductionStageTwo ::
-  KernelConstants ->
-  [PatElem LetDecMem] ->
-  Imp.TExp Int32 ->
-  Imp.TExp Int32 ->
-  [Imp.TExp Int64] ->
-  Imp.TExp Int64 ->
-  Imp.TExp Int64 ->
-  SegBinOpSlug ->
-  [LParam GPUMem] ->
-  [LParam GPUMem] ->
-  Lambda GPUMem ->
-  [SubExp] ->
-  Imp.TExp Int32 ->
-  VName ->
-  Imp.TExp Int32 ->
-  VName ->
-  [VName] ->
-  [VName] ->
-  InKernelGen ()
-reductionStageTwo
-  constants
-  segred_pes
-  group_id
-  flat_segment_id
-  segment_gtids
-  first_group_for_segment
-  groups_per_segment
-  slug
-  red_x_params
-  red_y_params
-  red_op_renamed
-  nes
-  num_counters
-  counter
-  counter_i
-  sync_arr
-  group_res_arrs
-  red_arrs = do
-    let local_tid = kernelLocalThreadId constants
-        group_size = kernelGroupSize constants
-    old_counter <- dPrim "old_counter" int32
-    (counter_mem, _, counter_offset) <-
-      fullyIndexArray
-        counter
-        [ sExt64 $
-            counter_i * num_counters
-              + flat_segment_id `rem` num_counters
-        ]
-    comment "first thread in group saves group result to global memory" $
-      sWhen (local_tid .==. 0) $ do
-        forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->
-          copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is
-        sOp $ Imp.MemFence Imp.FenceGlobal
-        -- Increment the counter, thus stating that our result is
-        -- available.
-        sOp
-          $ Imp.Atomic DefaultSpace
-          $ Imp.AtomicAdd
-            Int32
-            (tvVar old_counter)
-            counter_mem
-            counter_offset
-          $ untyped (1 :: Imp.TExp Int32)
-        -- Now check if we were the last group to write our result.  If
-        -- so, it is our responsibility to produce the final result.
-        sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1
-
-    sOp $ Imp.Barrier Imp.FenceGlobal
-
-    is_last_group <- dPrim "is_last_group" Bool
-    copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0]
-    sWhen (tvExp is_last_group) $ do
-      -- The final group has written its result (and it was
-      -- us!), so read in all the group results and perform the
-      -- final stage of the reduction.  But first, we reset the
-      -- counter so it is ready for next time.  This is done
-      -- with an atomic to avoid warnings about write/write
-      -- races in oclgrind.
-      sWhen (local_tid .==. 0) $
-        sOp $
-          Imp.Atomic DefaultSpace $
-            Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $
-              untyped $
-                negate groups_per_segment
-
-      sLoopNest (slugShape slug) $ \vec_is -> do
-        unless (null $ slugShape slug) $
-          sOp (Imp.Barrier Imp.FenceLocal)
-
-        -- There is no guarantee that the number of workgroups for the
-        -- segment is less than the workgroup size, so each thread may
-        -- have to read multiple elements.  We do this in a sequential
-        -- way that may induce non-coalesced accesses, but the total
-        -- number of accesses should be tiny here.
-        comment "read in the per-group-results" $ do
-          read_per_thread <-
-            dPrimVE "read_per_thread" $
-              groups_per_segment `divUp` sExt64 group_size
-
-          forM_ (zip red_x_params nes) $ \(p, ne) ->
-            copyDWIMFix (paramName p) [] ne []
-
-          sFor "i" read_per_thread $ \i -> do
-            group_res_id <-
-              dPrimVE "group_res_id" $
-                sExt64 local_tid * read_per_thread + i
-            index_of_group_res <-
-              dPrimVE "index_of_group_res" $
-                first_group_for_segment + group_res_id
-
-            sWhen (group_res_id .<. groups_per_segment) $ do
-              forM_ (zip red_y_params group_res_arrs) $
-                \(p, group_res_arr) ->
-                  copyDWIMFix
-                    (paramName p)
-                    []
-                    (Var group_res_arr)
-                    ([0, index_of_group_res] ++ vec_is)
-
-              compileStms mempty (bodyStms $ slugBody slug) $
-                forM_ (zip red_x_params $ map resSubExp $ bodyResult $ slugBody slug) $ \(p, se) ->
-                  copyDWIMFix (paramName p) [] se []
-
-        forM_ (zip red_x_params red_arrs) $ \(p, arr) ->
-          when (primType $ paramType p) $
-            copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) []
-
-        sOp $ Imp.ErrorSync Imp.FenceLocal
-
-        sComment "reduce the per-group results" $ do
-          groupReduce (sExt32 group_size) red_op_renamed red_arrs
-
-          sComment "and back to memory with the final result" $
-            sWhen (local_tid .==. 0) $
-              forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) ->
-                copyDWIMFix
-                  (patElemName pe)
-                  (segment_gtids ++ vec_is)
-                  (Var $ paramName p)
-                  []
+--
+-- An optimization specfically targeted at non-segmented and large-segments
+-- segmented reductions with non-commutative is made: The stage one main loop is
+-- essentially stripmined by a factor `chunk`, inserting collective copies via
+-- local memory of each reduction parameter going into the intra-group (partial)
+-- reductions. This saves a factor `chunk` number of intra-group reductions at
+-- the cost of some overhead in collective copies.
+module Futhark.CodeGen.ImpGen.GPU.SegRed
+  ( compileSegRed,
+    compileSegRed',
+    DoSegBody,
+  )
+where
+
+import Control.Monad
+import Data.List (genericLength, zip4)
+import Data.Map qualified as M
+import Data.Maybe
+import Futhark.CodeGen.ImpCode.GPU qualified as Imp
+import Futhark.CodeGen.ImpGen
+import Futhark.CodeGen.ImpGen.GPU.Base
+import Futhark.Error
+import Futhark.IR.GPUMem
+import Futhark.IR.Mem.LMAD qualified as LMAD
+import Futhark.Transform.Rename
+import Futhark.Util (chunks, mapAccumLM)
+import Futhark.Util.IntegralExp (divUp, nextMul, quot, rem)
+import Prelude hiding (quot, rem)
+
+forM2_ :: (Monad m) => [a] -> [b] -> (a -> b -> m c) -> m ()
+forM2_ xs ys f = forM_ (zip xs ys) (uncurry f)
+
+-- | The maximum number of operators we support in a single SegRed.
+-- This limit arises out of the static allocation of counters.
+maxNumOps :: Int
+maxNumOps = 10
+
+-- | Code generation for the body of the SegRed, taking a continuation
+-- for saving the results of the body.  The results should be
+-- represented as a pairing of a t'SubExp' along with a list of
+-- indexes into that t'SubExp' for reading the result.
+type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen ()
+
+-- | Datatype used to distinguish between and work with the different sets of
+-- intermediate memory we need for the different ReduceKinds.
+data SegRedIntermediateArrays
+  = GeneralSegRedInterms
+      { groupRedArrs :: [VName]
+      }
+  | NoncommPrimSegRedInterms
+      { collCopyArrs :: [VName],
+        groupRedArrs :: [VName],
+        privateChunks :: [VName]
+      }
+
+-- | Compile 'SegRed' instance to host-level code with calls to
+-- various kernels.
+compileSegRed ::
+  Pat LetDecMem ->
+  SegLevel ->
+  SegSpace ->
+  [SegBinOp GPUMem] ->
+  KernelBody GPUMem ->
+  CallKernelGen ()
+compileSegRed pat lvl space segbinops map_kbody = do
+  emit $ Imp.DebugPrint "\n# SegRed" Nothing
+  KernelAttrs {kAttrNumGroups = num_groups, kAttrGroupSize = group_size} <-
+    lvlKernelAttrs lvl
+  let grid = KernelGrid num_groups group_size
+
+  compileSegRed' pat grid space segbinops $ \red_cont ->
+    sComment "apply map function" $
+      compileStms mempty (kernelBodyStms map_kbody) $ do
+        let (red_res, map_res) = splitAt (segBinOpResults segbinops) $ kernelBodyResult map_kbody
+
+        let mapout_arrs = drop (segBinOpResults segbinops) $ patElems pat
+        unless (null mapout_arrs) $
+          sComment "write map-out result(s)" $ do
+            zipWithM_ (compileThreadResult space) mapout_arrs map_res
+
+        red_cont $ map ((,[]) . kernelResultSubExp) red_res
+  emit $ Imp.DebugPrint "" Nothing
+
+paramOf :: SegBinOp GPUMem -> [Param LParamMem]
+paramOf (SegBinOp _ op ne _) = take (length ne) $ lambdaParams op
+
+isPrimSegBinOp :: SegBinOp GPUMem -> Bool
+isPrimSegBinOp segbinop =
+  all primType (lambdaReturnType $ segBinOpLambda segbinop)
+    && shapeRank (segBinOpShape segbinop) == 0
+
+-- | Like 'compileSegRed', but where the body is a monadic action.
+compileSegRed' ::
+  Pat LetDecMem ->
+  KernelGrid ->
+  SegSpace ->
+  [SegBinOp GPUMem] ->
+  DoSegBody ->
+  CallKernelGen ()
+compileSegRed' pat grid space segbinops map_body_cont
+  | genericLength segbinops > maxNumOps =
+      compilerLimitationS $
+        "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."
+  | otherwise = do
+      chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const
+      case unSegSpace space of
+        [(_, Constant (IntValue (Int64Value 1))), _] ->
+          compileReduction (chunk_v, chunk_const) nonsegmentedReduction
+        _ -> do
+          let segment_size = pe64 $ last $ segSpaceDims space
+              use_small_segments = segment_size * 2 .<. pe64 (unCount group_size) * tvExp chunk_v
+          sIf
+            use_small_segments
+            (compileReduction (chunk_v, chunk_const) smallSegmentsReduction)
+            (compileReduction (chunk_v, chunk_const) largeSegmentsReduction)
+  where
+    compileReduction chunk f =
+      f pat num_groups group_size chunk space segbinops map_body_cont
+
+    param_types = map paramType $ concatMap paramOf segbinops
+
+    num_groups = gridNumGroups grid
+    group_size = gridGroupSize grid
+
+    chunk_const =
+      if Noncommutative `elem` map segBinOpComm segbinops
+        && all isPrimSegBinOp segbinops
+        then getChunkSize param_types
+        else Imp.ValueExp $ IntValue $ intValue Int64 (1 :: Int64)
+
+-- | Prepare intermediate arrays for the reduction.  Prim-typed
+-- arguments go in local memory (so we need to do the allocation of
+-- those arrays inside the kernel), while array-typed arguments go in
+-- global memory.  Allocations for the latter have already been
+-- performed.  This policy is baked into how the allocations are done
+-- in ExplicitAllocations.
+--
+-- For more info about the intermediate arrays used for the different reduction
+-- kernels, see note [IntermArrays].
+makeIntermArrays ::
+  Imp.TExp Int64 ->
+  SubExp ->
+  SubExp ->
+  [SegBinOp GPUMem] ->
+  InKernelGen [SegRedIntermediateArrays]
+makeIntermArrays group_id group_size chunk segbinops
+  | Noncommutative <- mconcat (map segBinOpComm segbinops),
+    all isPrimSegBinOp segbinops =
+      noncommPrimSegRedInterms
+  | otherwise =
+      generalSegRedInterms group_id group_size segbinops
+  where
+    params = map paramOf segbinops
+
+    noncommPrimSegRedInterms = do
+      group_worksize <- tvSize <$> dPrimV "group_worksize" group_worksize_E
+
+      -- compute total amount of lmem.
+      let sum_ x y = nextMul x y + group_size_E * y
+          group_reds_lmem_requirement = foldl sum_ 0 $ concat elem_sizes
+          collcopy_lmem_requirement = group_worksize_E * max_elem_size
+          lmem_total_size =
+            Imp.bytes $
+              collcopy_lmem_requirement `sMax64` group_reds_lmem_requirement
+
+      -- offsets into the total pool of lmem for each group reduction array.
+      (_, offsets) <-
+        forAccumLM2D 0 elem_sizes $ \byte_offs elem_size ->
+          (,byte_offs `quot` elem_size)
+            <$> dPrimVE "offset" (sum_ byte_offs elem_size)
+
+      -- total pool of local mem.
+      lmem <- sAlloc "local_mem" lmem_total_size (Space "local")
+      let arrInLMem ptype name len_se offset =
+            sArray
+              (name ++ "_" ++ prettyString ptype)
+              ptype
+              (Shape [len_se])
+              lmem
+              $ LMAD.iota offset [pe64 len_se]
+
+      forM (zipWith zip params offsets) $ \ps_and_offsets -> do
+        (coll_copy_arrs, group_red_arrs, priv_chunks) <-
+          fmap unzip3 $ forM ps_and_offsets $ \(p, offset) -> do
+            let ptype = elemType $ paramType p
+            (,,)
+              <$> arrInLMem ptype "coll_copy_arr" group_worksize 0
+              <*> arrInLMem ptype "group_red_arr" group_size offset
+              <*> sAllocArray
+                ("chunk_" ++ prettyString ptype)
+                ptype
+                (Shape [chunk])
+                (ScalarSpace [chunk] ptype)
+        pure $ NoncommPrimSegRedInterms coll_copy_arrs group_red_arrs priv_chunks
+    group_size_E = pe64 group_size
+    group_worksize_E = group_size_E * pe64 chunk
+
+    paramSize = primByteSize . elemType . paramType
+    elem_sizes = map (map paramSize) params
+    max_elem_size = maximum $ concat elem_sizes
+
+    forAccumLM2D acc ls f = mapAccumLM (mapAccumLM f) acc ls
+
+generalSegRedInterms ::
+  Imp.TExp Int64 ->
+  SubExp ->
+  [SegBinOp GPUMem] ->
+  InKernelGen [SegRedIntermediateArrays]
+generalSegRedInterms group_id group_size segbinops =
+  fmap (map GeneralSegRedInterms) $
+    forM (map paramOf segbinops) $
+      mapM $ \p ->
+        case paramDec p of
+          MemArray pt shape _ (ArrayIn mem _) -> do
+            let shape' = Shape [group_size] <> shape
+            let shape_E = map pe64 $ shapeDims shape'
+            sArray ("red_arr_" ++ prettyString pt) pt shape' mem $
+              LMAD.iota (group_id * product shape_E) shape_E
+          _ -> do
+            let pt = elemType $ paramType p
+                shape = Shape [group_size]
+            sAllocArray ("red_arr_" ++ prettyString pt) pt shape $ Space "local"
+
+-- | Arrays for storing group results.
+--
+-- The group-result arrays have an extra dimension because they are
+-- also used for keeping vectorised accumulators for first-stage
+-- reduction, if necessary.  If necessary, this dimension has size
+-- group_size, and otherwise 1.  When actually storing group results,
+-- the first index is set to 0.
+groupResultArrays ::
+  SubExp ->
+  SubExp ->
+  [SegBinOp GPUMem] ->
+  CallKernelGen [[VName]]
+groupResultArrays num_virtgroups group_size segbinops =
+  forM segbinops $ \(SegBinOp _ lam _ shape) ->
+    forM (lambdaReturnType lam) $ \t -> do
+      let pt = elemType t
+          extra_dim
+            | primType t, shapeRank shape == 0 = intConst Int64 1
+            | otherwise = group_size
+          full_shape = Shape [extra_dim, num_virtgroups] <> shape <> arrayShape t
+          -- Move the groupsize dimension last to ensure coalesced
+          -- memory access.
+          perm = [1 .. shapeRank full_shape - 1] ++ [0]
+      sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm
+
+type DoCompileSegRed =
+  Pat LetDecMem ->
+  Count NumGroups SubExp ->
+  Count GroupSize SubExp ->
+  (TV Int64, Imp.KernelConstExp) ->
+  SegSpace ->
+  [SegBinOp GPUMem] ->
+  DoSegBody ->
+  CallKernelGen ()
+
+nonsegmentedReduction :: DoCompileSegRed
+nonsegmentedReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do
+  let (gtids, dims) = unzip $ unSegSpace space
+      chunk = tvExp chunk_v
+      num_groups_se = unCount num_groups
+      group_size_se = unCount group_size
+      group_size' = pe64 group_size_se
+      global_tid = Imp.le64 $ segFlat space
+      n = pe64 $ last dims
+
+  counters <- genZeroes "counters" maxNumOps
+
+  reds_group_res_arrs <- groupResultArrays num_groups_se group_size_se segbinops
+
+  num_threads <-
+    fmap tvSize $ dPrimV "num_threads" $ pe64 num_groups_se * group_size'
+
+  let attrs =
+        (defKernelAttrs num_groups group_size)
+          { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const
+          }
+
+  sKernelThread "segred_nonseg" (segFlat space) attrs $ do
+    constants <- kernelConstants <$> askEnv
+    let ltid = kernelLocalThreadId constants
+    let group_id = kernelGroupId constants
+
+    interms <- makeIntermArrays (sExt64 group_id) group_size_se (tvSize chunk_v) segbinops
+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
+
+    -- Since this is the nonsegmented case, all outer segment IDs must
+    -- necessarily be 0.
+    forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64)
+
+    q <- dPrimVE "q" $ n `divUp` (sExt64 (kernelNumThreads constants) * chunk)
+
+    slugs <-
+      mapM (segBinOpSlug ltid group_id) $
+        zip3 segbinops interms reds_group_res_arrs
+    new_lambdas <-
+      reductionStageOne
+        gtids
+        n
+        global_tid
+        q
+        chunk
+        (pe64 num_threads)
+        slugs
+        map_body_cont
+
+    let segred_pess =
+          chunks
+            (map (length . segBinOpNeutral) segbinops)
+            segred_pes
+
+    forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $
+      \(pes, slug, new_lambda, i) ->
+        reductionStageTwo
+          pes
+          group_id
+          [0]
+          0
+          (sExt64 $ kernelNumGroups constants)
+          slug
+          new_lambda
+          counters
+          sync_arr
+          (fromInteger i)
+
+smallSegmentsReduction :: DoCompileSegRed
+smallSegmentsReduction (Pat segred_pes) num_groups group_size _ space segbinops map_body_cont = do
+  let (gtids, dims) = unzip $ unSegSpace space
+      dims' = map pe64 dims
+      segment_size = last dims'
+
+  -- Careful to avoid division by zero now.
+  segment_size_nonzero <-
+    dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size
+
+  let group_size_se = unCount group_size
+      num_groups_se = unCount group_size
+      num_groups' = pe64 num_groups_se
+      group_size' = pe64 group_size_se
+  num_threads <- fmap tvSize $ dPrimV "num_threads" $ num_groups' * group_size'
+  let num_segments = product $ init dims'
+      segments_per_group = group_size' `quot` segment_size_nonzero
+      required_groups = sExt32 $ num_segments `divUp` segments_per_group
+
+  emit $ Imp.DebugPrint "# SegRed-small" Nothing
+  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments
+  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size
+  emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group
+  emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups
+
+  sKernelThread "segred_small" (segFlat space) (defKernelAttrs num_groups group_size) $ do
+    constants <- kernelConstants <$> askEnv
+    let group_id = kernelGroupSize constants
+        ltid = sExt64 $ kernelLocalThreadId constants
+
+    interms <- generalSegRedInterms group_id group_size_se segbinops
+    let reds_arrs = map groupRedArrs interms
+
+    -- We probably do not have enough actual workgroups to cover the
+    -- entire iteration space.  Some groups thus have to perform double
+    -- duty; we put an outer loop to accomplish this.
+    virtualiseGroups SegVirt required_groups $ \virtgroup_id -> do
+      -- Compute the 'n' input indices.  The outer 'n-1' correspond to
+      -- the segment ID, and are computed from the group id.  The inner
+      -- is computed from the local thread id, and may be out-of-bounds.
+      let segment_index =
+            (ltid `quot` segment_size_nonzero)
+              + (sExt64 virtgroup_id * sExt64 segments_per_group)
+          index_within_segment = ltid `rem` segment_size
+
+      dIndexSpace (zip (init gtids) (init dims')) segment_index
+      dPrimV_ (last gtids) index_within_segment
+
+      let in_bounds =
+            map_body_cont $ \red_res ->
+              sComment "save results to be reduced" $ do
+                let red_dests = map (,[ltid]) (concat reds_arrs)
+                forM2_ red_dests red_res $ \(d, d_is) (res, res_is) ->
+                  copyDWIMFix d d_is res res_is
+          out_of_bounds =
+            forM2_ segbinops reds_arrs $ \(SegBinOp _ _ nes _) red_arrs ->
+              forM2_ red_arrs nes $ \arr ne ->
+                copyDWIMFix arr [ltid] ne []
+
+      sComment "apply map function if in bounds" $
+        sIf
+          ( segment_size
+              .>. 0
+              .&&. isActive (init $ zip gtids dims)
+              .&&. ltid
+              .<. segment_size
+              * segments_per_group
+          )
+          in_bounds
+          out_of_bounds
+
+      sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.
+      let crossesSegment from to =
+            (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size)
+      sWhen (segment_size .>. 0) $
+        sComment "perform segmented scan to imitate reduction" $
+          forM2_ segbinops reds_arrs $ \(SegBinOp _ red_op _ _) red_arrs ->
+            groupScan
+              (Just crossesSegment)
+              (sExt64 $ pe64 num_threads)
+              (segment_size * segments_per_group)
+              red_op
+              red_arrs
+
+      sOp $ Imp.Barrier Imp.FenceLocal
+
+      sComment "save final values of segments"
+        $ sWhen
+          ( sExt64 virtgroup_id
+              * segments_per_group
+              + sExt64 ltid
+                .<. num_segments
+                .&&. ltid
+                .<. segments_per_group
+          )
+        $ forM2_ segred_pes (concat reds_arrs)
+        $ \pe arr -> do
+          -- Figure out which segment result this thread should write...
+          let flat_segment_index =
+                sExt64 virtgroup_id * segments_per_group + sExt64 ltid
+              gtids' =
+                unflattenIndex (init dims') flat_segment_index
+          copyDWIMFix
+            (patElemName pe)
+            gtids'
+            (Var arr)
+            [(ltid + 1) * segment_size_nonzero - 1]
+
+      -- Finally another barrier, because we will be writing to the
+      -- local memory array first thing in the next iteration.
+      sOp $ Imp.Barrier Imp.FenceLocal
+
+largeSegmentsReduction :: DoCompileSegRed
+largeSegmentsReduction (Pat segred_pes) num_groups group_size (chunk_v, chunk_const) space segbinops map_body_cont = do
+  let (gtids, dims) = unzip $ unSegSpace space
+      dims' = map pe64 dims
+      num_segments = product $ init dims'
+      segment_size = last dims'
+      num_groups' = pe64 $ unCount num_groups
+      group_size_se = unCount group_size
+      group_size' = pe64 group_size_se
+      chunk = tvExp chunk_v
+
+  groups_per_segment <-
+    dPrimVE "groups_per_segment" $
+      num_groups' `divUp` sMax64 1 num_segments
+
+  q <-
+    dPrimVE "q" $
+      segment_size `divUp` (group_size' * groups_per_segment * chunk)
+
+  num_virtgroups <-
+    dPrimV "num_virtgroups" $
+      groups_per_segment * num_segments
+  threads_per_segment <-
+    dPrimVE "threads_per_segment" $
+      groups_per_segment * group_size'
+
+  emit $ Imp.DebugPrint "# SegRed-large" Nothing
+  emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments
+  emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size
+  emit $ Imp.DebugPrint "num_virtgroups" $ Just $ untyped $ tvExp num_virtgroups
+  emit $ Imp.DebugPrint "num_groups" $ Just $ untyped num_groups'
+  emit $ Imp.DebugPrint "group_size" $ Just $ untyped group_size'
+  emit $ Imp.DebugPrint "q" $ Just $ untyped q
+  emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment
+
+  reds_group_res_arrs <- groupResultArrays (tvSize num_virtgroups) group_size_se segbinops
+
+  -- In principle we should have a counter for every segment.  Since
+  -- the number of segments is a dynamic quantity, we would have to
+  -- allocate and zero out an array here, which is expensive.
+  -- However, we exploit the fact that the number of segments being
+  -- reduced at any point in time is limited by the number of
+  -- workgroups. If we bound the number of workgroups, we can get away
+  -- with using that many counters.  FIXME: Is this limit checked
+  -- anywhere?  There are other places in the compiler that will fail
+  -- if the group count exceeds the maximum group size, which is at
+  -- most 1024 anyway.
+  let num_counters = maxNumOps * 1024
+  counters <- genZeroes "counters" $ fromIntegral num_counters
+
+  let attrs =
+        (defKernelAttrs num_groups group_size)
+          { kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const
+          }
+
+  sKernelThread "segred_large" (segFlat space) attrs $ do
+    constants <- kernelConstants <$> askEnv
+    let group_id = sExt64 $ kernelGroupId constants
+        ltid = kernelLocalThreadId constants
+
+    interms <- makeIntermArrays group_id group_size_se (tvSize chunk_v) segbinops
+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
+
+    -- We probably do not have enough actual workgroups to cover the
+    -- entire iteration space.  Some groups thus have to perform double
+    -- duty; we put an outer loop to accomplish this.
+    virtualiseGroups SegVirt (sExt32 (tvExp num_virtgroups)) $ \virtgroup_id -> do
+      let segment_gtids = init gtids
+
+      flat_segment_id <-
+        dPrimVE "flat_segment_id" $
+          sExt64 virtgroup_id `quot` groups_per_segment
+
+      global_tid <-
+        dPrimVE "global_tid" $
+          (sExt64 virtgroup_id * sExt64 group_size' + sExt64 ltid)
+            `rem` threads_per_segment
+
+      let first_group_for_segment = flat_segment_id * groups_per_segment
+      dIndexSpace (zip segment_gtids (init dims')) flat_segment_id
+      dPrim_ (last gtids) int64
+      let n = pe64 $ last dims
+
+      slugs <-
+        mapM (segBinOpSlug ltid virtgroup_id) $
+          zip3 segbinops interms reds_group_res_arrs
+      new_lambdas <-
+        reductionStageOne
+          gtids
+          n
+          global_tid
+          q
+          chunk
+          threads_per_segment
+          slugs
+          map_body_cont
+
+      let segred_pess =
+            chunks
+              (map (length . segBinOpNeutral) segbinops)
+              segred_pes
+
+          multiple_groups_per_segment =
+            forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $
+              \(pes, slug, new_lambda, i) -> do
+                let counter_idx =
+                      fromIntegral (i * num_counters)
+                        + flat_segment_id
+                          `rem` fromIntegral num_counters
+                reductionStageTwo
+                  pes
+                  virtgroup_id
+                  (map Imp.le64 segment_gtids)
+                  first_group_for_segment
+                  groups_per_segment
+                  slug
+                  new_lambda
+                  counters
+                  sync_arr
+                  counter_idx
+
+          one_group_per_segment =
+            sComment "first thread in group saves final result to memory" $
+              forM2_ slugs segred_pess $ \slug pes ->
+                sWhen (ltid .==. 0) $
+                  forM2_ pes (slugAccs slug) $ \v (acc, acc_is) ->
+                    copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is
+
+      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment
+
+-- | Auxiliary information for a single reduction. A slug holds the `SegBinOp`
+-- operator for a single reduction, the different arrays required throughout
+-- stages one and two, and a global mem destination for the final result of the
+-- particular reduction.
+data SegBinOpSlug = SegBinOpSlug
+  { slugOp :: SegBinOp GPUMem,
+    -- | Intermediate arrays needed for the given reduction.
+    slugInterms :: SegRedIntermediateArrays,
+    -- | Place(s) to store group accumulator(s) in stage 1 reduction.
+    slugAccs :: [(VName, [Imp.TExp Int64])],
+    -- | Global memory destination(s) for the final result(s) for this
+    -- particular reduction.
+    groupResArrs :: [VName]
+  }
+
+segBinOpSlug ::
+  Imp.TExp Int32 ->
+  Imp.TExp Int32 ->
+  (SegBinOp GPUMem, SegRedIntermediateArrays, [VName]) ->
+  InKernelGen SegBinOpSlug
+segBinOpSlug ltid group_id (op, interms, group_res_arrs) = do
+  accs <- zipWithM mkAcc (lambdaParams (segBinOpLambda op)) group_res_arrs
+  pure $ SegBinOpSlug op interms accs group_res_arrs
+  where
+    mkAcc p group_res_arr
+      | Prim t <- paramType p,
+        shapeRank (segBinOpShape op) == 0 = do
+          group_res_acc <- dPrim (baseString (paramName p) <> "_group_res_acc") t
+          pure (tvVar group_res_acc, [])
+      -- if this is a non-primitive reduction, the global mem result array will
+      -- double as accumulator.
+      | otherwise =
+          pure (group_res_arr, [sExt64 ltid, sExt64 group_id])
+
+slugLambda :: SegBinOpSlug -> Lambda GPUMem
+slugLambda = segBinOpLambda . slugOp
+
+slugBody :: SegBinOpSlug -> Body GPUMem
+slugBody = lambdaBody . slugLambda
+
+slugParams :: SegBinOpSlug -> [LParam GPUMem]
+slugParams = lambdaParams . slugLambda
+
+slugNeutral :: SegBinOpSlug -> [SubExp]
+slugNeutral = segBinOpNeutral . slugOp
+
+slugShape :: SegBinOpSlug -> Shape
+slugShape = segBinOpShape . slugOp
+
+slugsComm :: [SegBinOpSlug] -> Commutativity
+slugsComm = mconcat . map (segBinOpComm . slugOp)
+
+slugSplitParams :: SegBinOpSlug -> ([LParam GPUMem], [LParam GPUMem])
+slugSplitParams slug = splitAt (length (slugNeutral slug)) $ slugParams slug
+
+slugGroupRedArrs :: SegBinOpSlug -> [VName]
+slugGroupRedArrs = groupRedArrs . slugInterms
+
+slugPrivChunks :: SegBinOpSlug -> [VName]
+slugPrivChunks = privateChunks . slugInterms
+
+slugCollCopyArrs :: SegBinOpSlug -> [VName]
+slugCollCopyArrs = collCopyArrs . slugInterms
+
+reductionStageOne ::
+  [VName] ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  [SegBinOpSlug] ->
+  DoSegBody ->
+  InKernelGen [Lambda GPUMem]
+reductionStageOne gtids n global_tid q chunk threads_per_segment slugs body_cont = do
+  constants <- kernelConstants <$> askEnv
+  let glb_ind_var = mkTV (last gtids) int64
+      ltid = sExt64 $ kernelLocalThreadId constants
+
+  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs
+
+  sComment "ne-initialise the outer (per-group) accumulator(s)" $ do
+    forM_ slugs $ \slug ->
+      forM2_ (slugAccs slug) (slugNeutral slug) $ \(acc, acc_is) ne ->
+        sLoopNest (slugShape slug) $ \vec_is ->
+          copyDWIMFix acc (acc_is ++ vec_is) ne []
+
+  new_lambdas <- mapM (renameLambda . slugLambda) slugs
+  let group_size = sExt32 $ kernelGroupSize constants
+  let doGroupReduce =
+        forM2_ slugs new_lambdas $ \slug new_lambda -> do
+          let accs = slugAccs slug
+          let params = slugParams slug
+          sLoopNest (slugShape slug) $ \vec_is -> do
+            let group_red_arrs = slugGroupRedArrs slug
+            sComment "store accs. prims go in lmem; non-prims in params (in global mem)" $
+              forM_ (zip3 group_red_arrs accs params) $
+                \(arr, (acc, acc_is), p) ->
+                  if isPrimParam p
+                    then copyDWIMFix arr [ltid] (Var acc) (acc_is ++ vec_is)
+                    else copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)
+
+            sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.
+            groupReduce group_size new_lambda group_red_arrs
+            sOp $ Imp.Barrier Imp.FenceLocal
+
+            sComment "thread 0 updates per-group acc(s); rest reset to ne" $ do
+              sIf
+                (ltid .==. 0)
+                ( forM2_ accs (lambdaParams new_lambda) $
+                    \(acc, acc_is) p ->
+                      copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) []
+                )
+                ( forM2_ accs (slugNeutral slug) $
+                    \(acc, acc_is) ne ->
+                      copyDWIMFix acc (acc_is ++ vec_is) ne []
+                )
+
+  case (slugsComm slugs, all (isPrimSegBinOp . slugOp) slugs) of
+    (Noncommutative, True) ->
+      noncommPrimParamsStageOneBody
+        slugs
+        body_cont
+        glb_ind_var
+        global_tid
+        q
+        n
+        chunk
+        doGroupReduce
+    _ ->
+      generalStageOneBody
+        slugs
+        body_cont
+        glb_ind_var
+        global_tid
+        q
+        n
+        threads_per_segment
+        doGroupReduce
+
+  pure new_lambdas
+
+generalStageOneBody ::
+  [SegBinOpSlug] ->
+  DoSegBody ->
+  TV Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  InKernelGen () ->
+  InKernelGen ()
+generalStageOneBody slugs body_cont glb_ind_var global_tid q n threads_per_segment doGroupReduce = do
+  let is_comm = slugsComm slugs == Commutative
+
+  constants <- kernelConstants <$> askEnv
+  let group_size = kernelGroupSize constants
+      ltid = sExt64 $ kernelLocalThreadId constants
+
+  -- this group's id within its designated segment, and this group's initial
+  -- global offset.
+  group_id_in_segment <- dPrimVE "group_id_in_segment" $ global_tid `quot` group_size
+  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_size
+
+  sFor "i" q $ \i -> do
+    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_size
+    glb_ind_var
+      <-- if is_comm
+        then global_tid + threads_per_segment * i
+        else group_offset + ltid
+
+    sWhen (tvExp glb_ind_var .<. n) $
+      sComment "apply map function(s)" $
+        body_cont $ \all_red_res -> do
+          let maps_res = chunks (map (length . slugNeutral) slugs) all_red_res
+
+          forM2_ slugs maps_res $ \slug map_res ->
+            sLoopNest (slugShape slug) $ \vec_is -> do
+              let (acc_params, next_params) = slugSplitParams slug
+              sComment "load accumulator(s)" $
+                forM2_ acc_params (slugAccs slug) $ \p (acc, acc_is) ->
+                  copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is)
+              sComment "load next value(s)" $
+                forM2_ next_params map_res $ \p (res, res_is) ->
+                  copyDWIMFix (paramName p) [] res (res_is ++ vec_is)
+              sComment "apply reduction operator(s)"
+                $ compileStms mempty (bodyStms $ slugBody slug)
+                $ sComment "store in accumulator(s)"
+                $ forM2_
+                  (slugAccs slug)
+                  (map resSubExp $ bodyResult $ slugBody slug)
+                $ \(acc, acc_is) se ->
+                  copyDWIMFix acc (acc_is ++ vec_is) se []
+
+    unless is_comm doGroupReduce
+  sOp $ Imp.ErrorSync Imp.FenceLocal
+  when is_comm doGroupReduce
+
+noncommPrimParamsStageOneBody ::
+  [SegBinOpSlug] ->
+  DoSegBody ->
+  TV Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  InKernelGen () ->
+  InKernelGen ()
+noncommPrimParamsStageOneBody slugs body_cont glb_ind_var global_tid q n chunk doLMemGroupReduce = do
+  constants <- kernelConstants <$> askEnv
+  let group_size = kernelGroupSize constants
+      ltid = sExt64 $ kernelLocalThreadId constants
+
+  -- this group's id within its designated segment; the stride made per group in
+  -- the outer `i < q` loop; and this group's initial global offset.
+  group_id_in_segment <- dPrimVE "group_offset_in_segment" $ global_tid `quot` group_size
+  group_stride <- dPrimVE "group_stride" $ group_size * chunk
+  group_base_offset <- dPrimVE "group_base_offset" $ group_id_in_segment * q * group_stride
+
+  let chunkLoop = sFor "k" chunk
+
+  sFor "i" q $ \i -> do
+    -- group offset in this iteration.
+    group_offset <- dPrimVE "group_offset" $ group_base_offset + i * group_stride
+    chunkLoop $ \k -> do
+      loc_ind <- dPrimVE "loc_ind" $ k * group_size + ltid
+      glb_ind_var <-- group_offset + loc_ind
+
+      sIf
+        (tvExp glb_ind_var .<. n)
+        ( body_cont $ \all_red_res -> do
+            let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res
+            forM2_ slugs slugs_res $ \slug slug_res -> do
+              let priv_chunks = slugPrivChunks slug
+              sComment "write map result(s) to private chunk(s)" $
+                forM2_ priv_chunks slug_res $ \priv_chunk (res, res_is) ->
+                  copyDWIMFix priv_chunk [k] res res_is
+        )
+        -- if out of bounds, fill chunk(s) with neutral element(s)
+        ( forM_ slugs $ \slug ->
+            forM2_ (slugPrivChunks slug) (slugNeutral slug) $
+              \priv_chunk ne ->
+                copyDWIMFix priv_chunk [k] ne []
+        )
+
+    sOp $ Imp.ErrorSync Imp.FenceLocal
+
+    sComment "effectualize collective copies in local memory" $ do
+      forM_ slugs $ \slug -> do
+        let coll_copy_arrs = slugCollCopyArrs slug
+        let priv_chunks = slugPrivChunks slug
+
+        forM2_ coll_copy_arrs priv_chunks $ \lmem_arr priv_chunk -> do
+          chunkLoop $ \k -> do
+            lmem_idx <- dPrimVE "lmem_idx" $ ltid + k * group_size
+            copyDWIMFix lmem_arr [lmem_idx] (Var priv_chunk) [k]
+
+          sOp $ Imp.Barrier Imp.FenceLocal
+
+          chunkLoop $ \k -> do
+            lmem_idx <- dPrimVE "lmem_idx" $ ltid * chunk + k
+            copyDWIMFix priv_chunk [k] (Var lmem_arr) [lmem_idx]
+
+          sOp $ Imp.Barrier Imp.FenceLocal
+
+    sComment "per-thread sequential reduction of private chunk(s)" $ do
+      chunkLoop $ \k ->
+        forM_ slugs $ \slug -> do
+          let accs = map fst $ slugAccs slug
+          let (acc_ps, next_ps) = slugSplitParams slug
+          let ps_accs_chunks = zip4 acc_ps next_ps accs (slugPrivChunks slug)
+
+          sComment "load params for all reductions" $ do
+            forM_ ps_accs_chunks $ \(acc_p, next_p, acc, priv_chunk) -> do
+              copyDWIM (paramName acc_p) [] (Var acc) []
+              copyDWIMFix (paramName next_p) [] (Var priv_chunk) [k]
+
+          sComment "apply reduction operator(s)" $ do
+            let binop_ress = map resSubExp $ bodyResult $ slugBody slug
+            compileStms mempty (bodyStms $ slugBody slug) $
+              forM2_ accs binop_ress $ \acc binop_res ->
+                copyDWIM acc [] binop_res []
+    doLMemGroupReduce
+  sOp $ Imp.ErrorSync Imp.FenceLocal
+
+reductionStageTwo ::
+  [PatElem LetDecMem] ->
+  Imp.TExp Int32 ->
+  [Imp.TExp Int64] ->
+  Imp.TExp Int64 ->
+  Imp.TExp Int64 ->
+  SegBinOpSlug ->
+  Lambda GPUMem ->
+  VName ->
+  VName ->
+  Imp.TExp Int64 ->
+  InKernelGen ()
+reductionStageTwo segred_pes group_id segment_gtids first_group_for_segment groups_per_segment slug new_lambda counters sync_arr counter_idx = do
+  constants <- kernelConstants <$> askEnv
+
+  let ltid32 = kernelLocalThreadId constants
+      ltid = sExt64 ltid32
+      group_size = kernelGroupSize constants
+
+  let (acc_params, next_params) = slugSplitParams slug
+  let nes = slugNeutral slug
+  let red_arrs = slugGroupRedArrs slug
+  let group_res_arrs = groupResArrs slug
+
+  old_counter <- dPrim "old_counter" int32
+  (counter_mem, _, counter_offset) <-
+    fullyIndexArray
+      counters
+      [counter_idx]
+  sComment "first thread in group saves group result to global memory" $
+    sWhen (ltid32 .==. 0) $ do
+      forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->
+        copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is
+      sOp $ Imp.MemFence Imp.FenceGlobal
+      -- Increment the counter, thus stating that our result is
+      -- available.
+      sOp
+        $ Imp.Atomic DefaultSpace
+        $ Imp.AtomicAdd
+          Int32
+          (tvVar old_counter)
+          counter_mem
+          counter_offset
+        $ untyped (1 :: Imp.TExp Int32)
+      -- Now check if we were the last group to write our result.  If
+      -- so, it is our responsibility to produce the final result.
+      sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1
+
+  sOp $ Imp.Barrier Imp.FenceGlobal
+
+  is_last_group <- dPrim "is_last_group" Bool
+  copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0]
+  sWhen (tvExp is_last_group) $ do
+    -- The final group has written its result (and it was
+    -- us!), so read in all the group results and perform the
+    -- final stage of the reduction.  But first, we reset the
+    -- counter so it is ready for next time.  This is done
+    -- with an atomic to avoid warnings about write/write
+    -- races in oclgrind.
+    sWhen (ltid32 .==. 0) $
+      sOp $
+        Imp.Atomic DefaultSpace $
+          Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $
+            untyped $
+              negate groups_per_segment
+
+    sLoopNest (slugShape slug) $ \vec_is -> do
+      unless (null $ slugShape slug) $
+        sOp (Imp.Barrier Imp.FenceLocal)
+
+      -- There is no guarantee that the number of workgroups for the
+      -- segment is less than the workgroup size, so each thread may
+      -- have to read multiple elements.  We do this in a sequential
+      -- way that may induce non-coalesced accesses, but the total
+      -- number of accesses should be tiny here.
+      --
+      -- TODO: here we *could* insert a collective copy of the num_groups
+      -- per-group results. However, it may not be beneficial, since num_groups
+      -- is not necessarily larger than group_size, meaning the number of
+      -- uncoalesced reads here may be insignificant. In fact, if we happen to
+      -- have a num_groups < group_size, then the collective copy would add
+      -- unnecessary overhead. Also, this code is only executed by a single
+      -- group.
+      sComment "read in the per-group-results" $ do
+        read_per_thread <-
+          dPrimVE "read_per_thread" $
+            groups_per_segment `divUp` sExt64 group_size
+
+        forM2_ acc_params nes $ \p ne ->
+          copyDWIM (paramName p) [] ne []
+
+        sFor "i" read_per_thread $ \i -> do
+          group_res_id <-
+            dPrimVE "group_res_id" $
+              ltid * read_per_thread + i
+          index_of_group_res <-
+            dPrimVE "index_of_group_res" $
+              first_group_for_segment + group_res_id
+
+          sWhen (group_res_id .<. groups_per_segment) $ do
+            forM2_ next_params group_res_arrs $
+              \p group_res_arr ->
+                copyDWIMFix
+                  (paramName p)
+                  []
+                  (Var group_res_arr)
+                  ([0, index_of_group_res] ++ vec_is)
+
+            compileStms mempty (bodyStms $ slugBody slug) $
+              forM2_ acc_params (map resSubExp $ bodyResult $ slugBody slug) $ \p se ->
+                copyDWIMFix (paramName p) [] se []
+
+      forM2_ acc_params red_arrs $ \p arr ->
+        when (isPrimParam p) $
+          copyDWIMFix arr [ltid] (Var $ paramName p) []
+
+      sOp $ Imp.ErrorSync Imp.FenceLocal
+
+      sComment "reduce the per-group results" $ do
+        groupReduce (sExt32 group_size) new_lambda red_arrs
+
+        sComment "and back to memory with the final result" $
+          sWhen (ltid32 .==. 0) $
+            forM2_ segred_pes (lambdaParams new_lambda) $ \pe p ->
+              copyDWIMFix
+                (patElemName pe)
+                (segment_gtids ++ vec_is)
+                (Var $ paramName p)
+                []
+
+-- Note [IntermArrays]
+--
+-- Intermediate memory for the nonsegmented and large segments non-commutative
+-- reductions with all primitive parameters:
+--
+--   These kernels need local memory for 1) the initial collective copy, 2) the
+--   (virtualized) group reductions, and (TODO: this one not implemented yet!)
+--   3) the final single-group collective copy. There are no dependencies
+--   between these three stages, so we can reuse the same pool of local mem for
+--   all three. These intermediates all go into local mem because of the
+--   assumption of primitive parameter types.
+--
+--   Let `elem_sizes` be a list of element type sizes for the reduction
+--   operators in a given redomap fusion. Then the amount of local mem needed
+--   across the three steps are:
+--
+--   1) The initial collective copy from global to thread-private memory
+--   requires `group_size * CHUNK * max elem_sizes`, since the collective copies
+--   are performed in sequence (ie. inputs to different reduction operators need
+--   not be held in local mem simultaneously).
+--   2) The intra-group reductions of local memory held per-thread results
+--   require `group_size * sum elem_sizes` bytes, since per-thread results for
+--   all fused reductions are group-reduced simultaneously.
+--   3) If group_size < num_groups, then after the final single-group collective
+--   copy, a thread-sequential reduction reduces the number of per-group partial
+--   results from num_groups down to group_size for each reduction array, such
+--   that they will each fit in the final intra-group reduction. This requires
+--   `num_groups * max elem_sizes`.
+--
+--   In summary, the total amount of local mem needed is the maximum between:
+--   1) initial collective copy: group_size * CHUNK * max elem_sizes
+--   2) intra-group reductions:  group_size * sum elem_sizes
+--   3) final collective copy:   num_groups * max elem_sizes
+--
+--   The amount of local mem will most likely be decided by 1) in most cases,
+--   unless the number of fused operators is very high *or* if we have a
+--   `num_groups > group_size * CHUNK`, but this is unlikely, in which case 2)
+--   and 3), respectively, will dominate.
+--
+--   Aside from local memory, these kernels also require a CHUNK-sized array of
+--   thread-private register memory per reduction operator.
+--
+-- For all other reductions, ie. commutative reductions, reductions with at
+-- least one non-primitive operator, and small segments reductions:
+--
+--   These kernels use local memory only for the intra-group reductions, and
+--   since they do not use chunking or CHUNK, they all require onlly `group_size
+--   * max elem_sizes` bytes of local memory and no thread-private register mem.
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
@@ -7,6 +7,7 @@
 
 import Control.Monad
 import Data.List (zip4, zip7)
+import Data.Map qualified as M
 import Data.Maybe
 import Futhark.CodeGen.ImpCode.GPU qualified as Imp
 import Futhark.CodeGen.ImpGen
@@ -15,7 +16,7 @@
 import Futhark.IR.Mem.LMAD qualified as LMAD
 import Futhark.Transform.Rename
 import Futhark.Util (mapAccumLM, takeLast)
-import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, quot)
+import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, nextMul, quot)
 import Prelude hiding (mod, quot, rem)
 
 xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem]
@@ -24,9 +25,6 @@
 yParams scan =
   drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan))
 
-alignTo :: (IntegralExp a) => a -> a -> a
-alignTo x a = (x `divUp` a) * a
-
 createLocalArrays ::
   Count GroupSize SubExp ->
   SubExp ->
@@ -36,7 +34,7 @@
   let groupSizeE = pe64 groupSize
       workSize = pe64 chunk * groupSizeE
       prefixArraysSize =
-        foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $
+        foldl (\acc tySize -> nextMul acc tySize + tySize * groupSizeE) 0 $
           map primByteSize types
       maxTransposedArraySize =
         foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types
@@ -44,7 +42,7 @@
       warpSize :: (Num a) => a
       warpSize = 32
       maxWarpExchangeSize =
-        foldl (\acc tySize -> alignTo acc tySize + tySize * fromInteger warpSize) 0 $
+        foldl (\acc tySize -> nextMul acc tySize + tySize * fromInteger warpSize) 0 $
           map primByteSize types
       maxLookbackSize = maxWarpExchangeSize + warpSize
       size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize
@@ -52,7 +50,7 @@
   (_, byteOffsets) <-
     mapAccumLM
       ( \off tySize -> do
-          off' <- dPrimVE "byte_offsets" $ alignTo off tySize + pe64 groupSize * tySize
+          off' <- dPrimVE "byte_offsets" $ nextMul off tySize + pe64 groupSize * tySize
           pure (off', off)
       )
       0
@@ -61,7 +59,7 @@
   (_, warpByteOffsets) <-
     mapAccumLM
       ( \off tySize -> do
-          off' <- dPrimVE "warp_byte_offset" $ alignTo off tySize + warpSize * tySize
+          off' <- dPrimVE "warp_byte_offset" $ nextMul off tySize + warpSize * tySize
           pure (off', off)
       )
       warpSize
@@ -208,12 +206,10 @@
       | otherwise =
           copyDWIMFix (paramName p) [] (Var arr) [gtid - behind + arrs_full_size]
 
-    writeResult x y arr
-      | primType $ paramType x = do
-          copyDWIMFix arr [ltid] (Var $ paramName x) []
-          copyDWIM (paramName y) [] (Var $ paramName x) []
-      | otherwise =
-          copyDWIM (paramName y) [] (Var $ paramName x) []
+    writeResult x y arr = do
+      when (isPrimParam x) $
+        copyDWIMFix arr [ltid] (Var $ paramName x) []
+      copyDWIM (paramName y) [] (Var $ paramName x) []
 
 -- | Compile 'SegScan' instance to host-level code with calls to a
 -- single-pass kernel.
@@ -232,27 +228,17 @@
 
       n = product $ map pe64 $ segSpaceDims space
 
-      tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scan_op
-      tys_sizes = map primByteSize tys
-
-      sumT, maxT :: Integer
-      sumT = sum tys_sizes
-      sumT' = sum (map (max 4 . primByteSize) tys) `div` 4
-      maxT = maximum tys_sizes
-
-      -- TODO: Make these constants dynamic by querying device
-      k_reg = 64
-      k_mem = 95
-
-      mem_constraint = max k_mem sumT `div` maxT
-      reg_constraint = (k_reg - 1 - sumT') `div` (2 * sumT')
+      tys' = lambdaReturnType $ segBinOpLambda scan_op
 
-      chunk :: (Num a) => a
-      chunk = fromIntegral $ max 1 $ min mem_constraint reg_constraint
+      tys = map elemType tys'
 
       group_size_e = pe64 $ unCount $ kAttrGroupSize attrs
       num_physgroups_e = pe64 $ unCount $ kAttrNumGroups attrs
 
+  let chunk_const = getChunkSize tys'
+  chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const
+  let chunk = tvExp chunk_v
+
   num_virtgroups <-
     tvSize <$> dPrimV "num_virtgroups" (n `divUp` (group_size_e * chunk))
   let num_virtgroups_e = pe64 num_virtgroups
@@ -266,11 +252,7 @@
       not_segmented_e = fromBool $ not segmented
       segment_size = last dims'
 
-  let debug_ s v = emit $ Imp.DebugPrint s $ Just $ untyped (v :: Imp.TExp Int32)
-  debug_ "Sequential elements per thread (chunk) " chunk
-  debug_ "Memory constraint" $ fromIntegral mem_constraint
-  debug_ "Register constraint" $ fromIntegral reg_constraint
-  debug_ "sumT'" $ fromIntegral sumT'
+  emit $ Imp.DebugPrint "Sequential elements per thread (chunk)" $ Just $ untyped chunk
 
   statusFlags <- sAllocArray "status_flags" int8 (Shape [num_virtgroups]) (Space "device")
   sReplicate statusFlags $ intConst Int8 statusX
@@ -284,11 +266,18 @@
 
   global_id <- genZeroes "global_dynid" 1
 
-  sKernelThread "segscan" (segFlat space) attrs $ do
+  let attrs' = attrs {kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const}
+
+  sKernelThread "segscan" (segFlat space) attrs' $ do
+    chunk32 <- dPrimVE "chunk_size_32b" $ sExt32 $ tvExp chunk_v
+
     constants <- kernelConstants <$> askEnv
 
+    let ltid32 = kernelLocalThreadId constants
+        ltid = sExt64 ltid32
+
     (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <-
-      createLocalArrays (kAttrGroupSize attrs) (intConst Int64 chunk) tys
+      createLocalArrays (kAttrGroupSize attrs) (tvSize chunk_v) tys
 
     -- We wrap the entire kernel body in a virtualisation loop to handle the
     -- case where we do not have enough workgroups to cover the iteration space.
@@ -308,7 +297,7 @@
     sFor "virtloop_i" iters $ const $ do
       dyn_id <- dPrim "dynamic_id" int32
       sComment "First thread in block fetches this block's dynamic_id" $ do
-        sWhen (kernelLocalThreadId constants .==. 0) $ do
+        sWhen (ltid32 .==. 0) $ do
           (globalIdMem, _, globalIdOff) <- fullyIndexArray global_id [0]
           sOp $
             Imp.Atomic DefaultSpace $
@@ -333,34 +322,32 @@
       copyDWIMFix (tvVar dyn_id) [] (Var sharedId) [0]
       sOp local_barrier
 
-      blockOff <-
-        dPrimV "blockOff" $
-          sExt64 (tvExp dyn_id) * chunk * group_size_e -- kernelGroupSize constants
-      sgmIdx <- dPrimVE "sgm_idx" $ tvExp blockOff `mod` segment_size
+      block_offset <-
+        dPrimVE "block_offset" $
+          sExt64 (tvExp dyn_id) * chunk * group_size_e
+      sgm_idx <- dPrimVE "sgm_idx" $ block_offset `mod` segment_size
       boundary <-
         dPrimVE "boundary" $
           sExt32 $
-            sMin64 (chunk * group_size_e) (segment_size - sgmIdx)
+            sMin64 (chunk * group_size_e) (segment_size - sgm_idx)
       segsize_compact <-
         dPrimVE "segsize_compact" $
           sExt32 $
             sMin64 (chunk * group_size_e) segment_size
-      privateArrays <-
+      private_chunks <-
         forM tys $ \ty ->
           sAllocArray
             "private"
             ty
-            (Shape [intConst Int64 chunk])
-            (ScalarSpace [intConst Int64 chunk] ty)
+            (Shape [tvSize chunk_v])
+            (ScalarSpace [tvSize chunk_v] ty)
 
+      thd_offset <- dPrimVE "thd_offset" $ block_offset + ltid
+
       sComment "Load and map" $
         sFor "i" chunk $ \i -> do
           -- The map's input index
-          virt_tid <-
-            dPrimVE "virt_tid" $
-              tvExp blockOff
-                + sExt64 (kernelLocalThreadId constants)
-                + i * kernelGroupSize constants
+          virt_tid <- dPrimVE "virt_tid" $ thd_offset + i * group_size_e
           dIndexSpace (zip gtids dims') virt_tid
           -- Perform the map
           let in_bounds =
@@ -373,66 +360,62 @@
                     copyDWIMFix (patElemName dest) (map Imp.le64 gtids) (kernelResultSubExp src) []
 
                   -- Write to-scan results to private memory.
-                  forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->
+                  forM_ (zip private_chunks $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->
                     copyDWIMFix dest [i] src []
 
               out_of_bounds =
-                forM_ (zip privateArrays scanop_nes) $ \(dest, ne) ->
+                forM_ (zip private_chunks scanop_nes) $ \(dest, ne) ->
                   copyDWIMFix dest [i] ne []
 
           sIf (virt_tid .<. n) in_bounds out_of_bounds
 
       sOp $ Imp.ErrorSync Imp.FenceLocal
       sComment "Transpose scan inputs" $ do
-        forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do
+        forM_ (zip transposedArrays private_chunks) $ \(trans, priv) -> do
           sFor "i" chunk $ \i -> do
-            sharedIdx <-
-              dPrimVE "sharedIdx" $
-                sExt64 (kernelLocalThreadId constants)
-                  + i * kernelGroupSize constants
+            sharedIdx <- dPrimVE "sharedIdx" $ ltid + i * group_size_e
             copyDWIMFix trans [sharedIdx] (Var priv) [i]
           sOp local_barrier
           sFor "i" chunk $ \i -> do
-            sharedIdx <- dPrimV "sharedIdx" $ kernelLocalThreadId constants * chunk + i
+            sharedIdx <- dPrimV "sharedIdx" $ ltid * chunk + i
             copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx]
           sOp local_barrier
 
       sComment "Per thread scan" $ do
         -- We don't need to touch the first element, so only m-1
         -- iterations here.
-        globalIdx <-
-          dPrimVE "gidx" $
-            (kernelLocalThreadId constants * chunk) + 1
         sFor "i" (chunk - 1) $ \i -> do
           let xs = map paramName $ xParams scan_op
               ys = map paramName $ yParams scan_op
           -- determine if start of segment
           new_sgm <-
             if segmented
-              then dPrimVE "new_sgm" $ (globalIdx + sExt32 i - boundary) `mod` segsize_compact .==. 0
+              then do
+                gidx <- dPrimVE "gidx" $ (ltid32 * chunk32) + 1
+                dPrimVE "new_sgm" $ (gidx + sExt32 i - boundary) `mod` segsize_compact .==. 0
               else pure false
           -- skip scan of first element in segment
           sUnless new_sgm $ do
-            forM_ (zip4 privateArrays xs ys tys) $ \(src, x, y, ty) -> do
+            forM_ (zip4 private_chunks xs ys tys) $ \(src, x, y, ty) -> do
               dPrim_ x ty
               dPrim_ y ty
               copyDWIMFix x [] (Var src) [i]
               copyDWIMFix y [] (Var src) [i + 1]
 
             compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $
-              forM_ (zip privateArrays $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) ->
+              forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) ->
                 copyDWIMFix dest [i + 1] res []
 
       sComment "Publish results in shared memory" $ do
-        forM_ (zip prefixArrays privateArrays) $ \(dest, src) ->
-          copyDWIMFix dest [sExt64 $ kernelLocalThreadId constants] (Var src) [chunk - 1]
+        forM_ (zip prefixArrays private_chunks) $ \(dest, src) ->
+          copyDWIMFix dest [ltid] (Var src) [chunk - 1]
         sOp local_barrier
 
       let crossesSegment = do
             guard segmented
             Just $ \from to ->
-              let from' = (from + 1) * chunk - 1
-                  to' = (to + 1) * chunk - 1
+              let from' = (from + 1) * chunk32 - 1
+                  to' = (to + 1) * chunk32 - 1
                in (to' - from') .>. (to' + segsize_compact - boundary) `mod` segsize_compact
 
       scan_op1 <- renameLambda $ segBinOpLambda scan_op
@@ -441,19 +424,19 @@
       sComment "Scan results (with warp scan)" $ do
         groupScan
           crossesSegment
+          group_size_e
           num_virt_threads
-          (kernelGroupSize constants)
           scan_op1
           prefixArrays
 
         sOp $ Imp.ErrorSync Imp.FenceLocal
 
         let firstThread acc prefixes =
-              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelGroupSize constants) - 1]
+              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 group_size_e - 1]
             notFirstThread acc prefixes =
-              copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelLocalThreadId constants) - 1]
+              copyDWIMFix (tvVar acc) [] (Var prefixes) [ltid - 1]
         sIf
-          (kernelLocalThreadId constants .==. 0)
+          (ltid32 .==. 0)
           (zipWithM_ firstThread accs prefixArrays)
           (zipWithM_ notFirstThread accs prefixArrays)
 
@@ -461,9 +444,9 @@
 
       prefixes <- forM (zip scanop_nes tys) $ \(ne, ty) ->
         dPrimV "prefix" $ TPrimExp $ toExp' ty ne
-      blockNewSgm <- dPrimVE "block_new_sgm" $ sgmIdx .==. 0
+      blockNewSgm <- dPrimVE "block_new_sgm" $ sgm_idx .==. 0
       sComment "Perform lookback" $ do
-        sWhen (blockNewSgm .&&. kernelLocalThreadId constants .==. 0) $ do
+        sWhen (blockNewSgm .&&. ltid32 .==. 0) $ do
           everythingVolatile $
             forM_ (zip accs incprefixArrays) $ \(acc, incprefixArray) ->
               copyDWIMFix incprefixArray [tvExp dyn_id] (tvSize acc) []
@@ -475,8 +458,8 @@
         -- end sWhen
 
         let warpSize = kernelWaveSize constants
-        sWhen (bNot blockNewSgm .&&. kernelLocalThreadId constants .<. warpSize) $ do
-          sWhen (kernelLocalThreadId constants .==. 0) $ do
+        sWhen (bNot blockNewSgm .&&. ltid32 .<. warpSize) $ do
+          sWhen (ltid32 .==. 0) $ do
             sIf
               (not_segmented_e .||. boundary .==. sExt32 (group_size_e * chunk))
               ( do
@@ -505,7 +488,7 @@
 
           sIf
             (tvExp status .==. statusP)
-            ( sWhen (kernelLocalThreadId constants .==. 0) $
+            ( sWhen (ltid32 .==. 0) $
                 everythingVolatile $
                   forM_ (zip prefixes incprefixArrays) $ \(prefix, incprefixArray) ->
                     copyDWIMFix (tvVar prefix) [] (Var incprefixArray) [tvExp dyn_id - 1]
@@ -518,11 +501,11 @@
                 let loopStop = warpSize * (-1)
                     sameSegment readIdx
                       | segmented =
-                          let startIdx = sExt64 (tvExp readIdx + 1) * kernelGroupSize constants * chunk - 1
-                           in tvExp blockOff - startIdx .<=. sgmIdx
+                          let startIdx = sExt64 (tvExp readIdx + 1) * group_size_e * chunk - 1
+                           in block_offset - startIdx .<=. sgm_idx
                       | otherwise = true
                 sWhile (tvExp readOffset .>. loopStop) $ do
-                  readI <- dPrimV "read_i" $ tvExp readOffset + kernelLocalThreadId constants
+                  readI <- dPrimV "read_i" $ tvExp readOffset + ltid32
                   aggrs <- forM (zip scanop_nes tys) $ \(ne, ty) ->
                     dPrimV "aggr" $ TPrimExp $ toExp' ty ne
                   flag <- dPrimV "flag" (statusX :: Imp.TExp Int8)
@@ -546,12 +529,12 @@
                   -- end sWhen
 
                   forM_ (zip exchanges aggrs) $ \(exchange, aggr) ->
-                    copyDWIMFix exchange [sExt64 $ kernelLocalThreadId constants] (tvSize aggr) []
-                  copyDWIMFix warpscan [sExt64 $ kernelLocalThreadId constants] (tvSize flag) []
+                    copyDWIMFix exchange [ltid] (tvSize aggr) []
+                  copyDWIMFix warpscan [ltid] (tvSize flag) []
 
                   -- execute warp-parallel reduction but only if the last read flag in not STATUS_P
                   copyDWIMFix (tvVar flag) [] (Var warpscan) [sExt64 warpSize - 1]
-                  sWhen (tvExp flag .<. (2 :: Imp.TExp Int8)) $ do
+                  sWhen (tvExp flag .<. statusP) $ do
                     lam' <- renameLambda scan_op1
                     inBlockScanLookback
                       constants
@@ -573,7 +556,7 @@
                     )
 
                   -- update prefix if flag different than STATUS_X:
-                  sWhen (tvExp flag .>. (statusX :: Imp.TExp Int8)) $ do
+                  sWhen (tvExp flag .>. statusX) $ do
                     lam <- renameLambda scan_op1
                     let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams lam
                     forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr)
@@ -586,7 +569,7 @@
 
           -- end sWhile
           -- end sIf
-          sWhen (kernelLocalThreadId constants .==. 0) $ do
+          sWhen (ltid32 .==. 0) $ do
             scan_op2 <- renameLambda scan_op1
             let xs = map paramName $ take (length tys) $ lambdaParams scan_op2
                 ys = map paramName $ drop (length tys) $ lambdaParams scan_op2
@@ -629,7 +612,7 @@
             dPrimV_ y' $ tvExp acc
 
         sIf
-          (kernelLocalThreadId constants * chunk .<. boundary .&&. bNot blockNewSgm)
+          (ltid32 * chunk32 .<. boundary .&&. bNot blockNewSgm)
           ( compileStms mempty (bodyStms $ lambdaBody scan_op4) $
               forM_ (zip3 xs tys $ map resSubExp $ bodyResult $ lambdaBody scan_op4) $
                 \(x, ty, res) -> x <~~ toExp' ty res
@@ -639,38 +622,34 @@
         -- elements left before new segment.
         stop <-
           dPrimVE "stopping_point" $
-            segsize_compact - (kernelLocalThreadId constants * chunk - 1 + segsize_compact - boundary) `rem` segsize_compact
+            segsize_compact - (ltid32 * chunk32 - 1 + segsize_compact - boundary) `rem` segsize_compact
         sFor "i" chunk $ \i -> do
           sWhen (sExt32 i .<. stop - 1) $ do
-            forM_ (zip privateArrays ys) $ \(src, y) ->
+            forM_ (zip private_chunks ys) $ \(src, y) ->
               -- only include prefix for the first segment part per thread
               copyDWIMFix y [] (Var src) [i]
             compileStms mempty (bodyStms $ lambdaBody scan_op3) $
-              forM_ (zip privateArrays $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $
+              forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $
                 \(dest, res) ->
                   copyDWIMFix dest [i] res []
 
       sComment "Transpose scan output and Write it to global memory in coalesced fashion" $ do
-        forM_ (zip3 transposedArrays privateArrays $ map patElemName all_pes) $ \(locmem, priv, dest) -> do
+        forM_ (zip3 transposedArrays private_chunks $ map patElemName all_pes) $ \(locmem, priv, dest) -> do
           -- sOp local_barrier
           sFor "i" chunk $ \i -> do
             sharedIdx <-
               dPrimV "sharedIdx" $
-                sExt64 (kernelLocalThreadId constants * chunk) + i
+                sExt64 (ltid * chunk) + i
             copyDWIMFix locmem [tvExp sharedIdx] (Var priv) [i]
           sOp local_barrier
           sFor "i" chunk $ \i -> do
-            flat_idx <-
-              dPrimVE "flat_idx" $
-                tvExp blockOff
-                  + kernelGroupSize constants * i
-                  + sExt64 (kernelLocalThreadId constants)
+            flat_idx <- dPrimVE "flat_idx" $ thd_offset + i * group_size_e
             dIndexSpace (zip gtids dims') flat_idx
             sWhen (flat_idx .<. n) $ do
               copyDWIMFix
                 dest
                 (map Imp.le64 gtids)
                 (Var locmem)
-                [sExt64 $ flat_idx - tvExp blockOff]
+                [sExt64 $ flat_idx - block_offset]
           sOp local_barrier
 {-# NOINLINE compileSegScan #-}
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs
@@ -10,7 +10,6 @@
 where
 
 import Control.Monad
-import Control.Monad.Identity
 import Control.Monad.Reader
 import Control.Monad.State
 import Data.Bifunctor (second)
@@ -58,7 +57,7 @@
 translateGPU target prog =
   let env = envFromProg prog
       ( prog',
-        ToOpenCL kernels device_funs used_types sizes failures
+        ToOpenCL kernels device_funs used_types sizes failures constants
         ) =
           (`runState` initialOpenCL) . (`runReaderT` env) $ do
             let ImpGPU.Definitions
@@ -86,13 +85,15 @@
             T.unlines device_defs
           ]
    in ImpOpenCL.Program
-        opencl_code
-        opencl_prelude
-        kernels'
-        (S.toList used_types)
-        (findParamUsers env prog' (cleanSizes sizes))
-        failures
-        prog'
+        { openClProgram = opencl_code,
+          openClPrelude = opencl_prelude,
+          openClMacroDefs = constants,
+          openClKernelNames = kernels',
+          openClUsedTypes = S.toList used_types,
+          openClParams = findParamUsers env prog' (cleanSizes sizes),
+          openClFailures = failures,
+          hostDefinitions = prog'
+        }
   where
     genPrelude TargetOpenCL = genOpenClPrelude
     genPrelude TargetCUDA = const genCUDAPrelude
@@ -172,11 +173,12 @@
     clDevFuns :: M.Map Name (C.Definition, T.Text),
     clUsedTypes :: S.Set PrimType,
     clSizes :: M.Map Name SizeClass,
-    clFailures :: [FailureMsg]
+    clFailures :: [FailureMsg],
+    clConstants :: [(Name, KernelConstExp)]
   }
 
 initialOpenCL :: ToOpenCL
-initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty
+initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty mempty
 
 data Env = Env
   { envFuns :: ImpGPU.Functions ImpGPU.HostOp,
@@ -331,7 +333,7 @@
 isConst :: GroupDim -> Maybe T.Text
 isConst (Left (ValueExp (IntValue x))) =
   Just $ prettyText $ intToInt64 x
-isConst (Right (SizeConst v)) =
+isConst (Right (SizeConst v _)) =
   Just $ zEncodeText $ nameToText v
 isConst (Right (SizeMaxConst size_class)) =
   Just $ "max_" <> prettyText size_class
@@ -354,7 +356,8 @@
           mapM_ GC.item body
       kstate = GC.compUserState cstate
 
-      (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel
+      (kernel_consts, (const_defs, const_undefs)) =
+        second unzip $ unzip $ mapMaybe (constDef (kernelName kernel)) $ kernelUses kernel
 
   let (local_memory_bytes, (local_memory_params, local_memory_args, local_memory_init)) =
         second unzip3 $
@@ -458,7 +461,8 @@
     s
       { clGPU = M.insert name (safety, kernel_fun) $ clGPU s,
         clUsedTypes = typesInKernel kernel <> clUsedTypes s,
-        clFailures = kernelFailures kstate
+        clFailures = kernelFailures kstate,
+        clConstants = kernel_consts <> clConstants s
       }
 
   -- The error handling stuff is automatically added later.
@@ -506,17 +510,19 @@
 -- Constants are #defined as macros.  Since a constant name in one
 -- kernel might potentially (although unlikely) also be used for
 -- something else in another kernel, we #undef them after the kernel.
-constDef :: KernelUse -> Maybe (C.BlockItem, C.BlockItem)
-constDef (ConstUse v e) =
+constDef :: Name -> KernelUse -> Maybe ((Name, KernelConstExp), (C.BlockItem, C.BlockItem))
+constDef kernel_name (ConstUse v e) =
   Just
-    ( [C.citem|$escstm:(T.unpack def)|],
-      [C.citem|$escstm:(T.unpack undef)|]
+    ( (nameFromText v', e),
+      ( [C.citem|$escstm:(T.unpack def)|],
+        [C.citem|$escstm:(T.unpack undef)|]
+      )
     )
   where
-    e' = compilePrimExp e
-    def = "#define " <> idText (C.toIdent v mempty) <> " (" <> expText e' <> ")"
+    v' = zEncodeText $ nameToText kernel_name <> "." <> prettyText v
+    def = "#define " <> idText (C.toIdent v mempty) <> " (" <> v' <> ")"
     undef = "#undef " <> idText (C.toIdent v mempty)
-constDef _ = Nothing
+constDef _ _ = Nothing
 
 commonPrelude :: T.Text
 commonPrelude =
@@ -549,14 +555,6 @@
   "#define FUTHARK_HIP\n"
     <> preludeCU
     <> commonPrelude
-
-compilePrimExp :: PrimExp KernelConst -> C.Exp
-compilePrimExp e = runIdentity $ GC.compilePrimExp compileKernelConst e
-  where
-    compileKernelConst (SizeConst key) =
-      pure [C.cexp|$id:(zEncodeText (prettyText key))|]
-    compileKernelConst (SizeMaxConst size_class) =
-      pure [C.cexp|$id:("max_" <> prettyString size_class)|]
 
 kernelArgs :: Kernel -> [KernelArg]
 kernelArgs = mapMaybe useToArg . kernelUses
diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
--- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
+++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs
@@ -37,7 +37,6 @@
 import Futhark.CodeGen.ImpGen
 import Futhark.Error
 import Futhark.IR.MCMem
-import Futhark.MonadFreshNames
 import Futhark.Transform.Rename
 import Prelude hiding (quot, rem)
 
diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs
--- a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs
+++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs
@@ -11,7 +11,6 @@
 import Futhark.CodeGen.ImpGen.Multicore.Base
 import Futhark.CodeGen.ImpGen.Multicore.SegRed (compileSegRed')
 import Futhark.IR.MCMem
-import Futhark.MonadFreshNames
 import Futhark.Transform.Rename (renameLambda)
 import Futhark.Util (chunks, splitFromEnd, takeLast)
 import Futhark.Util.IntegralExp (rem)
diff --git a/src/Futhark/Compiler/Program.hs b/src/Futhark/Compiler/Program.hs
--- a/src/Futhark/Compiler/Program.hs
+++ b/src/Futhark/Compiler/Program.hs
@@ -155,7 +155,6 @@
   -- First we try to find a file of the given name in the search path,
   -- then we look at the builtin library if we have to.  For the
   -- builtins, we don't use the search path.
-  let filepath = includeToFilePath include
   r <- contentsAndModTime filepath vfs
   case (r, lookup prelude_str prelude) of
     (Just (Right (s, mod_time)), _) ->
@@ -167,6 +166,8 @@
     (Nothing, Nothing) ->
       pure $ Left $ ProgError loc $ pretty not_found
   where
+    filepath = includeToFilePath include
+
     prelude_str = "/" Posix.</> includeToString include Posix.<.> "fut"
 
     loaded path s mod_time =
@@ -178,7 +179,7 @@
         }
 
     not_found =
-      "Could not find import " <> E.quote (includeToText include) <> "."
+      "Could not read file " <> E.quote (T.pack filepath) <> "."
 
 handleFile :: ReaderState -> VFS -> LoadedFile T.Text -> IO UncheckedImport
 handleFile state_mvar vfs (LoadedFile file_name import_name file_contents mod_time) = do
diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs
--- a/src/Futhark/Doc/Generator.hs
+++ b/src/Futhark/Doc/Generator.hs
@@ -124,7 +124,7 @@
 
         forEnv env =
           mconcat (map vname' $ M.toList $ envNameMap env)
-            <> mconcat (map forMty $ M.elems $ envSigTable env)
+            <> mconcat (map forMty $ M.elems $ envModTypeTable env)
         forMod (ModEnv env) = forEnv env
         forMod ModFun {} = mempty
         forMty = forMod . mtyMod
@@ -373,7 +373,7 @@
 
 synopsisDec :: S.Set VName -> FileModule -> Dec -> Maybe (DocM Html)
 synopsisDec visible fm dec = case dec of
-  SigDec s -> synopsisModType mempty s
+  ModTypeDec s -> synopsisModType mempty s
   ModDec m -> synopsisMod fm m
   ValDec v -> synopsisValBind v
   TypeDec t -> synopsisType t
@@ -386,8 +386,8 @@
           pure $
             fullRow $
               keyword "open" <> fromString (" <" <> prettyString x <> ">")
-  LocalDec (SigDec s) _
-    | sigName s `S.member` visible ->
+  LocalDec (ModTypeDec s) _
+    | modTypeName s `S.member` visible ->
         synopsisModType (keyword "local" <> " ") s
   LocalDec {} -> Nothing
   ImportDec {} -> Nothing
@@ -404,7 +404,7 @@
     keyword "import "
       <> (H.a ! A.href dest) (fromString (show (includeToString file)))
 synopsisOpened (ModAscript _ se _ _) = Just $ do
-  se' <- synopsisSigExp se
+  se' <- synopsisModTypeExp se
   pure $ "... : " <> se'
 synopsisOpened _ = Nothing
 
@@ -427,18 +427,18 @@
       mconcat (intersperse " -> " $ params' ++ [rettype'])
     )
 
-synopsisModType :: Html -> SigBind -> Maybe (DocM Html)
+synopsisModType :: Html -> ModTypeBind -> Maybe (DocM Html)
 synopsisModType prefix sb = Just $ do
-  let name' = vnameSynopsisDef $ sigName sb
+  let name' = vnameSynopsisDef $ modTypeName sb
   fullRow <$> do
-    se' <- synopsisSigExp $ sigExp sb
+    se' <- synopsisModTypeExp $ modTypeExp sb
     pure $ prefix <> keyword "module type " <> name' <> " = " <> se'
 
 synopsisMod :: FileModule -> ModBind -> Maybe (DocM Html)
 synopsisMod fm (ModBind name ps sig _ _ _) =
   case sig of
-    Nothing -> (proceed <=< envSig) <$> M.lookup name modtable
-    Just (s, _) -> Just $ proceed =<< synopsisSigExp s
+    Nothing -> (proceed <=< envModType) <$> M.lookup name modtable
+    Just (s, _) -> Just $ proceed =<< synopsisModTypeExp s
   where
     proceed sig' = do
       let name' = vnameSynopsisDef name
@@ -446,8 +446,8 @@
       pure $ specRow (keyword "module " <> name') ": " (ps' <> sig')
 
     FileModule _abs Env {envModTable = modtable} _ _ = fm
-    envSig (ModEnv e) = renderEnv e
-    envSig (ModFun (FunSig _ _ (MTy _ m))) = envSig m
+    envModType (ModEnv e) = renderEnv e
+    envModType (ModFun (FunModType _ _ (MTy _ m))) = envModType m
 
 synopsisType :: TypeBind -> Maybe (DocM Html)
 synopsisType tb = Just $ do
@@ -548,7 +548,7 @@
 modParamHtml :: [ModParamBase Info VName] -> DocM Html
 modParamHtml [] = pure mempty
 modParamHtml (ModParam pname psig _ _ : mps) =
-  liftM2 f (synopsisSigExp psig) (modParamHtml mps)
+  liftM2 f (synopsisModTypeExp psig) (modParamHtml mps)
   where
     f se params =
       "("
@@ -558,26 +558,26 @@
         <> ") -> "
         <> params
 
-synopsisSigExp :: SigExpBase Info VName -> DocM Html
-synopsisSigExp e = case e of
-  SigVar v _ _ -> qualNameHtml v
-  SigParens e' _ -> parens <$> synopsisSigExp e'
-  SigSpecs ss _ -> braces . (H.table ! A.class_ "specs") . mconcat <$> mapM synopsisSpec ss
-  SigWith s (TypeRef v ps t _) _ -> do
-    s' <- synopsisSigExp s
+synopsisModTypeExp :: ModTypeExpBase Info VName -> DocM Html
+synopsisModTypeExp e = case e of
+  ModTypeVar v _ _ -> qualNameHtml v
+  ModTypeParens e' _ -> parens <$> synopsisModTypeExp e'
+  ModTypeSpecs ss _ -> braces . (H.table ! A.class_ "specs") . mconcat <$> mapM synopsisSpec ss
+  ModTypeWith s (TypeRef v ps t _) _ -> do
+    s' <- synopsisModTypeExp s
     t' <- typeExpHtml t
     v' <- qualNameHtml v
     let ps' = mconcat $ map ((" " <>) . typeParamHtml) ps
     pure $ s' <> keyword " with " <> v' <> ps' <> " = " <> t'
-  SigArrow Nothing e1 e2 _ ->
-    liftM2 f (synopsisSigExp e1) (synopsisSigExp e2)
+  ModTypeArrow Nothing e1 e2 _ ->
+    liftM2 f (synopsisModTypeExp e1) (synopsisModTypeExp e2)
     where
       f e1' e2' = e1' <> " -> " <> e2'
-  SigArrow (Just v) e1 e2 _ ->
+  ModTypeArrow (Just v) e1 e2 _ ->
     do
       let name = vnameHtml v
-      e1' <- synopsisSigExp e1
-      e2' <- noLink [v] $ synopsisSigExp e2
+      e1' <- synopsisModTypeExp e1
+      e2' <- noLink [v] $ synopsisModTypeExp e2
       pure $ "(" <> name <> ": " <> e1' <> ") -> " <> e2'
 
 keyword :: String -> Html
@@ -625,8 +625,8 @@
         (mconcat (map (" " <>) tparams') <> ": ")
         rettype'
   ModSpec name sig _ _ ->
-    specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig
-  IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e
+    specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisModTypeExp sig
+  IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisModTypeExp e
 
 typeExpHtml :: TypeExp Info VName -> DocM Html
 typeExpHtml e = case e of
@@ -802,7 +802,7 @@
 describeGenericMod ::
   VName ->
   IndexWhat ->
-  SigExp ->
+  ModTypeExp ->
   Maybe DocComment ->
   (Html -> DocM Html) ->
   DocM Html
@@ -812,7 +812,7 @@
   decl_type <- f name'
 
   doc' <- case se of
-    SigSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs
+    ModTypeSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs
     _ -> docHtml doc
 
   let decl_doc = H.dd ! A.class_ "desc_doc" $ doc'
@@ -837,14 +837,14 @@
 describeDec _ (TypeDec vb) =
   Just $
     describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb)
-describeDec _ (SigDec (SigBind name se doc _)) = Just $
+describeDec _ (ModTypeDec (ModTypeBind name se doc _)) = Just $
   describeGenericMod name IndexModuleType se doc $ \name' ->
     pure $ keyword "module type " <> name'
 describeDec _ (ModDec mb) = Just $
   describeGeneric (modName mb) IndexModule (modDoc mb) $ \name' ->
     pure $ keyword "module " <> name'
 describeDec _ OpenDec {} = Nothing
-describeDec visible (LocalDec (SigDec (SigBind name se doc _)) _)
+describeDec visible (LocalDec (ModTypeDec (ModTypeBind name se doc _)) _)
   | name `S.member` visible = Just $
       describeGenericMod name IndexModuleType se doc $ \name' ->
         pure $ keyword "local module type " <> name'
@@ -883,12 +883,12 @@
 describeSpec (ModSpec name se doc _) =
   describeGenericMod name IndexModule se doc $ \name' ->
     case se of
-      SigSpecs {} -> pure $ keyword "module " <> name'
+      ModTypeSpecs {} -> pure $ keyword "module " <> name'
       _ -> do
-        se' <- synopsisSigExp se
+        se' <- synopsisModTypeExp se
         pure $ keyword "module " <> name' <> ": " <> se'
 describeSpec (IncludeSpec sig _) = do
-  sig' <- synopsisSigExp sig
+  sig' <- synopsisModTypeExp sig
   doc' <- docHtml Nothing
   let decl_header =
         (H.dt ! A.class_ "desc_header") $
diff --git a/src/Futhark/IR/GPU/Sizes.hs b/src/Futhark/IR/GPU/Sizes.hs
--- a/src/Futhark/IR/GPU/Sizes.hs
+++ b/src/Futhark/IR/GPU/Sizes.hs
@@ -38,6 +38,12 @@
     SizeLocalMemory
   | -- | A bespoke size with a default.
     SizeBespoke Name Int64
+  | -- | Amount of registers available per workgroup. Mostly
+    -- meaningful for querying the maximum.
+    SizeRegisters
+  | -- | Amount of L2 cache memory, in bytes. Mostly meaningful for
+    -- querying the maximum.
+    SizeCache
   deriving (Eq, Ord, Show)
 
 instance Pretty SizeClass where
@@ -54,6 +60,8 @@
   pretty SizeLocalMemory = "local_memory"
   pretty (SizeBespoke k def) =
     "bespoke" <> parens (pretty k <> comma <+> pretty def)
+  pretty SizeRegisters = "registers"
+  pretty SizeCache = "cache"
 
 -- | The default value for the size.  If 'Nothing', that means the backend gets to decide.
 sizeDefault :: SizeClass -> Maybe Int64
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
@@ -1157,8 +1157,8 @@
 simplifySegOp (SegHist lvl space ops ts kbody) = do
   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)
 
-  (ops', ops_hoisted) <- fmap unzip $
-    forM ops $
+  Engine.localVtable (flip (foldr ST.consume) $ concatMap histDest ops) $ do
+    (ops', ops_hoisted) <- fmap unzip . forM ops $
       \(HistOp w rf arrs nes dims lam) -> do
         w' <- Engine.simplify w
         rf' <- Engine.simplify rf
@@ -1174,12 +1174,12 @@
             op_hoisted
           )
 
-  (kbody', body_hoisted) <- simplifyKernelBody space kbody
+    (kbody', body_hoisted) <- simplifyKernelBody space kbody
 
-  pure
-    ( SegHist lvl' space' ops' ts' kbody',
-      mconcat ops_hoisted <> body_hoisted
-    )
+    pure
+      ( SegHist lvl' space' ops' ts' kbody',
+        mconcat ops_hoisted <> body_hoisted
+      )
   where
     scope = scopeOfSegSpace space
     scope_vtable = ST.fromScope scope
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
@@ -148,7 +148,7 @@
 
 maybeAscript ::
   SrcLoc ->
-  Maybe (SigExp, Info (M.Map VName VName)) ->
+  Maybe (ModTypeExp, Info (M.Map VName VName)) ->
   ModExp ->
   ModExp
 maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc
@@ -311,7 +311,7 @@
     evalModExp
       $ foldr
         addParam
-        (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb)
+        (maybeAscript (srclocOf mb) (modType mb) $ modExp mb)
       $ modParams mb
   mname <- transformName $ modName mb
   pure $ Scope (scopeSubsts $ modScope mod) $ M.singleton mname mod
@@ -331,7 +331,7 @@
       bindingNames [typeAlias tb] $ do
         transformTypeBind tb
         transformDecs ds'
-    SigDec {} : ds' ->
+    ModTypeDec {} : ds' ->
       transformDecs ds'
     ModDec mb : ds' ->
       bindingNames [modName mb] $ do
diff --git a/src/Futhark/LSP/Handlers.hs b/src/Futhark/LSP/Handlers.hs
--- a/src/Futhark/LSP/Handlers.hs
+++ b/src/Futhark/LSP/Handlers.hs
@@ -66,18 +66,13 @@
         file_path = uriToFilePath $ doc ^. uri
     tryReCompile state_mvar file_path
 
-onDocumentOpenHandler :: IORef State -> Handlers (LspM ())
-onDocumentOpenHandler state_mvar =
-  notificationHandler SMethod_TextDocumentDidOpen $ \msg -> do
-    let TNotificationMessage _ _ (DidOpenTextDocumentParams doc) = msg
-        file_path = uriToFilePath $ doc ^. uri
-    logStringStderr <& ("Opened document: " ++ show (doc ^. uri))
-    tryReCompile state_mvar file_path
+-- Some clients (Eglot) sends open/close events whether we want them
+-- or not, so we better be prepared to ignore them.
+onDocumentOpenHandler :: Handlers (LspM ())
+onDocumentOpenHandler = notificationHandler SMethod_TextDocumentDidOpen $ \_ -> pure ()
 
 onDocumentCloseHandler :: Handlers (LspM ())
-onDocumentCloseHandler =
-  notificationHandler SMethod_TextDocumentDidClose $ \_msg ->
-    logStringStderr <& "Closed document"
+onDocumentCloseHandler = notificationHandler SMethod_TextDocumentDidClose $ \_msg -> pure ()
 
 -- Sent by Eglot when first connecting - not sure when else it might
 -- be sent.
@@ -93,7 +88,7 @@
 handlers state_mvar _ =
   mconcat
     [ onInitializeHandler,
-      onDocumentOpenHandler state_mvar,
+      onDocumentOpenHandler,
       onDocumentCloseHandler,
       onDocumentSaveHandler state_mvar,
       onDocumentChangeHandler state_mvar,
diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs
deleted file mode 100644
--- a/src/Futhark/Optimise/InPlaceLowering.hs
+++ /dev/null
@@ -1,416 +0,0 @@
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE UndecidableInstances #-}
-
--- | This module implements an optimisation that moves in-place
--- updates into/before loops where possible, with the end goal of
--- minimising memory copies.  As an example, consider this program:
---
--- @
---   let r =
---     loop (r1 = r0) = for i < n do
---       let a = r1[i]
---       let r1[i] = a * i
---       in r1
---   ...
---   let x = y with [k] <- r in
---   ...
--- @
---
--- We want to turn this into the following:
---
--- @
---   let x0 = y with [k] <- r0
---   loop (x = x0) = for i < n do
---     let a = a[k,i]
---     let x[k,i] = a * i
---     in x
---   let r = x[k] in
---   ...
--- @
---
--- The intent is that we are also going to optimise the new data
--- movement (in the @x0@-binding), possibly by changing how @r0@ is
--- defined.  For the above transformation to be valid, a number of
--- conditions must be fulfilled:
---
---    (1) @r@ must not be consumed after the original in-place update.
---
---    (2) @k@ and @y@ must be available at the beginning of the loop.
---
---    (3) @x@ must be visible whenever @r@ is visible.  (This means
---    that both @x@ and @r@ must be bound in the same t'Body'.)
---
---    (4) If @x@ is consumed at a point after the loop, @r@ must not
---    be used after that point.
---
---    (5) The size of @r1@ is invariant inside the loop.
---
---    (6) The value @r@ must come from something that we can actually
---    optimise (e.g. not a function parameter).
---
---    (7) @y@ (or its aliases) may not be used inside the body of the
---    loop.
---
---    (8) The result of the loop may not alias the merge parameter
---    @r1@.
---
---    (9) @y@ or its aliases may not be used after the loop.
---
--- FIXME: the implementation is not finished yet.  Specifically, not
--- all of the above conditions are checked.
-module Futhark.Optimise.InPlaceLowering
-  ( inPlaceLoweringGPU,
-    inPlaceLoweringSeq,
-    inPlaceLoweringMC,
-  )
-where
-
-import Control.Monad
-import Control.Monad.RWS
-import Data.Map.Strict qualified as M
-import Data.Ord (comparing)
-import Futhark.Analysis.Alias
-import Futhark.Builder
-import Futhark.IR.Aliases
-import Futhark.IR.GPU
-import Futhark.IR.MC
-import Futhark.IR.Seq (Seq)
-import Futhark.Optimise.InPlaceLowering.LowerIntoStm
-import Futhark.Pass
-import Futhark.Util (nubByOrd)
-
--- | Apply the in-place lowering optimisation to the given program.
-inPlaceLoweringGPU :: Pass GPU GPU
-inPlaceLoweringGPU = inPlaceLowering onKernelOp lowerUpdateGPU
-
--- | Apply the in-place lowering optimisation to the given program.
-inPlaceLoweringSeq :: Pass Seq Seq
-inPlaceLoweringSeq = inPlaceLowering pure lowerUpdate
-
--- | Apply the in-place lowering optimisation to the given program.
-inPlaceLoweringMC :: Pass MC MC
-inPlaceLoweringMC = inPlaceLowering onMCOp lowerUpdate
-
--- | Apply the in-place lowering optimisation to the given program.
-inPlaceLowering ::
-  (Constraints rep) =>
-  OnOp rep ->
-  LowerUpdate rep (ForwardingM rep) ->
-  Pass rep rep
-inPlaceLowering onOp lower =
-  Pass "In-place lowering" "Lower in-place updates into loops" $
-    fmap removeProgAliases
-      . intraproceduralTransformationWithConsts optimiseConsts optimiseFunDef
-      . aliasAnalysis
-  where
-    optimiseConsts stms =
-      modifyNameSource $
-        runForwardingM lower onOp $
-          stmsFromList <$> optimiseStms (stmsToList stms) (pure ())
-
-    optimiseFunDef consts fundec =
-      modifyNameSource $
-        runForwardingM lower onOp $
-          descend (stmsToList consts) $
-            bindingFParams (funDefParams fundec) $ do
-              body <- optimiseBody $ funDefBody fundec
-              pure $ fundec {funDefBody = body}
-
-    descend [] m = m
-    descend (stm : stms) m = bindingStm stm $ descend stms m
-
-type Constraints rep = (Buildable rep, AliasableRep rep)
-
-optimiseBody ::
-  (Constraints rep) =>
-  Body (Aliases rep) ->
-  ForwardingM rep (Body (Aliases rep))
-optimiseBody (Body als stms res) = do
-  stms' <- deepen $ optimiseStms (stmsToList stms) $ mapM_ (seen . resSubExp) res
-  pure $ Body als (stmsFromList stms') res
-  where
-    seen Constant {} = pure ()
-    seen (Var v) = seenVar v
-
-optimiseStms ::
-  (Constraints rep) =>
-  [Stm (Aliases rep)] ->
-  ForwardingM rep () ->
-  ForwardingM rep [Stm (Aliases rep)]
-optimiseStms [] m = m >> pure []
-optimiseStms (stm : stms) m = do
-  (stms', bup) <- tapBottomUp $ bindingStm stm $ optimiseStms stms m
-  stm' <- optimiseInStm stm
-  -- XXX: unfortunate that we cannot handle duplicate update values.
-  -- Would be good to improve this.  See inplacelowering6.fut.
-  case nubByOrd (comparing updateValue)
-    . filter ((`notNameIn` bottomUpSeen bup) . updateSource) -- (9)
-    . filter ((`elem` boundHere) . updateValue)
-    $ forwardThese bup of
-    [] -> do
-      checkIfForwardableUpdate stm'
-      pure $ stm' : stms'
-    updates -> do
-      lower <- asks topLowerUpdate
-      scope <- askScope
-
-      -- If we forward any updates, we need to remove them from stms'.
-      let updated_names =
-            map updateName updates
-          notUpdated =
-            not . any (`elem` updated_names) . patNames . stmPat
-
-      -- Condition (5) and (7) are assumed to be checked by
-      -- lowerUpdate.
-      case lower scope stm' updates of
-        Just lowering -> do
-          new_stms <- lowering
-          new_stms' <- optimiseStms new_stms $ tell bup {forwardThese = []}
-          pure $ new_stms' ++ filter notUpdated stms'
-        Nothing -> do
-          checkIfForwardableUpdate stm'
-          pure $ stm' : stms'
-  where
-    boundHere = patNames $ stmPat stm
-
-    checkIfForwardableUpdate (Let pat (StmAux cs _ _) e)
-      | Pat [PatElem v dec] <- pat,
-        BasicOp (Update Unsafe src slice (Var ve)) <- e =
-          maybeForward ve v dec cs src slice
-    checkIfForwardableUpdate stm' =
-      mapM_ seenVar $ namesToList $ freeIn $ stmExp stm'
-
-optimiseInStm :: (Constraints rep) => Stm (Aliases rep) -> ForwardingM rep (Stm (Aliases rep))
-optimiseInStm (Let pat dec e) =
-  Let pat dec <$> optimiseExp e
-
-optimiseExp :: (Constraints rep) => Exp (Aliases rep) -> ForwardingM rep (Exp (Aliases rep))
-optimiseExp (Loop merge form body) =
-  bindingScope (scopeOfLoopForm form) . bindingFParams (map fst merge) $
-    Loop merge form <$> optimiseBody body
-optimiseExp (Op op) = do
-  f <- asks topOnOp
-  Op <$> f op
-optimiseExp e = mapExpM optimise e
-  where
-    optimise =
-      identityMapper
-        { mapOnBody = const optimiseBody
-        }
-
-onSegOp ::
-  (Constraints rep) =>
-  SegOp lvl (Aliases rep) ->
-  ForwardingM rep (SegOp lvl (Aliases rep))
-onSegOp op =
-  bindingScope (scopeOfSegSpace (segSpace op)) $ do
-    let mapper = identitySegOpMapper {mapOnSegOpBody = onKernelBody}
-        onKernelBody kbody = do
-          stms <-
-            deepen $
-              optimiseStms (stmsToList (kernelBodyStms kbody)) $
-                mapM_ seenVar $
-                  namesToList $
-                    freeIn $
-                      kernelBodyResult kbody
-          pure kbody {kernelBodyStms = stmsFromList stms}
-    mapSegOpM mapper op
-
-onMCOp :: OnOp MC
-onMCOp (ParOp par_op op) = ParOp <$> traverse onSegOp par_op <*> onSegOp op
-onMCOp op = pure op
-
-onKernelOp :: OnOp GPU
-onKernelOp (SegOp op) = SegOp <$> onSegOp op
-onKernelOp op = pure op
-
-data Entry rep = Entry
-  { entryNumber :: Int,
-    entryAliases :: Names,
-    entryDepth :: Int,
-    entryOptimisable :: Bool,
-    entryType :: NameInfo (Aliases rep)
-  }
-
-type VTable rep = M.Map VName (Entry rep)
-
-type OnOp rep = Op (Aliases rep) -> ForwardingM rep (Op (Aliases rep))
-
-data TopDown rep = TopDown
-  { topDownCounter :: Int,
-    topDownTable :: VTable rep,
-    topDownDepth :: Int,
-    topLowerUpdate :: LowerUpdate rep (ForwardingM rep),
-    topOnOp :: OnOp rep
-  }
-
-data BottomUp rep = BottomUp
-  { bottomUpSeen :: Names,
-    forwardThese :: [DesiredUpdate (LetDec (Aliases rep))]
-  }
-
-instance Semigroup (BottomUp rep) where
-  BottomUp seen1 forward1 <> BottomUp seen2 forward2 =
-    BottomUp (seen1 <> seen2) (forward1 <> forward2)
-
-instance Monoid (BottomUp rep) where
-  mempty = BottomUp mempty mempty
-
-newtype ForwardingM rep a = ForwardingM (RWS (TopDown rep) (BottomUp rep) VNameSource a)
-  deriving
-    ( Monad,
-      Applicative,
-      Functor,
-      MonadReader (TopDown rep),
-      MonadWriter (BottomUp rep),
-      MonadState VNameSource
-    )
-
-instance MonadFreshNames (ForwardingM rep) where
-  getNameSource = get
-  putNameSource = put
-
-instance (Constraints rep) => HasScope (Aliases rep) (ForwardingM rep) where
-  askScope = M.map entryType <$> asks topDownTable
-
-runForwardingM ::
-  LowerUpdate rep (ForwardingM rep) ->
-  OnOp rep ->
-  ForwardingM rep a ->
-  VNameSource ->
-  (a, VNameSource)
-runForwardingM f g (ForwardingM m) src =
-  let (x, src', _) = runRWS m emptyTopDown src
-   in (x, src')
-  where
-    emptyTopDown =
-      TopDown
-        { topDownCounter = 0,
-          topDownTable = M.empty,
-          topDownDepth = 0,
-          topLowerUpdate = f,
-          topOnOp = g
-        }
-
-bindingParams ::
-  (dec -> NameInfo (Aliases rep)) ->
-  [Param dec] ->
-  ForwardingM rep a ->
-  ForwardingM rep a
-bindingParams f params = local $ \(TopDown n vtable d x y) ->
-  let entry fparam =
-        ( paramName fparam,
-          Entry n mempty d False $ f $ paramDec fparam
-        )
-      entries = M.fromList $ map entry params
-   in TopDown (n + 1) (M.union entries vtable) d x y
-
-bindingFParams ::
-  [FParam (Aliases rep)] ->
-  ForwardingM rep a ->
-  ForwardingM rep a
-bindingFParams = bindingParams FParamName
-
-bindingScope ::
-  Scope (Aliases rep) ->
-  ForwardingM rep a ->
-  ForwardingM rep a
-bindingScope scope = local $ \(TopDown n vtable d x y) ->
-  let entries = M.map entry scope
-      infoAliases (LetName (aliases, _)) = unAliases aliases
-      infoAliases _ = mempty
-      entry info = Entry n (infoAliases info) d False info
-   in TopDown (n + 1) (entries <> vtable) d x y
-
-bindingStm ::
-  Stm (Aliases rep) ->
-  ForwardingM rep a ->
-  ForwardingM rep a
-bindingStm (Let pat _ _) = local $ \(TopDown n vtable d x y) ->
-  let entries = M.fromList $ map entry $ patElems pat
-      entry patElem =
-        let (aliases, _) = patElemDec patElem
-         in ( patElemName patElem,
-              Entry n (unAliases aliases) d True $ LetName $ patElemDec patElem
-            )
-   in TopDown (n + 1) (M.union entries vtable) d x y
-
-bindingNumber :: VName -> ForwardingM rep Int
-bindingNumber name = do
-  res <- asks $ fmap entryNumber . M.lookup name . topDownTable
-  case res of
-    Just n -> pure n
-    Nothing ->
-      error $
-        "bindingNumber: variable "
-          ++ prettyString name
-          ++ " not found."
-
-deepen :: ForwardingM rep a -> ForwardingM rep a
-deepen = local $ \env -> env {topDownDepth = topDownDepth env + 1}
-
-areAvailableBefore :: Names -> VName -> ForwardingM rep Bool
-areAvailableBefore names point = do
-  pointN <- bindingNumber point
-  nameNs <- mapM bindingNumber $ namesToList names
-  pure $ all (< pointN) nameNs
-
-isInCurrentBody :: VName -> ForwardingM rep Bool
-isInCurrentBody name = do
-  current <- asks topDownDepth
-  res <- asks $ fmap entryDepth . M.lookup name . topDownTable
-  case res of
-    Just d -> pure $ d == current
-    Nothing ->
-      error $
-        "isInCurrentBody: variable "
-          ++ prettyString name
-          ++ " not found."
-
-isOptimisable :: VName -> ForwardingM rep Bool
-isOptimisable name = do
-  res <- asks $ fmap entryOptimisable . M.lookup name . topDownTable
-  case res of
-    Just b -> pure b
-    Nothing ->
-      error $
-        "isOptimisable: variable "
-          ++ prettyString name
-          ++ " not found."
-
-seenVar :: forall rep. VName -> ForwardingM rep ()
-seenVar name = do
-  aliases <-
-    asks $
-      maybe mempty entryAliases
-        . M.lookup name
-        . topDownTable
-  tell $ (mempty :: BottomUp rep) {bottomUpSeen = oneName name <> aliases}
-
-tapBottomUp :: ForwardingM rep a -> ForwardingM rep (a, BottomUp rep)
-tapBottomUp m = do
-  (x, bup) <- listen m
-  pure (x, bup)
-
-maybeForward ::
-  (Constraints rep) =>
-  VName ->
-  VName ->
-  LetDec (Aliases rep) ->
-  Certs ->
-  VName ->
-  Slice SubExp ->
-  ForwardingM rep ()
-maybeForward v dest_nm dest_dec cs src slice = do
-  -- Checks condition (2)
-  available <-
-    (freeIn src <> freeIn slice <> freeIn cs)
-      `areAvailableBefore` v
-  -- Check condition (3)
-  samebody <- isInCurrentBody v
-  -- Check condition (6)
-  optimisable <- isOptimisable v
-  not_prim <- not . primType <$> lookupType v
-  when (available && samebody && optimisable && not_prim) $ do
-    let fwd = DesiredUpdate dest_nm dest_dec cs src slice v
-    tell mempty {forwardThese = [fwd]}
diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
deleted file mode 100644
--- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
+++ /dev/null
@@ -1,376 +0,0 @@
-{-# LANGUAGE TypeFamilies #-}
-
-module Futhark.Optimise.InPlaceLowering.LowerIntoStm
-  ( lowerUpdateGPU,
-    lowerUpdate,
-    LowerUpdate,
-    DesiredUpdate (..),
-  )
-where
-
-import Control.Monad
-import Control.Monad.Writer
-import Data.Either
-import Data.List (find, unzip5)
-import Data.Maybe (isNothing, mapMaybe)
-import Futhark.Analysis.PrimExp.Convert
-import Futhark.Construct
-import Futhark.IR.Aliases
-import Futhark.IR.GPU
-import Futhark.Optimise.InPlaceLowering.SubstituteIndices
-
-data DesiredUpdate dec = DesiredUpdate
-  { -- | Name of result.
-    updateName :: VName,
-    -- | Type of result.
-    updateType :: dec,
-    updateCerts :: Certs,
-    updateSource :: VName,
-    updateIndices :: Slice SubExp,
-    updateValue :: VName
-  }
-  deriving (Show)
-
-instance Functor DesiredUpdate where
-  f `fmap` u = u {updateType = f $ updateType u}
-
-updateHasValue :: VName -> DesiredUpdate dec -> Bool
-updateHasValue name = (name ==) . updateValue
-
-type LowerUpdate rep m =
-  Scope (Aliases rep) ->
-  Stm (Aliases rep) ->
-  [DesiredUpdate (LetDec (Aliases rep))] ->
-  Maybe (m [Stm (Aliases rep)])
-
-lowerUpdate ::
-  ( MonadFreshNames m,
-    Buildable rep,
-    LetDec rep ~ Type,
-    AliasableRep rep
-  ) =>
-  LowerUpdate rep m
-lowerUpdate scope (Let pat aux (Loop merge form body)) updates = do
-  canDo <- lowerUpdateIntoLoop scope updates pat merge form body
-  Just $ do
-    (prestms, poststms, pat', merge', body') <- canDo
-    pure $
-      prestms
-        ++ [ certify (stmAuxCerts aux) $
-               mkLet pat' $
-                 Loop merge' form body'
-           ]
-        ++ poststms
-lowerUpdate
-  _
-  (Let pat aux (BasicOp (SubExp (Var v))))
-  [DesiredUpdate bindee_nm bindee_dec cs src (Slice is) val]
-    | patNames pat == [src] =
-        let is' = fullSlice (typeOf bindee_dec) is
-         in Just . pure $
-              [ certify (stmAuxCerts aux <> cs) $
-                  mkLet [Ident bindee_nm $ typeOf bindee_dec] $
-                    BasicOp $
-                      Update Unsafe v is' $
-                        Var val
-              ]
-lowerUpdate _ _ _ =
-  Nothing
-
-lowerUpdateGPU :: (MonadFreshNames m) => LowerUpdate GPU m
-lowerUpdateGPU
-  scope
-  (Let pat aux (Op (SegOp (SegMap lvl space ts kbody))))
-  updates
-    | all ((`elem` patNames pat) . updateValue) updates,
-      not source_used_in_kbody = do
-        mk <- lowerUpdatesIntoSegMap scope pat ts updates space kbody
-        Just $ do
-          (pat', ts', kbody', poststms) <- mk
-          let cs = stmAuxCerts aux <> foldMap updateCerts updates
-          pure $
-            certify cs (Let pat' aux $ Op $ SegOp $ SegMap lvl space ts' kbody')
-              : stmsToList poststms
-    where
-      -- This check is a bit more conservative than ideal.  In a perfect
-      -- world, we would allow indexing a[i,j] if the update is also
-      -- to exactly a[i,j], as that will not create cross-iteration
-      -- dependencies.  (Although the type checker wouldn't be able to
-      -- permit this anyway.)
-      source_used_in_kbody =
-        mconcat (map (`lookupAliases` scope) (namesToList (freeIn kbody)))
-          `namesIntersect` mconcat (map ((`lookupAliases` scope) . updateSource) updates)
-lowerUpdateGPU scope stm updates = lowerUpdate scope stm updates
-
-lowerUpdatesIntoSegMap ::
-  (MonadFreshNames m) =>
-  Scope (Aliases GPU) ->
-  Pat (LetDec (Aliases GPU)) ->
-  [Type] ->
-  [DesiredUpdate (LetDec (Aliases GPU))] ->
-  SegSpace ->
-  KernelBody (Aliases GPU) ->
-  Maybe
-    ( m
-        ( Pat (LetDec (Aliases GPU)),
-          [Type],
-          KernelBody (Aliases GPU),
-          Stms (Aliases GPU)
-        )
-    )
-lowerUpdatesIntoSegMap scope pat ret_ts updates kspace kbody = do
-  -- The updates are all-or-nothing.  Being more liberal would require
-  -- changes to the in-place-lowering pass itself.
-  mk <- mapM onRet (zip3 (patElems pat) ret_ts (kernelBodyResult kbody))
-  pure $ do
-    (pes, ret_ts', bodystms, krets, poststms) <- unzip5 <$> sequence mk
-    pure
-      ( Pat pes,
-        ret_ts',
-        kbody
-          { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms,
-            kernelBodyResult = krets
-          },
-        mconcat poststms
-      )
-  where
-    (gtids, _dims) = unzip $ unSegSpace kspace
-
-    onRet (PatElem v v_dec, _, ret)
-      | Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <-
-          find ((== v) . updateValue) updates = do
-          Returns _ cs se <- Just ret
-
-          -- The slice we're writing per thread must fully cover the
-          -- underlying dimensions.
-          guard $
-            let (dims', slice') =
-                  unzip . drop (length gtids) . filter (isNothing . dimFix . snd) $
-                    zip (arrayDims (typeOf bindee_dec)) (unSlice slice)
-             in isFullSlice (Shape dims') (Slice slice')
-
-          Just $ do
-            (slice', bodystms) <-
-              flip runBuilderT scope $
-                traverse (toSubExp "index") $
-                  fixSlice (fmap pe64 slice) $
-                    map (pe64 . Var) gtids
-
-            let ret' = WriteReturns cs src [(fullSlice (typeOf bindee_dec) (map DimFix slice'), se)]
-
-            v_aliased <- newName v
-
-            pure
-              ( PatElem bindee_nm bindee_dec,
-                typeOf bindee_dec,
-                bodystms,
-                ret',
-                stmsFromList
-                  [ mkLet [Ident v_aliased $ typeOf v_dec] $ BasicOp $ Index bindee_nm slice,
-                    mkLet [Ident v $ typeOf v_dec] $ BasicOp $ Replicate mempty $ Var v_aliased
-                  ]
-              )
-    onRet (pe, ret_t, ret) =
-      Just $ pure (pe, ret_t, mempty, ret, mempty)
-
-lowerUpdateIntoLoop ::
-  ( Buildable rep,
-    BuilderOps rep,
-    Aliased rep,
-    LetDec rep ~ (als, Type),
-    MonadFreshNames m
-  ) =>
-  Scope rep ->
-  [DesiredUpdate (LetDec rep)] ->
-  Pat (LetDec rep) ->
-  [(FParam rep, SubExp)] ->
-  LoopForm ->
-  Body rep ->
-  Maybe
-    ( m
-        ( [Stm rep],
-          [Stm rep],
-          [Ident],
-          [(FParam rep, SubExp)],
-          Body rep
-        )
-    )
-lowerUpdateIntoLoop scope updates pat val form body = do
-  -- Algorithm:
-  --
-  --   0) Map each result of the loop body to a corresponding in-place
-  --      update, if one exists.
-  --
-  --   1) Create new merge variables corresponding to the arrays being
-  --      updated; extend the pattern and the @res@ list with these,
-  --      and remove the parts of the result list that have a
-  --      corresponding in-place update.
-  --
-  --      (The creation of the new merge variable identifiers is
-  --      actually done at the same time as step (0)).
-  --
-  --   2) Create in-place updates at the end of the loop body.
-  --
-  --   3) Create index expressions that read back the values written
-  --      in (2).  If the merge parameter corresponding to this value
-  --      is unique, also @copy@ this value.
-  --
-  --   4) Update the result of the loop body to properly pass the new
-  --      arrays and indexed elements to the next iteration of the
-  --      loop.
-  --
-  -- We also check that the merge parameters we work with have
-  -- loop-invariant shapes.
-
-  -- Safety condition (8).
-  forM_ (zip val $ bodyAliases body) $ \((p, _), als) ->
-    guard $ paramName p `notNameIn` als
-
-  mk_in_place_map <- summariseLoop scope updates usedInBody resmap val
-
-  Just $ do
-    in_place_map <- mk_in_place_map
-    (val', prestms, poststms) <- mkMerges in_place_map
-    let valpat = mkResAndPat in_place_map
-        idxsubsts = indexSubstitutions in_place_map
-    (idxsubsts', newstms) <- substituteIndices idxsubsts $ bodyStms body
-    (body_res, res_stms) <- manipulateResult in_place_map idxsubsts'
-    let body' = mkBody (newstms <> res_stms) body_res
-    pure (prestms, poststms, valpat, val', body')
-  where
-    usedInBody =
-      mconcat $ map (`lookupAliases` scope) $ namesToList $ freeIn body <> freeIn form
-    resmap = zip (bodyResult body) $ patIdents pat
-
-    mkMerges ::
-      (MonadFreshNames m, Buildable rep) =>
-      [LoopResultSummary (als, Type)] ->
-      m ([(Param DeclType, SubExp)], [Stm rep], [Stm rep])
-    mkMerges summaries = do
-      ((origmerge, extramerge), (prestms, poststms)) <-
-        runWriterT $ partitionEithers <$> mapM mkMerge summaries
-      pure (origmerge ++ extramerge, prestms, poststms)
-
-    mkMerge summary
-      | Just (update, mergename, mergedec) <- relatedUpdate summary = do
-          source <- newVName "modified_source"
-          precopy <- newVName $ baseString (updateValue update) <> "_precopy"
-          let source_t = snd $ updateType update
-              elm_t = source_t `setArrayDims` sliceDims (updateIndices update)
-          tell
-            ( [ mkLet [Ident source source_t] . BasicOp
-                  $ Update
-                    Unsafe
-                    (updateSource update)
-                    (fullSlice source_t $ unSlice $ updateIndices update)
-                  $ snd
-                  $ mergeParam summary
-              ],
-              [ mkLet [Ident precopy elm_t] . BasicOp $
-                  Index
-                    (updateName update)
-                    (fullSlice source_t $ unSlice $ updateIndices update),
-                mkLet [Ident (updateValue update) elm_t] $ BasicOp $ Replicate mempty $ Var precopy
-              ]
-            )
-          pure $
-            Right
-              ( Param mempty mergename (toDecl (typeOf mergedec) Unique),
-                Var source
-              )
-      | otherwise = pure $ Left $ mergeParam summary
-
-    mkResAndPat summaries =
-      let (origpat, extrapat) = partitionEithers $ map mkResAndPat' summaries
-       in origpat ++ extrapat
-
-    mkResAndPat' summary
-      | Just (update, _, _) <- relatedUpdate summary =
-          Right (Ident (updateName update) (snd $ updateType update))
-      | otherwise =
-          Left (inPatAs summary)
-
-summariseLoop ::
-  ( Aliased rep,
-    MonadFreshNames m
-  ) =>
-  Scope rep ->
-  [DesiredUpdate (als, Type)] ->
-  Names ->
-  [(SubExpRes, Ident)] ->
-  [(Param DeclType, SubExp)] ->
-  Maybe (m [LoopResultSummary (als, Type)])
-summariseLoop scope updates usedInBody resmap merge =
-  sequence <$> zipWithM summariseLoopResult resmap merge
-  where
-    summariseLoopResult (se, v) (fparam, mergeinit)
-      | Just update <- find (updateHasValue $ identName v) updates =
-          -- Safety condition (7)
-          if usedInBody `namesIntersect` lookupAliases (updateSource update) scope
-            then Nothing
-            else
-              if hasLoopInvariantShape fparam
-                then Just $ do
-                  lowered_array <- newVName "lowered_array"
-                  pure
-                    LoopResultSummary
-                      { resultSubExp = se,
-                        inPatAs = v,
-                        mergeParam = (fparam, mergeinit),
-                        relatedUpdate =
-                          Just
-                            ( update,
-                              lowered_array,
-                              updateType update
-                            )
-                      }
-                else Nothing
-    summariseLoopResult _ _ =
-      Nothing -- XXX: conservative; but this entire pass is going away.
-    hasLoopInvariantShape = all loopInvariant . arrayDims . paramType
-
-    merge_param_names = map (paramName . fst) merge
-
-    loopInvariant (Var v) = v `notElem` merge_param_names
-    loopInvariant Constant {} = True
-
-data LoopResultSummary dec = LoopResultSummary
-  { resultSubExp :: SubExpRes,
-    inPatAs :: Ident,
-    mergeParam :: (Param DeclType, SubExp),
-    relatedUpdate :: Maybe (DesiredUpdate dec, VName, dec)
-  }
-  deriving (Show)
-
-indexSubstitutions :: (Typed dec) => [LoopResultSummary dec] -> IndexSubstitutions
-indexSubstitutions = mapMaybe getSubstitution
-  where
-    getSubstitution res = do
-      (DesiredUpdate _ _ cs _ is _, nm, dec) <- relatedUpdate res
-      let name = paramName $ fst $ mergeParam res
-      pure (name, (cs, nm, typeOf dec, is))
-
-manipulateResult ::
-  (Buildable rep, MonadFreshNames m) =>
-  [LoopResultSummary (LetDec rep)] ->
-  IndexSubstitutions ->
-  m (Result, Stms rep)
-manipulateResult summaries substs = do
-  let (orig_ses, updated_ses) = partitionEithers $ map unchangedRes summaries
-  (subst_ses, res_stms) <- runWriterT $ zipWithM substRes updated_ses substs
-  pure (orig_ses ++ subst_ses, stmsFromList res_stms)
-  where
-    unchangedRes summary =
-      case relatedUpdate summary of
-        Nothing -> Left $ resultSubExp summary
-        Just _ -> Right $ resultSubExp summary
-    substRes (SubExpRes res_cs (Var res_v)) (subst_v, (_, nm, _, _))
-      | res_v == subst_v =
-          pure $ SubExpRes res_cs $ Var nm
-    substRes (SubExpRes res_cs res_se) (_, (cs, nm, dec, Slice is)) = do
-      v' <- newIdent' (++ "_updated") $ Ident nm $ typeOf dec
-      tell
-        [ certify (res_cs <> cs) . mkLet [v'] . BasicOp $
-            Update Unsafe nm (fullSlice (typeOf dec) is) res_se
-        ]
-      pure $ varRes $ identName v'
diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs
deleted file mode 100644
--- a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs
+++ /dev/null
@@ -1,191 +0,0 @@
-{-# LANGUAGE TypeFamilies #-}
-
--- | This module exports facilities for transforming array accesses in
--- a list of 'Stm's (intended to be the bindings in a body).  The
--- idea is that you can state that some variable @x@ is in fact an
--- array indexing @v[i0,i1,...]@.
-module Futhark.Optimise.InPlaceLowering.SubstituteIndices
-  ( substituteIndices,
-    IndexSubstitution,
-    IndexSubstitutions,
-  )
-where
-
-import Control.Monad
-import Data.Map.Strict qualified as M
-import Futhark.Construct
-import Futhark.IR
-import Futhark.IR.Prop.Aliases
-import Futhark.Transform.Substitute
-
--- | Essentially the components of an 'Index' expression.
-type IndexSubstitution = (Certs, VName, Type, Slice SubExp)
-
--- | A mapping from variable names to the indexing operation they
--- should be replaced with.
-type IndexSubstitutions = [(VName, IndexSubstitution)]
-
-typeEnvFromSubstitutions :: (LParamInfo rep ~ Type) => IndexSubstitutions -> Scope rep
-typeEnvFromSubstitutions = M.fromList . map (fromSubstitution . snd)
-  where
-    fromSubstitution (_, name, t, _) =
-      (name, LParamName t)
-
--- | Perform the substitution.
-substituteIndices ::
-  ( MonadFreshNames m,
-    BuilderOps rep,
-    Buildable rep,
-    Aliased rep
-  ) =>
-  IndexSubstitutions ->
-  Stms rep ->
-  m (IndexSubstitutions, Stms rep)
-substituteIndices substs stms =
-  runBuilderT (substituteIndicesInStms substs stms) types
-  where
-    types = typeEnvFromSubstitutions substs
-
-substituteIndicesInStms ::
-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>
-  IndexSubstitutions ->
-  Stms (Rep m) ->
-  m IndexSubstitutions
-substituteIndicesInStms = foldM substituteIndicesInStm
-
-substituteIndicesInStm ::
-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>
-  IndexSubstitutions ->
-  Stm (Rep m) ->
-  m IndexSubstitutions
-substituteIndicesInStm substs (Let pat _ (BasicOp (Rearrange perm v)))
-  | Just (cs, src, src_t, is) <- lookup v substs,
-    [v'] <- patNames pat = do
-      let extra_dims = arrayRank src_t - length perm
-          perm' = [0 .. extra_dims - 1] ++ map (+ extra_dims) perm
-      src' <-
-        letExp (baseString v' <> "_subst") $ BasicOp $ Rearrange perm' src
-      src_t' <- lookupType src'
-      pure $ (v', (cs, src', src_t', is)) : substs
-substituteIndicesInStm substs (Let pat rep e) = do
-  e' <- substituteIndicesInExp substs e
-  addStm $ Let pat rep e'
-  pure substs
-
-substituteIndicesInExp ::
-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>
-  IndexSubstitutions ->
-  Exp (Rep m) ->
-  m (Exp (Rep m))
-substituteIndicesInExp substs (Op op) = do
-  let used_in_op = filter ((`nameIn` freeIn op) . fst) substs
-  var_substs <- fmap mconcat $
-    forM used_in_op $ \(v, (cs, src, src_dec, Slice is)) -> do
-      v' <-
-        certifying cs $
-          letExp (baseString src <> "_op_idx") $
-            BasicOp $
-              Index src $
-                fullSlice (typeOf src_dec) is
-      pure $ M.singleton v v'
-  pure $ Op $ substituteNames var_substs op
-substituteIndicesInExp substs e = do
-  substs' <- copyAnyConsumed e
-  let substitute =
-        identityMapper
-          { mapOnSubExp = substituteIndicesInSubExp substs',
-            mapOnVName = substituteIndicesInVar substs',
-            mapOnBody = const $ substituteIndicesInBody substs'
-          }
-
-  mapExpM substitute e
-  where
-    copyAnyConsumed =
-      let consumingSubst substs' v
-            | Just (cs2, src2, src2dec, is2) <- lookup v substs = do
-                row <-
-                  certifying cs2 $
-                    letSubExp (baseString v ++ "_row") $
-                      BasicOp $
-                        Index src2 $
-                          fullSlice (typeOf src2dec) (unSlice is2)
-                row_copy <-
-                  letExp (baseString v ++ "_row_copy") . BasicOp $
-                    Replicate mempty row
-                pure $
-                  update
-                    v
-                    v
-                    ( mempty,
-                      row_copy,
-                      src2dec
-                        `setType` ( typeOf src2dec
-                                      `setArrayDims` sliceDims is2
-                                  ),
-                      Slice []
-                    )
-                    substs'
-          consumingSubst substs' _ =
-            pure substs'
-       in foldM consumingSubst substs . namesToList . consumedInExp
-
-substituteIndicesInSubExp ::
-  (MonadBuilder m) =>
-  IndexSubstitutions ->
-  SubExp ->
-  m SubExp
-substituteIndicesInSubExp substs (Var v) =
-  Var <$> substituteIndicesInVar substs v
-substituteIndicesInSubExp _ se =
-  pure se
-
-substituteIndicesInVar ::
-  (MonadBuilder m) =>
-  IndexSubstitutions ->
-  VName ->
-  m VName
-substituteIndicesInVar substs v
-  | Just (cs2, src2, _, Slice []) <- lookup v substs =
-      certifying cs2 $
-        letExp (baseString src2) $
-          BasicOp $
-            SubExp $
-              Var src2
-  | Just (cs2, src2, src2_dec, Slice is2) <- lookup v substs =
-      certifying cs2 $
-        letExp (baseString src2 <> "_v_idx") $
-          BasicOp $
-            Index src2 $
-              fullSlice (typeOf src2_dec) is2
-  | otherwise =
-      pure v
-
-substituteIndicesInBody ::
-  (MonadBuilder m, Buildable (Rep m), Aliased (Rep m)) =>
-  IndexSubstitutions ->
-  Body (Rep m) ->
-  m (Body (Rep m))
-substituteIndicesInBody substs (Body _ stms res) = do
-  (substs', stms') <-
-    inScopeOf stms $
-      collectStms $
-        substituteIndicesInStms substs stms
-  (res', res_stms) <-
-    inScopeOf stms' $
-      collectStms $
-        mapM (onSubExpRes substs') res
-  mkBodyM (stms' <> res_stms) res'
-  where
-    onSubExpRes substs' (SubExpRes cs se) =
-      SubExpRes cs <$> substituteIndicesInSubExp substs' se
-
-update ::
-  VName ->
-  VName ->
-  IndexSubstitution ->
-  IndexSubstitutions ->
-  IndexSubstitutions
-update needle name subst ((othername, othersubst) : substs)
-  | needle == othername = (name, subst) : substs
-  | otherwise = (othername, othersubst) : update needle name subst substs
-update needle _ _ [] = error $ "Cannot find substitution for " ++ prettyString needle
diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs
--- a/src/Futhark/Passes.hs
+++ b/src/Futhark/Passes.hs
@@ -25,7 +25,6 @@
 import Futhark.Optimise.Fusion
 import Futhark.Optimise.GenRedOpt
 import Futhark.Optimise.HistAccs
-import Futhark.Optimise.InPlaceLowering
 import Futhark.Optimise.InliningDeadFun
 import Futhark.Optimise.MemoryBlockMerging qualified as MemoryBlockMerging
 import Futhark.Optimise.MergeGPUBodies
@@ -103,7 +102,6 @@
         mergeGPUBodies,
         simplifyGPU, -- Cleanup merged GPUBody kernels.
         sinkGPU, -- Sink reads within GPUBody kernels.
-        inPlaceLoweringGPU,
         babysitKernels,
         -- Important to simplify after babysitting in order to fix up
         -- redundant manifests.
@@ -118,8 +116,7 @@
   standardPipeline
     >>> onePass firstOrderTransform
     >>> passes
-      [ simplifySeq,
-        inPlaceLoweringSeq
+      [ simplifySeq
       ]
 
 -- | Run 'seqPipeline', then add memory information (and
@@ -183,8 +180,7 @@
         unstreamMC,
         performCSE True,
         simplifyMC,
-        sinkMC,
-        inPlaceLoweringMC
+        sinkMC
       ]
 
 -- | Run 'mcPipeline' and then add memory information.
diff --git a/src/Futhark/Util/IntegralExp.hs b/src/Futhark/Util/IntegralExp.hs
--- a/src/Futhark/Util/IntegralExp.hs
+++ b/src/Futhark/Util/IntegralExp.hs
@@ -39,6 +39,9 @@
   divUp x y =
     (x + y - 1) `Futhark.Util.IntegralExp.div` y
 
+  nextMul :: e -> e -> e
+  nextMul x y = x `divUp` y * y
+
 -- | This wrapper allows you to use a type that is an instance of the
 -- true class whenever the simile class is required.
 newtype Wrapped a = Wrapped {wrappedValue :: a}
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
@@ -90,14 +90,14 @@
 
 type Stack = [StackFrame]
 
-type Sizes = M.Map VName Int64
+type Exts = M.Map VName Value
 
 -- | The monad in which evaluation takes place.
 newtype EvalM a
   = EvalM
       ( ReaderT
           (Stack, M.Map ImportName Env)
-          (StateT Sizes (F ExtOp))
+          (StateT Exts (F ExtOp))
           a
       )
   deriving
@@ -106,7 +106,7 @@
       Functor,
       MonadFree ExtOp,
       MonadReader (Stack, M.Map ImportName Env),
-      MonadState Sizes
+      MonadState Exts
     )
 
 runEvalM :: M.Map ImportName Env -> EvalM a -> F ExtOp a
@@ -129,11 +129,11 @@
 lookupImport :: ImportName -> EvalM (Maybe Env)
 lookupImport f = asks $ M.lookup f . snd
 
-putExtSize :: VName -> Int64 -> EvalM ()
+putExtSize :: VName -> Value -> EvalM ()
 putExtSize v x = modify $ M.insert v x
 
-getSizes :: EvalM Sizes
-getSizes = get
+getExts :: EvalM Exts
+getExts = get
 
 -- | Disregard any existential sizes computed during this action.
 -- This is used so that existentials computed during one iteration of
@@ -145,8 +145,13 @@
   put s
   pure x
 
-extSizeEnv :: EvalM Env
-extSizeEnv = i64Env <$> getSizes
+extEnv :: EvalM Env
+extEnv = valEnv . M.map f <$> getExts
+  where
+    f v =
+      ( Nothing,
+        v
+      )
 
 valueStructType :: ValueType -> StructType
 valueStructType = first $ flip sizeFromInteger mempty . toInteger
@@ -296,7 +301,8 @@
 -- | An expression evaluator that embeds an environment.
 type Eval = Exp -> EvalM Value
 
--- | A TermValue with a 'Nothing' type annotation is an intrinsic.
+-- | A TermValue with a 'Nothing' type annotation is an intrinsic or
+-- an existential.
 data TermBinding
   = TermValue (Maybe T.BoundV) Value
   | -- | A polymorphic value that must be instantiated.  The
@@ -636,7 +642,7 @@
 
 evalWithExts :: Env -> EvalM Eval
 evalWithExts env = do
-  size_env <- extSizeEnv
+  size_env <- extEnv
   pure $ eval $ size_env <> env
 
 -- | Evaluate all possible sizes, except those that contain free
@@ -750,16 +756,17 @@
 evalArg env e ext = do
   v <- eval env e
   case ext of
-    Just ext' -> putExtSize ext' $ asInt64 v
-    Nothing -> pure ()
+    Just ext' -> putExtSize ext' v
+    _ -> pure ()
   pure v
 
 returned :: Env -> TypeBase Size als -> [VName] -> Value -> EvalM Value
 returned _ _ [] v = pure v
 returned env ret retext v = do
-  mapM_ (uncurry putExtSize) . M.toList $
-    resolveExistentials retext (expandType env $ toStruct ret) $
-      valueShape v
+  mapM_ (uncurry putExtSize . second (ValuePrim . SignedValue . Int64Value))
+    . M.toList
+    $ resolveExistentials retext (expandType env $ toStruct ret)
+    $ valueShape v
   pure v
 
 evalAppExp :: Env -> AppExp -> EvalM Value
@@ -1193,7 +1200,7 @@
 evalDec :: Env -> Dec -> EvalM Env
 evalDec env (ValDec (ValBind _ v _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do
   binding <- evalFunctionBinding env tparams ps ret fbody
-  sizes <- extSizeEnv
+  sizes <- extEnv
   pure $
     env {envTerm = M.insert v binding $ envTerm env} <> sizes
 evalDec env (OpenDec me _) = do
@@ -1204,7 +1211,7 @@
 evalDec env (ImportDec name name' loc) =
   evalDec env $ LocalDec (OpenDec (ModImport name name' loc) loc) loc
 evalDec env (LocalDec d _) = evalDec env d
-evalDec env SigDec {} = pure env
+evalDec env ModTypeDec {} = pure env
 evalDec env (TypeDec (TypeBind v l ps _ (Info (RetType dims t)) _ _)) = do
   let abbr = T.TypeAbbr l ps . RetType dims $ expandType env t
   pure env {envType = M.insert v abbr $ envType env}
@@ -1941,7 +1948,7 @@
     -- We need to extract any new existential sizes and add them as
     -- ordinary bindings to the context, or we will not be able to
     -- look up their values later.
-    sizes <- extSizeEnv
+    sizes <- extEnv
     pure $ env <> sizes
 
 interpretDec :: Ctx -> Dec -> F ExtOp Ctx
diff --git a/src/Language/Futhark/Parser/Monad.hs b/src/Language/Futhark/Parser/Monad.hs
--- a/src/Language/Futhark/Parser/Monad.hs
+++ b/src/Language/Futhark/Parser/Monad.hs
@@ -57,7 +57,7 @@
 addDoc :: DocComment -> UncheckedDec -> UncheckedDec
 addDoc doc (ValDec val) = ValDec (val {valBindDoc = Just doc})
 addDoc doc (TypeDec tp) = TypeDec (tp {typeDoc = Just doc})
-addDoc doc (SigDec sig) = SigDec (sig {sigDoc = Just doc})
+addDoc doc (ModTypeDec sig) = ModTypeDec (sig {modTypeDoc = Just doc})
 addDoc doc (ModDec mod) = ModDec (mod {modDoc = Just doc})
 addDoc _ dec = dec
 
diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y
--- a/src/Language/Futhark/Parser/Parser.y
+++ b/src/Language/Futhark/Parser/Parser.y
@@ -217,7 +217,7 @@
 Dec_ :: { UncheckedDec }
     : Val               { ValDec $1 }
     | TypeAbbr          { TypeDec $1 }
-    | SigBind           { SigDec $1 }
+    | ModTypeBind       { ModTypeDec $1 }
     | ModBind           { ModDec $1 }
     | open ModExp       { OpenDec $2 (srclocOf $1) }
     | import stringlit
@@ -228,29 +228,29 @@
 
 ;
 
-SigExp :: { UncheckedSigExp }
-        : QualName            { let (v, loc) = $1 in SigVar v NoInfo (srclocOf loc) }
-        | '{' Specs '}'  { SigSpecs $2 (srcspan $1 $>) }
-        | SigExp with TypeRef { SigWith $1 $3 (srcspan $1 $>) }
-        | '(' SigExp ')'      { SigParens $2 (srcspan $1 $>) }
-        | '(' id ':' SigExp ')' '->' SigExp
-                              { let L _ (ID name) = $2
-                                in SigArrow (Just name) $4 $7 (srcspan $1 $>) }
-        | SigExp '->' SigExp  { SigArrow Nothing $1 $3 (srcspan $1 $>) }
+ModTypeExp :: { UncheckedModTypeExp }
+        : QualName                { let (v, loc) = $1 in ModTypeVar v NoInfo (srclocOf loc) }
+        | '{' Specs '}'           { ModTypeSpecs $2 (srcspan $1 $>) }
+        | ModTypeExp with TypeRef { ModTypeWith $1 $3 (srcspan $1 $>) }
+        | '(' ModTypeExp ')'      { ModTypeParens $2 (srcspan $1 $>) }
+        | '(' id ':' ModTypeExp ')' '->' ModTypeExp
+                                  { let L _ (ID name) = $2
+                                    in ModTypeArrow (Just name) $4 $7 (srcspan $1 $>) }
+        | ModTypeExp '->' ModTypeExp  { ModTypeArrow Nothing $1 $3 (srcspan $1 $>) }
 
 TypeRef :: { TypeRefBase NoInfo Name }
          : QualName TypeParams '=' TypeExpTerm
            { TypeRef (fst $1) $2 $4 (srcspan (snd $1) $>) }
 
-SigBind :: { SigBindBase NoInfo Name }
-         : module type id '=' SigExp
+ModTypeBind :: { ModTypeBindBase NoInfo Name }
+         : module type id '=' ModTypeExp
           { let L _ (ID name) = $3
-            in SigBind name $5 Nothing (srcspan $1 $>) }
+            in ModTypeBind name $5 Nothing (srcspan $1 $>) }
 
 ModExp :: { UncheckedModExp }
-        : ModExp ':' SigExp
+        : ModExp ':' ModTypeExp
           { ModAscript $1 $3 NoInfo (srcspan $1 $>) }
-        | '\\' ModParam maybeAscription(SimpleSigExp) '->' ModExp
+        | '\\' ModParam maybeAscription(SimpleModTypeExp) '->' ModExp
           { ModLambda $2 (fmap (,NoInfo) $3) $5 (srcspan $1 $>) }
         | import stringlit
           { let L _ (STRINGLIT s) = $2 in ModImport (T.unpack s) NoInfo (srcspan $1 $>) }
@@ -273,18 +273,18 @@
               { let (v, loc) = $1 in ModVar v (srclocOf loc) }
             | '{' Decs '}' { ModDecs $2 (srcspan $1 $>) }
 
-SimpleSigExp :: { UncheckedSigExp }
-             : QualName            { let (v, loc) = $1 in SigVar v NoInfo (srclocOf loc) }
-             | '(' SigExp ')'      { $2 }
+SimpleModTypeExp :: { UncheckedModTypeExp }
+             : QualName            { let (v, loc) = $1 in ModTypeVar v NoInfo (srclocOf loc) }
+             | '(' ModTypeExp ')'      { $2 }
 
 ModBind :: { ModBindBase NoInfo Name }
-         : module id ModParams maybeAscription(SigExp) '=' ModExp
+         : module id ModParams maybeAscription(ModTypeExp) '=' ModExp
            { let L floc (ID fname) = $2;
              in ModBind fname $3 (fmap (,NoInfo) $4) $6 Nothing (srcspan $1 $>)
            }
 
 ModParam :: { ModParamBase NoInfo Name }
-          : '(' id ':' SigExp ')' { let L _ (ID name) = $2 in ModParam name $4 NoInfo (srcspan $1 $>) }
+          : '(' id ':' ModTypeExp ')' { let L _ (ID name) = $2 in ModParam name $4 NoInfo (srcspan $1 $>) }
 
 ModParams :: { [ModParamBase NoInfo Name] }
            : ModParam ModParams { $1 : $2 }
@@ -308,10 +308,10 @@
         { let L _ (ID name) = $3
           in TypeSpec $2 name $4 Nothing (srcspan $1 $>) }
 
-      | module id ':' SigExp
+      | module id ':' ModTypeExp
         { let L _ (ID name) = $2
           in ModSpec name $4 Nothing (srcspan $1 $>) }
-      | include SigExp
+      | include ModTypeExp
         { IncludeSpec $2 (srcspan $1 $>) }
       | Doc Spec
         { addDocSpec $1 $2 }
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
@@ -460,7 +460,7 @@
 instance (Eq vn, IsName vn, Annot f) => Pretty (DecBase f vn) where
   pretty (ValDec dec) = pretty dec
   pretty (TypeDec dec) = pretty dec
-  pretty (SigDec sig) = pretty sig
+  pretty (ModTypeDec sig) = pretty sig
   pretty (ModDec sd) = pretty sd
   pretty (OpenDec x _) = "open" <+> pretty x
   pretty (LocalDec dec _) = "local" <+> pretty dec
@@ -544,19 +544,19 @@
   pretty (IncludeSpec e _) =
     "include" <+> pretty e
 
-instance (Eq vn, IsName vn, Annot f) => Pretty (SigExpBase f vn) where
-  pretty (SigVar v _ _) = pretty v
-  pretty (SigParens e _) = parens $ pretty e
-  pretty (SigSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map pretty ss)
-  pretty (SigWith s (TypeRef v ps td _) _) =
+instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeExpBase f vn) where
+  pretty (ModTypeVar v _ _) = pretty v
+  pretty (ModTypeParens e _) = parens $ pretty e
+  pretty (ModTypeSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map pretty ss)
+  pretty (ModTypeWith s (TypeRef v ps td _) _) =
     pretty s <+> "with" <+> pretty v <+> hsep (map pretty ps) <> " =" <+> pretty td
-  pretty (SigArrow (Just v) e1 e2 _) =
+  pretty (ModTypeArrow (Just v) e1 e2 _) =
     parens (prettyName v <> colon <+> pretty e1) <+> "->" <+> pretty e2
-  pretty (SigArrow Nothing e1 e2 _) =
+  pretty (ModTypeArrow Nothing e1 e2 _) =
     pretty e1 <+> "->" <+> pretty e2
 
-instance (Eq vn, IsName vn, Annot f) => Pretty (SigBindBase f vn) where
-  pretty (SigBind name e _ _) =
+instance (Eq vn, IsName vn, Annot f) => Pretty (ModTypeBindBase f vn) where
+  pretty (ModTypeBind name e _ _) =
     "module type" <+> prettyName name <+> equals <+> pretty e
 
 instance (Eq vn, IsName vn, Annot f) => Pretty (ModParamBase f vn) where
diff --git a/src/Language/Futhark/Primitive.hs b/src/Language/Futhark/Primitive.hs
--- a/src/Language/Futhark/Primitive.hs
+++ b/src/Language/Futhark/Primitive.hs
@@ -1285,7 +1285,7 @@
           FloatType Float32,
           \case
             [FloatValue (Float32Value x), IntValue (Int32Value y)] ->
-              Just $ FloatValue $ Float32Value $ x * (2 ** fromIntegral y)
+              Just $ FloatValue $ Float32Value $ ldexpf x $ fromIntegral y
             _ -> Nothing
         )
       ),
@@ -1294,7 +1294,7 @@
           FloatType Float64,
           \case
             [FloatValue (Float64Value x), IntValue (Int32Value y)] ->
-              Just $ FloatValue $ Float64Value $ x * (2 ** fromIntegral y)
+              Just $ FloatValue $ Float64Value $ ldexp x $ fromIntegral y
             _ -> Nothing
         )
       ),
diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs
--- a/src/Language/Futhark/Prop.hs
+++ b/src/Language/Futhark/Prop.hs
@@ -89,12 +89,12 @@
     UncheckedSlice,
     UncheckedExp,
     UncheckedModExp,
-    UncheckedSigExp,
+    UncheckedModTypeExp,
     UncheckedTypeParam,
     UncheckedPat,
     UncheckedValBind,
     UncheckedTypeBind,
-    UncheckedSigBind,
+    UncheckedModTypeBind,
     UncheckedModBind,
     UncheckedDec,
     UncheckedSpec,
@@ -110,9 +110,9 @@
     Pat,
     ModExp,
     ModParam,
-    SigExp,
+    ModTypeExp,
     ModBind,
-    SigBind,
+    ModTypeBind,
     ValBind,
     Dec,
     Spec,
@@ -1209,7 +1209,7 @@
 decImports :: DecBase f vn -> [(String, Loc)]
 decImports (OpenDec x _) = modExpImports x
 decImports (ModDec md) = modExpImports $ modExp md
-decImports SigDec {} = []
+decImports ModTypeDec {} = []
 decImports TypeDec {} = []
 decImports ValDec {} = []
 decImports (LocalDec d _) = decImports d
@@ -1236,31 +1236,31 @@
       where
         onDec OpenDec {} = mempty
         onDec ModDec {} = mempty
-        onDec (SigDec sb) =
-          M.singleton (sigName sb) (onSigExp (sigExp sb))
+        onDec (ModTypeDec sb) =
+          M.singleton (modTypeName sb) (onModTypeExp (modTypeExp sb))
         onDec TypeDec {} = mempty
         onDec ValDec {} = mempty
         onDec (LocalDec d _) = onDec d
         onDec ImportDec {} = mempty
 
-        onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v
-        onSigExp (SigParens e _) = onSigExp e
-        onSigExp (SigSpecs ss _) = foldMap onSpec ss
-        onSigExp (SigWith e _ _) = onSigExp e
-        onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2
+        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v
+        onModTypeExp (ModTypeParens e _) = onModTypeExp e
+        onModTypeExp (ModTypeSpecs ss _) = foldMap onSpec ss
+        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e
+        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp e2
 
         onSpec ValSpec {} = mempty
         onSpec TypeSpec {} = mempty
         onSpec TypeAbbrSpec {} = mempty
-        onSpec (ModSpec vn e _ _) = S.singleton vn <> onSigExp e
-        onSpec (IncludeSpec e _) = onSigExp e
+        onSpec (ModSpec vn e _ _) = S.singleton vn <> onModTypeExp e
+        onSpec (IncludeSpec e _) = onModTypeExp e
 
     mtypes_used = foldMap onDec $ progDecs prog
       where
         onDec (OpenDec x _) = onModExp x
         onDec (ModDec md) =
-          maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md)
-        onDec SigDec {} = mempty
+          maybe mempty (onModTypeExp . fst) (modType md) <> onModExp (modExp md)
+        onDec ModTypeDec {} = mempty
         onDec TypeDec {} = mempty
         onDec ValDec {} = mempty
         onDec LocalDec {} = mempty
@@ -1271,17 +1271,17 @@
         onModExp ModImport {} = mempty
         onModExp (ModDecs ds _) = mconcat $ map onDec ds
         onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2
-        onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se
+        onModExp (ModAscript me se _ _) = onModExp me <> onModTypeExp se
         onModExp (ModLambda p r me _) =
-          onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me
+          onModParam p <> maybe mempty (onModTypeExp . fst) r <> onModExp me
 
-        onModParam = onSigExp . modParamType
+        onModParam = onModTypeExp . modParamType
 
-        onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v
-        onSigExp (SigParens e _) = onSigExp e
-        onSigExp SigSpecs {} = mempty
-        onSigExp (SigWith e _ _) = onSigExp e
-        onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2
+        onModTypeExp (ModTypeVar v _ _) = S.singleton $ qualLeaf v
+        onModTypeExp (ModTypeParens e _) = onModTypeExp e
+        onModTypeExp ModTypeSpecs {} = mempty
+        onModTypeExp (ModTypeWith e _ _) = onModTypeExp e
+        onModTypeExp (ModTypeArrow _ e1 e2 _) = onModTypeExp e1 <> onModTypeExp e2
 
 -- | Extract a leading @((name, namespace, file), remainder)@ from a
 -- documentation comment string.  These are formatted as
@@ -1321,7 +1321,7 @@
     holesInDec (OpenDec me _) = holesInModExp me
     holesInDec (LocalDec d _) = holesInDec d
     holesInDec TypeDec {} = mempty
-    holesInDec SigDec {} = mempty
+    holesInDec ModTypeDec {} = mempty
     holesInDec ImportDec {} = mempty
 
     holesInModExp (ModDecs ds _) = foldMap holesInDec ds
@@ -1455,7 +1455,7 @@
 type ModBind = ModBindBase Info VName
 
 -- | A type-checked module type binding.
-type SigBind = SigBindBase Info VName
+type ModTypeBind = ModTypeBindBase Info VName
 
 -- | A type-checked module expression.
 type ModExp = ModExpBase Info VName
@@ -1464,7 +1464,7 @@
 type ModParam = ModParamBase Info VName
 
 -- | A type-checked module type expression.
-type SigExp = SigExpBase Info VName
+type ModTypeExp = ModTypeExpBase Info VName
 
 -- | A type-checked declaration.
 type Dec = DecBase Info VName
@@ -1509,7 +1509,7 @@
 type UncheckedModExp = ModExpBase NoInfo Name
 
 -- | A module type expression with no type annotations.
-type UncheckedSigExp = SigExpBase NoInfo Name
+type UncheckedModTypeExp = ModTypeExpBase NoInfo Name
 
 -- | A type parameter with no type annotations.
 type UncheckedTypeParam = TypeParamBase Name
@@ -1524,7 +1524,7 @@
 type UncheckedTypeBind = TypeBindBase NoInfo Name
 
 -- | A module type binding with no type annotations.
-type UncheckedSigBind = SigBindBase NoInfo Name
+type UncheckedModTypeBind = ModTypeBindBase NoInfo Name
 
 -- | A module binding with no type annotations.
 type UncheckedModBind = ModBindBase NoInfo Name
diff --git a/src/Language/Futhark/Query.hs b/src/Language/Futhark/Query.hs
--- a/src/Language/Futhark/Query.hs
+++ b/src/Language/Futhark/Query.hs
@@ -120,7 +120,7 @@
 modParamDefs :: ModParam -> Defs
 modParamDefs (ModParam p se _ loc) =
   M.singleton p (DefBound $ BoundModule $ locOf loc)
-    <> sigExpDefs se
+    <> modTypeExpDefs se
 
 modExpDefs :: ModExp -> Defs
 modExpDefs ModVar {} =
@@ -143,7 +143,7 @@
   M.singleton (modName mbind) (DefBound $ BoundModule $ locOf mbind)
     <> mconcat (map modParamDefs (modParams mbind))
     <> modExpDefs (modExp mbind)
-    <> case modSignature mbind of
+    <> case modType mbind of
       Nothing -> mempty
       Just (_, Info substs) ->
         M.map DefIndirect substs
@@ -159,28 +159,28 @@
       M.singleton v $ DefBound $ BoundType $ locOf loc
     ModSpec v se _ loc ->
       M.singleton v (DefBound $ BoundModuleType $ locOf loc)
-        <> sigExpDefs se
-    IncludeSpec se _ -> sigExpDefs se
+        <> modTypeExpDefs se
+    IncludeSpec se _ -> modTypeExpDefs se
 
-sigExpDefs :: SigExp -> Defs
-sigExpDefs se =
+modTypeExpDefs :: ModTypeExp -> Defs
+modTypeExpDefs se =
   case se of
-    SigVar _ (Info substs) _ -> M.map DefIndirect substs
-    SigParens e _ -> sigExpDefs e
-    SigSpecs specs _ -> mconcat $ map specDefs specs
-    SigWith e _ _ -> sigExpDefs e
-    SigArrow _ e1 e2 _ -> sigExpDefs e1 <> sigExpDefs e2
+    ModTypeVar _ (Info substs) _ -> M.map DefIndirect substs
+    ModTypeParens e _ -> modTypeExpDefs e
+    ModTypeSpecs specs _ -> mconcat $ map specDefs specs
+    ModTypeWith e _ _ -> modTypeExpDefs e
+    ModTypeArrow _ e1 e2 _ -> modTypeExpDefs e1 <> modTypeExpDefs e2
 
-sigBindDefs :: SigBind -> Defs
+sigBindDefs :: ModTypeBind -> Defs
 sigBindDefs sbind =
-  M.singleton (sigName sbind) (DefBound $ BoundModuleType $ locOf sbind)
-    <> sigExpDefs (sigExp sbind)
+  M.singleton (modTypeName sbind) (DefBound $ BoundModuleType $ locOf sbind)
+    <> modTypeExpDefs (modTypeExp sbind)
 
 decDefs :: Dec -> Defs
 decDefs (ValDec vbind) = valBindDefs vbind
 decDefs (TypeDec vbind) = typeBindDefs vbind
 decDefs (ModDec mbind) = modBindDefs mbind
-decDefs (SigDec mbind) = sigBindDefs mbind
+decDefs (ModTypeDec mbind) = sigBindDefs mbind
 decDefs (OpenDec me _) = modExpDefs me
 decDefs (LocalDec dec _) = decDefs dec
 decDefs ImportDec {} = mempty
@@ -312,19 +312,19 @@
     ValSpec _ _ te _ _ _ -> atPosInTypeExp te pos
     TypeAbbrSpec tbind -> atPosInTypeBind tbind pos
     TypeSpec {} -> Nothing
-    ModSpec _ se _ _ -> atPosInSigExp se pos
-    IncludeSpec se _ -> atPosInSigExp se pos
+    ModSpec _ se _ _ -> atPosInModTypeExp se pos
+    IncludeSpec se _ -> atPosInModTypeExp se pos
 
-atPosInSigExp :: SigExp -> Pos -> Maybe RawAtPos
-atPosInSigExp se pos =
+atPosInModTypeExp :: ModTypeExp -> Pos -> Maybe RawAtPos
+atPosInModTypeExp se pos =
   case se of
-    SigVar qn _ loc -> do
+    ModTypeVar qn _ loc -> do
       guard $ loc `contains` pos
       Just $ RawAtName qn $ locOf loc
-    SigParens e _ -> atPosInSigExp e pos
-    SigSpecs specs _ -> msum $ map (`atPosInSpec` pos) specs
-    SigWith e _ _ -> atPosInSigExp e pos
-    SigArrow _ e1 e2 _ -> atPosInSigExp e1 pos `mplus` atPosInSigExp e2 pos
+    ModTypeParens e _ -> atPosInModTypeExp e pos
+    ModTypeSpecs specs _ -> msum $ map (`atPosInSpec` pos) specs
+    ModTypeWith e _ _ -> atPosInModTypeExp e pos
+    ModTypeArrow _ e1 e2 _ -> atPosInModTypeExp e1 pos `mplus` atPosInModTypeExp e2 pos
 
 atPosInValBind :: ValBind -> Pos -> Maybe RawAtPos
 atPosInValBind vbind pos =
@@ -341,12 +341,12 @@
     `mplus` atPosInModExp e pos
     `mplus` case sig of
       Nothing -> Nothing
-      Just (se, _) -> atPosInSigExp se pos
+      Just (se, _) -> atPosInModTypeExp se pos
   where
-    inParam (ModParam _ se _ _) = atPosInSigExp se pos
+    inParam (ModParam _ se _ _) = atPosInModTypeExp se pos
 
-atPosInSigBind :: SigBind -> Pos -> Maybe RawAtPos
-atPosInSigBind = atPosInSigExp . sigExp
+atPosInModTypeBind :: ModTypeBind -> Pos -> Maybe RawAtPos
+atPosInModTypeBind = atPosInModTypeExp . modTypeExp
 
 atPosInDec :: Dec -> Pos -> Maybe RawAtPos
 atPosInDec dec pos = do
@@ -355,7 +355,7 @@
     ValDec vbind -> atPosInValBind vbind pos
     TypeDec tbind -> atPosInTypeBind tbind pos
     ModDec mbind -> atPosInModBind mbind pos
-    SigDec sbind -> atPosInSigBind sbind pos
+    ModTypeDec sbind -> atPosInModTypeBind sbind pos
     OpenDec e _ -> atPosInModExp e pos
     LocalDec dec' _ -> atPosInDec dec' pos
     ImportDec {} -> Nothing
diff --git a/src/Language/Futhark/Semantic.hs b/src/Language/Futhark/Semantic.hs
--- a/src/Language/Futhark/Semantic.hs
+++ b/src/Language/Futhark/Semantic.hs
@@ -12,7 +12,7 @@
     Namespace (..),
     Env (..),
     TySet,
-    FunSig (..),
+    FunModType (..),
     NameMap,
     BoundV (..),
     Mod (..),
@@ -96,15 +96,15 @@
 -- or a parametric module ("functor" in SML).
 data Mod
   = ModEnv Env
-  | ModFun FunSig
+  | ModFun FunModType
   deriving (Show)
 
 -- | A parametric functor consists of a set of abstract types, the
 -- environment of its parameter, and the resulting module type.
-data FunSig = FunSig
-  { funSigAbs :: TySet,
-    funSigMod :: Mod,
-    funSigMty :: MTy
+data FunModType = FunModType
+  { funModTypeAbs :: TySet,
+    funModTypeMod :: Mod,
+    funModTypeMty :: MTy
   }
   deriving (Show)
 
@@ -137,7 +137,7 @@
 data Env = Env
   { envVtable :: M.Map VName BoundV,
     envTypeTable :: M.Map VName TypeBinding,
-    envSigTable :: M.Map VName MTy,
+    envModTypeTable :: M.Map VName MTy,
     envModTable :: M.Map VName Mod,
     envNameMap :: NameMap
   }
@@ -160,7 +160,7 @@
 
 instance Pretty Mod where
   pretty (ModEnv e) = pretty e
-  pretty (ModFun (FunSig _ mod mty)) = pretty mod <+> "->" </> pretty mty
+  pretty (ModFun (FunModType _ mod mty)) = pretty mod <+> "->" </> pretty mty
 
 instance Pretty Env where
   pretty (Env vtable ttable sigtable modtable _) =
diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs
--- a/src/Language/Futhark/Syntax.hs
+++ b/src/Language/Futhark/Syntax.hs
@@ -66,9 +66,9 @@
     -- * Module language
     ImportName (..),
     SpecBase (..),
-    SigExpBase (..),
+    ModTypeExpBase (..),
     TypeRefBase (..),
-    SigBindBase (..),
+    ModTypeBindBase (..),
     ModExpBase (..),
     ModBindBase (..),
     ModParamBase (..),
@@ -1123,8 +1123,8 @@
   | TypeAbbrSpec (TypeBindBase f vn)
   | -- | Abstract type.
     TypeSpec Liftedness vn [TypeParamBase vn] (Maybe DocComment) SrcLoc
-  | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc
-  | IncludeSpec (SigExpBase f vn) SrcLoc
+  | ModSpec vn (ModTypeExpBase f vn) (Maybe DocComment) SrcLoc
+  | IncludeSpec (ModTypeExpBase f vn) SrcLoc
 
 deriving instance Show (SpecBase Info VName)
 
@@ -1138,16 +1138,16 @@
   locOf (IncludeSpec _ loc) = locOf loc
 
 -- | A module type expression.
-data SigExpBase f vn
-  = SigVar (QualName vn) (f (M.Map VName VName)) SrcLoc
-  | SigParens (SigExpBase f vn) SrcLoc
-  | SigSpecs [SpecBase f vn] SrcLoc
-  | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc
-  | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc
+data ModTypeExpBase f vn
+  = ModTypeVar (QualName vn) (f (M.Map VName VName)) SrcLoc
+  | ModTypeParens (ModTypeExpBase f vn) SrcLoc
+  | ModTypeSpecs [SpecBase f vn] SrcLoc
+  | ModTypeWith (ModTypeExpBase f vn) (TypeRefBase f vn) SrcLoc
+  | ModTypeArrow (Maybe vn) (ModTypeExpBase f vn) (ModTypeExpBase f vn) SrcLoc
 
-deriving instance Show (SigExpBase Info VName)
+deriving instance Show (ModTypeExpBase Info VName)
 
-deriving instance Show (SigExpBase NoInfo Name)
+deriving instance Show (ModTypeExpBase NoInfo Name)
 
 -- | A type refinement.
 data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeExp f vn) SrcLoc
@@ -1159,27 +1159,27 @@
 instance Located (TypeRefBase f vn) where
   locOf (TypeRef _ _ _ loc) = locOf loc
 
-instance Located (SigExpBase f vn) where
-  locOf (SigVar _ _ loc) = locOf loc
-  locOf (SigParens _ loc) = locOf loc
-  locOf (SigSpecs _ loc) = locOf loc
-  locOf (SigWith _ _ loc) = locOf loc
-  locOf (SigArrow _ _ _ loc) = locOf loc
+instance Located (ModTypeExpBase f vn) where
+  locOf (ModTypeVar _ _ loc) = locOf loc
+  locOf (ModTypeParens _ loc) = locOf loc
+  locOf (ModTypeSpecs _ loc) = locOf loc
+  locOf (ModTypeWith _ _ loc) = locOf loc
+  locOf (ModTypeArrow _ _ _ loc) = locOf loc
 
 -- | Module type binding.
-data SigBindBase f vn = SigBind
-  { sigName :: vn,
-    sigExp :: SigExpBase f vn,
-    sigDoc :: Maybe DocComment,
-    sigLoc :: SrcLoc
+data ModTypeBindBase f vn = ModTypeBind
+  { modTypeName :: vn,
+    modTypeExp :: ModTypeExpBase f vn,
+    modTypeDoc :: Maybe DocComment,
+    modTypeLoc :: SrcLoc
   }
 
-deriving instance Show (SigBindBase Info VName)
+deriving instance Show (ModTypeBindBase Info VName)
 
-deriving instance Show (SigBindBase NoInfo Name)
+deriving instance Show (ModTypeBindBase NoInfo Name)
 
-instance Located (SigBindBase f vn) where
-  locOf = locOf . sigLoc
+instance Located (ModTypeBindBase f vn) where
+  locOf = locOf . modTypeLoc
 
 -- | Canonical reference to a Futhark code file.  Does not include the
 -- @.fut@ extension.  This is most often a path relative to the
@@ -1205,10 +1205,10 @@
       (f (M.Map VName VName))
       (f (M.Map VName VName))
       SrcLoc
-  | ModAscript (ModExpBase f vn) (SigExpBase f vn) (f (M.Map VName VName)) SrcLoc
+  | ModAscript (ModExpBase f vn) (ModTypeExpBase f vn) (f (M.Map VName VName)) SrcLoc
   | ModLambda
       (ModParamBase f vn)
-      (Maybe (SigExpBase f vn, f (M.Map VName VName)))
+      (Maybe (ModTypeExpBase f vn, f (M.Map VName VName)))
       (ModExpBase f vn)
       SrcLoc
 
@@ -1229,7 +1229,7 @@
 data ModBindBase f vn = ModBind
   { modName :: vn,
     modParams :: [ModParamBase f vn],
-    modSignature :: Maybe (SigExpBase f vn, f (M.Map VName VName)),
+    modType :: Maybe (ModTypeExpBase f vn, f (M.Map VName VName)),
     modExp :: ModExpBase f vn,
     modDoc :: Maybe DocComment,
     modLocation :: SrcLoc
@@ -1245,7 +1245,7 @@
 -- | A module parameter.
 data ModParamBase f vn = ModParam
   { modParamName :: vn,
-    modParamType :: SigExpBase f vn,
+    modParamType :: ModTypeExpBase f vn,
     modParamAbs :: f [VName],
     modParamLocation :: SrcLoc
   }
@@ -1261,7 +1261,7 @@
 data DecBase f vn
   = ValDec (ValBindBase f vn)
   | TypeDec (TypeBindBase f vn)
-  | SigDec (SigBindBase f vn)
+  | ModTypeDec (ModTypeBindBase f vn)
   | ModDec (ModBindBase f vn)
   | OpenDec (ModExpBase f vn) SrcLoc
   | LocalDec (DecBase f vn) SrcLoc
@@ -1274,7 +1274,7 @@
 instance Located (DecBase f vn) where
   locOf (ValDec d) = locOf d
   locOf (TypeDec d) = locOf d
-  locOf (SigDec d) = locOf d
+  locOf (ModTypeDec d) = locOf d
   locOf (ModDec d) = locOf d
   locOf (OpenDec _ loc) = locOf loc
   locOf (LocalDec _ loc) = locOf loc
diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs
--- a/src/Language/Futhark/TypeChecker.hs
+++ b/src/Language/Futhark/TypeChecker.hs
@@ -178,7 +178,7 @@
       check Term (valBindName vb) (srclocOf vb)
     f (TypeDec (TypeBind name _ _ _ _ _ loc)) =
       check Type name loc
-    f (SigDec (SigBind name _ _ loc)) =
+    f (ModTypeDec (ModTypeBind name _ _ loc)) =
       check Signature name loc
     f (ModDec (ModBind name _ _ _ _ loc)) =
       check Term name loc
@@ -279,7 +279,7 @@
 checkSpecs (ModSpec name sig doc loc : specs) =
   bindSpaced [(Term, name)] $ do
     name' <- checkName Term name loc
-    (_sig_abs, mty, sig') <- checkSigExp sig
+    (_sig_abs, mty, sig') <- checkModTypeExp sig
     let senv =
           mempty
             { envNameMap = M.singleton (Term, name) $ qualName name',
@@ -292,7 +292,7 @@
         ModSpec name' sig' doc loc : specs'
       )
 checkSpecs (IncludeSpec e loc : specs) = do
-  (e_abs, env_abs, e_env, e') <- checkSigExpToEnv e
+  (e_abs, env_abs, e_env, e') <- checkModTypeExpToEnv e
 
   mapM_ (warnIfShadowing . fmap baseName) $ M.keys env_abs
 
@@ -309,28 +309,28 @@
     warnAbout qn =
       warn loc $ "Inclusion shadows type" <+> dquotes (pretty qn) <+> "."
 
-checkSigExp :: SigExpBase NoInfo Name -> TypeM (TySet, MTy, SigExpBase Info VName)
-checkSigExp (SigParens e loc) = do
-  (abs, mty, e') <- checkSigExp e
-  pure (abs, mty, SigParens e' loc)
-checkSigExp (SigVar name NoInfo loc) = do
+checkModTypeExp :: ModTypeExpBase NoInfo Name -> TypeM (TySet, MTy, ModTypeExpBase Info VName)
+checkModTypeExp (ModTypeParens e loc) = do
+  (abs, mty, e') <- checkModTypeExp e
+  pure (abs, mty, ModTypeParens e' loc)
+checkModTypeExp (ModTypeVar name NoInfo loc) = do
   (name', mty) <- lookupMTy loc name
   (mty', substs) <- newNamesForMTy mty
-  pure (mtyAbs mty', mty', SigVar name' (Info substs) loc)
-checkSigExp (SigSpecs specs loc) = do
+  pure (mtyAbs mty', mty', ModTypeVar name' (Info substs) loc)
+checkModTypeExp (ModTypeSpecs specs loc) = do
   checkForDuplicateSpecs specs
   (abstypes, env, specs') <- checkSpecs specs
-  pure (abstypes, MTy abstypes $ ModEnv env, SigSpecs specs' loc)
-checkSigExp (SigWith s (TypeRef tname ps te trloc) loc) = do
-  (abs, s_abs, s_env, s') <- checkSigExpToEnv s
+  pure (abstypes, MTy abstypes $ ModEnv env, ModTypeSpecs specs' loc)
+checkModTypeExp (ModTypeWith s (TypeRef tname ps te trloc) loc) = do
+  (abs, s_abs, s_env, s') <- checkModTypeExpToEnv s
   checkTypeParams ps $ \ps' -> do
     (ext, te', te_t, _) <- bindingTypeParams ps' $ checkTypeDecl te
     unless (null ext) $
       typeError te' mempty "Anonymous dimensions are not allowed here."
     (tname', s_abs', s_env') <- refineEnv loc s_abs s_env tname ps' te_t
-    pure (abs, MTy s_abs' $ ModEnv s_env', SigWith s' (TypeRef tname' ps' te' trloc) loc)
-checkSigExp (SigArrow maybe_pname e1 e2 loc) = do
-  (e1_abs, MTy s_abs e1_mod, e1') <- checkSigExp e1
+    pure (abs, MTy s_abs' $ ModEnv s_env', ModTypeWith s' (TypeRef tname' ps' te' trloc) loc)
+checkModTypeExp (ModTypeArrow maybe_pname e1 e2 loc) = do
+  (e1_abs, MTy s_abs e1_mod, e1') <- checkModTypeExp e1
   (env_for_e2, maybe_pname') <-
     case maybe_pname of
       Just pname -> bindSpaced [(Term, pname)] $ do
@@ -344,34 +344,34 @@
           )
       Nothing ->
         pure (mempty, Nothing)
-  (e2_abs, e2_mod, e2') <- localEnv env_for_e2 $ checkSigExp e2
+  (e2_abs, e2_mod, e2') <- localEnv env_for_e2 $ checkModTypeExp e2
   pure
     ( e1_abs <> e2_abs,
-      MTy mempty $ ModFun $ FunSig s_abs e1_mod e2_mod,
-      SigArrow maybe_pname' e1' e2' loc
+      MTy mempty $ ModFun $ FunModType s_abs e1_mod e2_mod,
+      ModTypeArrow maybe_pname' e1' e2' loc
     )
 
-checkSigExpToEnv ::
-  SigExpBase NoInfo Name ->
-  TypeM (TySet, TySet, Env, SigExpBase Info VName)
-checkSigExpToEnv e = do
-  (abs, MTy mod_abs mod, e') <- checkSigExp e
+checkModTypeExpToEnv ::
+  ModTypeExpBase NoInfo Name ->
+  TypeM (TySet, TySet, Env, ModTypeExpBase Info VName)
+checkModTypeExpToEnv e = do
+  (abs, MTy mod_abs mod, e') <- checkModTypeExp e
   case mod of
     ModEnv env -> pure (abs, mod_abs, env, e')
     ModFun {} -> unappliedFunctor $ srclocOf e
 
-checkSigBind :: SigBindBase NoInfo Name -> TypeM (TySet, Env, SigBindBase Info VName)
-checkSigBind (SigBind name e doc loc) = do
-  (abs, env, e') <- checkSigExp e
+checkModTypeBind :: ModTypeBindBase NoInfo Name -> TypeM (TySet, Env, ModTypeBindBase Info VName)
+checkModTypeBind (ModTypeBind name e doc loc) = do
+  (abs, env, e') <- checkModTypeExp e
   bindSpaced [(Signature, name)] $ do
     name' <- checkName Signature name loc
     pure
       ( abs,
         mempty
-          { envSigTable = M.singleton name' env,
+          { envModTypeTable = M.singleton name' env,
             envNameMap = M.singleton (Signature, name) (qualName name')
           },
-        SigBind name' e' doc loc
+        ModTypeBind name' e' doc loc
       )
 
 checkOneModExp ::
@@ -418,7 +418,7 @@
       typeError loc mempty "Cannot apply non-parametric module."
 checkOneModExp (ModAscript me se NoInfo loc) = do
   (me_abs, me_mod, me') <- checkOneModExp me
-  (se_abs, se_mty, se') <- checkSigExp se
+  (se_abs, se_mty, se') <- checkModTypeExp se
   match_subst <- badOnLeft $ matchMTys me_mod se_mty (locOf loc)
   pure (se_abs <> me_abs, se_mty, ModAscript me' se' (Info match_subst) loc)
 checkOneModExp (ModLambda param maybe_fsig_e body_e loc) =
@@ -427,7 +427,7 @@
       checkModBody (fst <$> maybe_fsig_e) body_e loc
     pure
       ( abs,
-        MTy mempty $ ModFun $ FunSig param_abs param_mod mty,
+        MTy mempty $ ModFun $ FunModType param_abs param_mod mty,
         ModLambda param' maybe_fsig_e' body_e' loc
       )
 
@@ -443,7 +443,7 @@
   (ModParamBase Info VName -> TySet -> Mod -> TypeM a) ->
   TypeM a
 withModParam (ModParam pname psig_e NoInfo loc) m = do
-  (_abs, MTy p_abs p_mod, psig_e') <- checkSigExp psig_e
+  (_abs, MTy p_abs p_mod, psig_e') <- checkModTypeExp psig_e
   bindSpaced [(Term, pname)] $ do
     pname' <- checkName Term pname loc
     let in_body_env = mempty {envModTable = M.singleton pname' p_mod}
@@ -460,12 +460,12 @@
     withModParams ps $ \ps' -> m $ (p', pabs, pmod) : ps'
 
 checkModBody ::
-  Maybe (SigExpBase NoInfo Name) ->
+  Maybe (ModTypeExpBase NoInfo Name) ->
   ModExpBase NoInfo Name ->
   SrcLoc ->
   TypeM
     ( TySet,
-      Maybe (SigExp, Info (M.Map VName VName)),
+      Maybe (ModTypeExp, Info (M.Map VName VName)),
       ModExp,
       MTy
     )
@@ -480,7 +480,7 @@
           body_mty
         )
     Just fsig_e -> do
-      (fsig_abs, fsig_mty, fsig_e') <- checkSigExp fsig_e
+      (fsig_abs, fsig_mty, fsig_e') <- checkModTypeExp fsig_e
       fsig_subst <- badOnLeft $ matchMTys body_mty fsig_mty (locOf loc)
       pure
         ( fsig_abs <> body_e_abs,
@@ -508,13 +508,13 @@
       withModParams ps $ \params_stuff -> do
         let (ps', ps_abs, ps_mod) = unzip3 params_stuff
         (abs, maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc
-        let addParam (x, y) mty' = MTy mempty $ ModFun $ FunSig x y mty'
+        let addParam (x, y) mty' = MTy mempty $ ModFun $ FunModType x y mty'
         pure
           ( abs,
             p' : ps',
             maybe_fsig_e',
             body_e',
-            FunSig p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod
+            FunModType p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod
           )
   bindSpaced [(Term, name)] $ do
     name' <- checkName Term name loc
@@ -755,9 +755,9 @@
 checkOneDec (ModDec struct) = do
   (abs, modenv, struct') <- checkModBind struct
   pure (abs, modenv, ModDec struct')
-checkOneDec (SigDec sig) = do
-  (abs, sigenv, sig') <- checkSigBind sig
-  pure (abs, sigenv, SigDec sig')
+checkOneDec (ModTypeDec sig) = do
+  (abs, sigenv, sig') <- checkModTypeBind sig
+  pure (abs, sigenv, ModTypeDec sig')
 checkOneDec (TypeDec tdec) = do
   (tenv, tdec') <- checkTypeBind tdec
   pure (mempty, tenv, TypeDec tdec')
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
@@ -25,8 +25,8 @@
 substituteTypesInMod :: TypeSubs -> Mod -> Mod
 substituteTypesInMod substs (ModEnv e) =
   ModEnv $ substituteTypesInEnv substs e
-substituteTypesInMod substs (ModFun (FunSig abs mod mty)) =
-  ModFun $ FunSig abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty)
+substituteTypesInMod substs (ModFun (FunModType abs mod mty)) =
+  ModFun $ FunModType abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty)
 
 substituteTypesInMTy :: TypeSubs -> MTy -> MTy
 substituteTypesInMTy substs (MTy abs mod) = MTy abs $ substituteTypesInMod substs mod
@@ -97,7 +97,7 @@
            in Env
                 { envVtable = vtable',
                   envTypeTable = ttable',
-                  envSigTable = mempty,
+                  envModTypeTable = mempty,
                   envModTable = mtable',
                   envNameMap = M.map (fmap substitute) names
                 }
@@ -122,10 +122,10 @@
         substituteInMod (ModEnv env) =
           ModEnv $ substituteInEnv env
         substituteInMod (ModFun funsig) =
-          ModFun $ substituteInFunSig funsig
+          ModFun $ substituteInFunModType funsig
 
-        substituteInFunSig (FunSig abs mod mty) =
-          FunSig
+        substituteInFunModType (FunModType abs mod mty) =
+          FunModType
             (M.mapKeys (fmap substitute) abs)
             (substituteInMod mod)
             (substituteInMTy substs mty)
@@ -169,7 +169,7 @@
 modTypeAbbrs :: Mod -> M.Map VName TypeBinding
 modTypeAbbrs (ModEnv env) =
   envTypeAbbrs env
-modTypeAbbrs (ModFun (FunSig _ mod mty)) =
+modTypeAbbrs (ModFun (FunModType _ mod mty)) =
   modTypeAbbrs mod <> mtyTypeAbbrs mty
 
 envTypeAbbrs :: Env -> M.Map VName TypeBinding
@@ -309,8 +309,8 @@
     resolveModNames (ModEnv mod_env) (ModEnv sig_env) =
       resolveEnvNames mod_env sig_env
     resolveModNames (ModFun mod_fun) (ModFun sig_fun) =
-      resolveModNames (funSigMod mod_fun) (funSigMod sig_fun)
-        <> resolveMTyNames' (funSigMty mod_fun) (funSigMty sig_fun)
+      resolveModNames (funModTypeMod mod_fun) (funModTypeMod sig_fun)
+        <> resolveMTyNames' (funModTypeMty mod_fun) (funModTypeMty sig_fun)
     resolveModNames _ _ =
       mempty
 
@@ -453,8 +453,8 @@
     matchMods
       old_abs_subst_to_type
       quals
-      (ModFun (FunSig mod_abs mod_pmod mod_mod))
-      (ModFun (FunSig sig_abs sig_pmod sig_mod))
+      (ModFun (FunModType mod_abs mod_pmod mod_mod))
+      (ModFun (FunModType sig_abs sig_pmod sig_mod))
       loc = do
         -- We need to use different substitutions when matching
         -- parameter and body signatures - this is because the
@@ -602,14 +602,14 @@
 -- | Apply a parametric module to an argument.
 applyFunctor ::
   Loc ->
-  FunSig ->
+  FunModType ->
   MTy ->
   TypeM
     ( MTy,
       M.Map VName VName,
       M.Map VName VName
     )
-applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do
+applyFunctor applyloc (FunModType p_abs p_mod body_mty) a_mty = do
   p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc
 
   -- Apply type abbreviations from a_mty to body_mty.
diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs
--- a/src/Language/Futhark/TypeChecker/Monad.hs
+++ b/src/Language/Futhark/TypeChecker/Monad.hs
@@ -30,7 +30,7 @@
     module Language.Futhark.Warnings,
     Env (..),
     TySet,
-    FunSig (..),
+    FunModType (..),
     ImportTable,
     NameMap,
     BoundV (..),
@@ -240,7 +240,7 @@
 lookupMTy :: SrcLoc -> QualName Name -> TypeM (QualName VName, MTy)
 lookupMTy loc qn = do
   (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Signature qn loc
-  (qn',) <$> maybe explode pure (M.lookup name $ envSigTable scope)
+  (qn',) <$> maybe explode pure (M.lookup name $ envModTypeTable scope)
   where
     explode = unknownVariable Signature qn loc
 
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
@@ -631,7 +631,7 @@
                 let f x = if x == v then Just (ExpSubst e') else Nothing
                 pure (applySubst f body_t, [])
           _ ->
-            unscopeType loc (patNames pat') body_t
+            unscopeType loc (map sizeName sizes' <> patNames pat') body_t
 
       pure $
         AppExp
@@ -1695,7 +1695,7 @@
     let keep_type_vars = overloadedTypeVars now_substs
 
     cur_lvl <- curLevel
-    let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= cur_lvl
+    let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= (cur_lvl - length params)
         new_substs = M.filterWithKey candidate now_substs
 
     (tparams', RetType ret_dims restype') <-
diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
--- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs
@@ -343,7 +343,7 @@
   checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do
     let descend ps' (p : ps) =
           checkPat [] p NoneInferred $ \p' ->
-            binding (patIdents $ fmap toStruct p') $ descend (p' : ps') ps
+            binding (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps
         descend ps' [] = m tps' $ reverse ps'
 
     descend [] orig_ps
diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs
--- a/src/Language/Futhark/TypeChecker/Types.hs
+++ b/src/Language/Futhark/TypeChecker/Types.hs
@@ -547,11 +547,11 @@
           }
 
 applyType ::
-  (Monoid als) =>
+  (Monoid u) =>
   [TypeParam] ->
-  TypeBase Size als ->
+  TypeBase Size u ->
   [StructTypeArg] ->
-  TypeBase Size als
+  TypeBase Size u
 applyType ps t args = substTypesAny (`M.lookup` substs) t
   where
     substs = M.fromList $ zipWith mkSubst ps args
@@ -564,10 +564,10 @@
       error $ "applyType mkSubst: cannot substitute " ++ prettyString a ++ " for " ++ prettyString p
 
 substTypesRet ::
-  (Monoid as) =>
-  (VName -> Maybe (Subst (RetTypeBase Size as))) ->
-  TypeBase Size as ->
-  RetTypeBase Size as
+  (Monoid u) =>
+  (VName -> Maybe (Subst (RetTypeBase Size u))) ->
+  TypeBase Size u ->
+  RetTypeBase Size u
 substTypesRet lookupSubst ot =
   uncurry (flip RetType) $ runState (onType ot) []
   where
@@ -614,8 +614,8 @@
           pure $ Scalar $ TypeVar u v targs'
     onType (Scalar (Record ts)) =
       Scalar . Record <$> traverse onType ts
-    onType (Scalar (Arrow als v d t1 t2)) =
-      Scalar <$> (Arrow als v d <$> onType t1 <*> onRetType t2)
+    onType (Scalar (Arrow u v d t1 t2)) =
+      Scalar <$> (Arrow u v d <$> onType t1 <*> onRetType t2)
     onType (Scalar (Sum ts)) =
       Scalar . Sum <$> traverse (traverse onType) ts
 
@@ -642,10 +642,10 @@
 -- | Perform substitutions, from type names to types, on a type. Works
 -- regardless of what shape and uniqueness information is attached to the type.
 substTypesAny ::
-  (Monoid as) =>
-  (VName -> Maybe (Subst (RetTypeBase Size as))) ->
-  TypeBase Size as ->
-  TypeBase Size as
+  (Monoid u) =>
+  (VName -> Maybe (Subst (RetTypeBase Size u))) ->
+  TypeBase Size u ->
+  TypeBase Size u
 substTypesAny lookupSubst ot =
   case substTypesRet lookupSubst ot of
     RetType [] ot' -> ot'