diff --git a/docs/man/futhark-repl.rst b/docs/man/futhark-repl.rst
--- a/docs/man/futhark-repl.rst
+++ b/docs/man/futhark-repl.rst
@@ -17,7 +17,8 @@
 Start an interactive Futhark session.  This will let you interactively
 enter expressions and declarations which are then immediately
 interpreted.  If the entered line can be either a declaration or an
-expression, it is assumed to be a declaration.
+expression, it is assumed to be a declaration.  The input msut fit on
+a single line.
 
 Futhark source files can be loaded using the ``:load`` command.  This
 will erase any interactively entered definitions.  Use the ``:help``
diff --git a/docs/usage.rst b/docs/usage.rst
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -482,9 +482,9 @@
 When using the OpenCL backend, extra API functions are provided for
 directly accessing or providing the OpenCL objects used by Futhark.
 Take care when using these functions.  In particular, a Futhark
-context can now be provided with the command queue to use::
+context can now be configured with the command queue to use::
 
-  struct futhark_context *futhark_context_new_with_command_queue(struct futhark_context_config *cfg, cl_command_queue queue);
+  void futhark_context_config_set_command_queue(struct futhark_context_config *cfg, cl_command_queue queue);
 
 As a ``cl_command_queue`` specifies an OpenCL device, this is also how
 manual platform and device selection is possible.  A function is also
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.1
+version:        0.25.2
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -248,6 +248,7 @@
       Futhark.CodeGen.ImpGen.CUDA
       Futhark.CodeGen.ImpGen.GPU
       Futhark.CodeGen.ImpGen.GPU.Base
+      Futhark.CodeGen.ImpGen.GPU.Copy
       Futhark.CodeGen.ImpGen.GPU.Group
       Futhark.CodeGen.ImpGen.GPU.SegHist
       Futhark.CodeGen.ImpGen.GPU.SegMap
@@ -494,8 +495,8 @@
     , haskeline
     , language-c-quote >= 0.12
     , lens
-    , lsp >= 2.0.0.0
-    , lsp-types >= 2.0.0.0
+    , lsp >= 2.1.0.0
+    , lsp-types >= 2.0.1.0
     , mainland-pretty >=0.7.1
     , cmark-gfm >=0.2.1
     , megaparsec >=9.0.0
diff --git a/rts/c/atomics.h b/rts/c/atomics.h
--- a/rts/c/atomics.h
+++ b/rts/c/atomics.h
@@ -1,5 +1,30 @@
 // Start of atomics.h
 
+inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x);
+inline int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,
+                                         int32_t cmp, int32_t val);
+inline int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p,
+                                        int32_t cmp, int32_t val);
+inline int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_add_i32_local(volatile __local int32_t *p, int32_t x);
+inline float atomic_fadd_f32_global(volatile __global float *p, float x);
+inline float atomic_fadd_f32_local(volatile __local float *p, float x);
+inline int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_smax_i32_local(volatile __local int32_t *p, int32_t x);
+inline int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_smin_i32_local(volatile __local int32_t *p, int32_t x);
+inline uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x);
+inline uint32_t atomic_umax_i32_local(volatile __local uint32_t *p, uint32_t x);
+inline uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x);
+inline uint32_t atomic_umin_i32_local(volatile __local uint32_t *p, uint32_t x);
+inline int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_and_i32_local(volatile __local int32_t *p, int32_t x);
+inline int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_or_i32_local(volatile __local int32_t *p, int32_t x);
+inline int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x);
+inline int32_t atomic_xor_i32_local(volatile __local int32_t *p, int32_t x);
+
 inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) {
 #ifdef FUTHARK_CUDA
   return atomicExch((int32_t*)p, x);
@@ -196,6 +221,36 @@
 
 // Start of 64 bit atomics
 
+#if defined(FUTHARK_CUDA) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)
+
+inline int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_xchg_i64_local(volatile __local int64_t *p, int64_t x);
+inline int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p,
+                                         int64_t cmp, int64_t val);
+inline int64_t atomic_cmpxchg_i64_local(volatile __local int64_t *p,
+                                        int64_t cmp, int64_t val);
+inline int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_add_i64_local(volatile __local int64_t *p, int64_t x);
+inline int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_smax_i64_local(volatile __local int64_t *p, int64_t x);
+inline int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_smin_i64_local(volatile __local int64_t *p, int64_t x);
+inline uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x);
+inline uint64_t atomic_umax_i64_local(volatile __local uint64_t *p, uint64_t x);
+inline uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x);
+uint64_t atomic_umin_i64_local(volatile __local uint64_t *p, uint64_t x);
+inline int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_and_i64_local(volatile __local int64_t *p, int64_t x);
+inline int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_or_i64_local(volatile __local int64_t *p, int64_t x);
+inline int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x);
+inline int64_t atomic_xor_i64_local(volatile __local int64_t *p, int64_t x);
+
+#ifdef FUTHARK_F64_ENABLED
+inline double atomic_fadd_f64_global(volatile __global double *p, double x);
+inline double atomic_fadd_f64_local(volatile __local double *p, double x);
+#endif
+
 inline int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x) {
 #ifdef FUTHARK_CUDA
   return atomicExch((uint64_t*)p, x);
@@ -393,5 +448,7 @@
   return atom_xor(p, x);
 #endif
 }
+
+#endif // defined(FUTHARK_CUDA) || defined(cl_khr_int64_base_atomics) && defined(cl_khr_int64_extended_atomics)
 
 // End of atomics.h
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
@@ -251,10 +251,10 @@
   lock_t error_lock;
   FILE *log;
   struct constants *constants;
-  struct free_list cu_free_list;
+  struct free_list free_list;
   int64_t peak_mem_usage_default;
   int64_t cur_mem_usage_default;
-  // Uniform above
+  // Uniform fields above.
 
   CUdeviceptr global_failure;
   CUdeviceptr global_failure_args;
@@ -270,8 +270,9 @@
   CUdevice dev;
   CUcontext cu_ctx;
   CUmodule module;
+  CUstream stream;
 
-  struct free_list free_list;
+  struct free_list cu_free_list;
 
   size_t max_block_size;
   size_t max_grid_size;
@@ -307,8 +308,8 @@
   struct futhark_context_config *cfg = ctx->cfg;
   char name[256];
   int count, chosen = -1, best_cc = -1;
-  int cc_major_best, cc_minor_best;
-  int cc_major, cc_minor;
+  int cc_major_best = 0, cc_minor_best = 0;
+  int cc_major = 0, cc_minor = 0;
   CUdevice dev;
 
   CUDA_SUCCEED_FATAL(cuDeviceGetCount(&count));
@@ -385,11 +386,11 @@
 }
 
 static const char *cuda_nvrtc_get_arch(CUdevice dev) {
-  struct {
+  static struct {
     int major;
     int minor;
     const char *arch_str;
-  } static const x[] = {
+  } const x[] = {
     { 3, 0, "compute_30" },
     { 3, 2, "compute_32" },
     { 3, 5, "compute_35" },
@@ -730,38 +731,15 @@
   return NULL;
 }
 
-static char* cuda_setup(struct futhark_context *ctx, const char *src_fragments[],
-                        const char *extra_opts[], const char* cache_fname) {
-  CUDA_SUCCEED_FATAL(cuInit(0));
-
-  if (cuda_device_setup(ctx) != 0) {
-    futhark_panic(-1, "No suitable CUDA device found.\n");
-  }
-  CUDA_SUCCEED_FATAL(cuCtxCreate(&ctx->cu_ctx, 0, ctx->dev));
-
-  free_list_init(&ctx->cu_free_list);
-
-  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);
-  ctx->max_block_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_block_size);
-  ctx->max_threshold = 0;
-  ctx->max_bespoke = 0;
-  ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);
-
-  cuda_size_setup(ctx);
-  return cuda_module_setup(ctx, src_fragments, extra_opts, cache_fname);
-}
-
 // Count up the runtime all the profiling_records that occured during execution.
 // Also clears the buffer of profiling_records.
-static cudaError_t cuda_tally_profiling_records(struct futhark_context *ctx) {
-  cudaError_t err;
+static CUresult cuda_tally_profiling_records(struct futhark_context *ctx) {
+  CUresult err;
   for (int i = 0; i < ctx->profiling_records_used; i++) {
     struct profiling_record record = ctx->profiling_records[i];
 
     float ms;
-    if ((err = cudaEventElapsedTime(&ms, record.events[0], record.events[1])) != cudaSuccess) {
+    if ((err = cuEventElapsedTime(&ms, record.events[0], record.events[1])) != CUDA_SUCCESS) {
       return err;
     }
 
@@ -769,10 +747,10 @@
     *record.runs += 1;
     *record.runtime += ms*1000;
 
-    if ((err = cudaEventDestroy(record.events[0])) != cudaSuccess) {
+    if ((err = cuEventDestroy(record.events[0])) != CUDA_SUCCESS) {
       return err;
     }
-    if ((err = cudaEventDestroy(record.events[1])) != cudaSuccess) {
+    if ((err = cuEventDestroy(record.events[1])) != CUDA_SUCCESS) {
       return err;
     }
 
@@ -781,7 +759,7 @@
 
   ctx->profiling_records_used = 0;
 
-  return cudaSuccess;
+  return CUDA_SUCCESS;
 }
 
 // Returns pointer to two events.
@@ -918,10 +896,28 @@
   ctx->peak_mem_usage_device = 0;
   ctx->cur_mem_usage_device = 0;
 
-  ctx->error = cuda_setup(ctx, cuda_program, ctx->cfg->nvrtc_opts, ctx->cfg->cache_fname);
+  CUDA_SUCCEED_FATAL(cuInit(0));
+  if (cuda_device_setup(ctx) != 0) {
+    futhark_panic(-1, "No suitable CUDA device found.\n");
+  }
+  CUDA_SUCCEED_FATAL(cuCtxCreate(&ctx->cu_ctx, 0, ctx->dev));
 
+  free_list_init(&ctx->cu_free_list);
+
+  ctx->max_shared_memory = device_query(ctx->dev, MAX_SHARED_MEMORY_PER_BLOCK);
+  ctx->max_block_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_block_size);
+  ctx->max_threshold = 0;
+  ctx->max_bespoke = 0;
+  ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE);
+  CUDA_SUCCEED_FATAL(cuStreamCreate(&ctx->stream, CU_STREAM_DEFAULT));
+  cuda_size_setup(ctx);
+  ctx->error = cuda_module_setup(ctx, cuda_program,
+                                 ctx->cfg->nvrtc_opts, ctx->cfg->cache_fname);
+
   if (ctx->error != NULL) {
-    futhark_panic(1, "%s\n", ctx->error);
+    futhark_panic(1, "During CUDA initialisation:\n%s\n", ctx->error);
   }
 
   int32_t no_error = -1;
@@ -938,6 +934,7 @@
   CUDA_SUCCEED_FATAL(cuda_free_all(ctx));
   (void)cuda_tally_profiling_records(ctx);
   free(ctx->profiling_records);
+  CUDA_SUCCEED_FATAL(cuStreamDestroy(ctx->stream));
   CUDA_SUCCEED_FATAL(cuModuleUnload(ctx->module));
   CUDA_SUCCEED_FATAL(cuCtxDestroy(ctx->cu_ctx));
 }
diff --git a/src/Futhark/CLI/C.hs b/src/Futhark/CLI/C.hs
--- a/src/Futhark/CLI/C.hs
+++ b/src/Futhark/CLI/C.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileCAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (sequentialCpuPipeline)
+import Futhark.Passes (seqmemPipeline)
 
 -- | Run @futhark c@
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile sequential C"
   "Generate sequential C code from optimised Futhark program."
-  sequentialCpuPipeline
+  seqmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileCAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/CUDA.hs b/src/Futhark/CLI/CUDA.hs
--- a/src/Futhark/CLI/CUDA.hs
+++ b/src/Futhark/CLI/CUDA.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileCUDAAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (gpuPipeline)
+import Futhark.Passes (gpumemPipeline)
 
 -- | Run @futhark cuda@.
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile CUDA"
   "Generate CUDA/C code from optimised Futhark program."
-  gpuPipeline
+  gpumemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileCUDAAction fcfg mode outpath) prog
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
@@ -639,7 +639,7 @@
       "GPU"
       GPU
       "Run the default optimised kernels pipeline"
-      kernelsPipeline
+      gpuPipeline
       []
       ["gpu"],
     pipelineOption
@@ -647,7 +647,7 @@
       "GPUMem"
       GPUMem
       "Run the full GPU compilation pipeline"
-      gpuPipeline
+      gpumemPipeline
       []
       ["gpu-mem"],
     pipelineOption
@@ -655,7 +655,7 @@
       "Seq"
       Seq
       "Run the sequential CPU compilation pipeline"
-      sequentialPipeline
+      seqPipeline
       []
       ["seq"],
     pipelineOption
@@ -663,7 +663,7 @@
       "SeqMem"
       SeqMem
       "Run the sequential CPU+memory compilation pipeline"
-      sequentialCpuPipeline
+      seqmemPipeline
       []
       ["seq-mem"],
     pipelineOption
@@ -679,7 +679,7 @@
       "MCMem"
       MCMem
       "Run the multicore+memory compilation pipeline"
-      multicorePipeline
+      mcmemPipeline
       []
       ["mc-mem"]
   ]
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
@@ -22,7 +22,7 @@
 import Futhark.Util (hashText, interactWithFileSafely)
 import Futhark.Util.Options
 import Futhark.Util.Pretty (prettyTextOneLine)
-import Language.Futhark.Parser.Lexer (scanTokens)
+import Language.Futhark.Parser.Lexer (scanTokensText)
 import Language.Futhark.Prop (isBuiltin)
 import Language.Futhark.Semantic (includeToString)
 import System.Environment (getExecutablePath)
@@ -115,7 +115,7 @@
 mainTokens = mainWithOptions () [] "program" $ \args () ->
   case args of
     [file] -> Just $ do
-      res <- interactWithFileSafely (scanTokens (startPos file) <$> BS.readFile file)
+      res <- interactWithFileSafely (scanTokensText (startPos file) <$> T.readFile file)
       case res of
         Nothing -> do
           hPutStrLn stderr $ file <> ": file not found."
diff --git a/src/Futhark/CLI/Multicore.hs b/src/Futhark/CLI/Multicore.hs
--- a/src/Futhark/CLI/Multicore.hs
+++ b/src/Futhark/CLI/Multicore.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileMulticoreAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (multicorePipeline)
+import Futhark.Passes (mcmemPipeline)
 
 -- | Run @futhark multicore@.
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile to multicore C"
   "Generate multicore C code from optimised Futhark program."
-  multicorePipeline
+  mcmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileMulticoreAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/MulticoreISPC.hs b/src/Futhark/CLI/MulticoreISPC.hs
--- a/src/Futhark/CLI/MulticoreISPC.hs
+++ b/src/Futhark/CLI/MulticoreISPC.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileMulticoreToISPCAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (multicorePipeline)
+import Futhark.Passes (mcmemPipeline)
 
 -- | Run @futhark multicore@.
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile to multicore ISPC"
   "Generate multicore ISPC code from optimised Futhark program."
-  multicorePipeline
+  mcmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileMulticoreToISPCAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/MulticoreWASM.hs b/src/Futhark/CLI/MulticoreWASM.hs
--- a/src/Futhark/CLI/MulticoreWASM.hs
+++ b/src/Futhark/CLI/MulticoreWASM.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileMulticoreToWASMAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (multicorePipeline)
+import Futhark.Passes (mcmemPipeline)
 
 -- | Run @futhark c@
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile to multicore WASM"
   "Generate multicore WASM with the multicore C backend code from optimised Futhark program."
-  multicorePipeline
+  mcmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileMulticoreToWASMAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/OpenCL.hs b/src/Futhark/CLI/OpenCL.hs
--- a/src/Futhark/CLI/OpenCL.hs
+++ b/src/Futhark/CLI/OpenCL.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileOpenCLAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (gpuPipeline)
+import Futhark.Passes (gpumemPipeline)
 
 -- | Run @futhark opencl@
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile OpenCL"
   "Generate OpenCL/C code from optimised Futhark program."
-  gpuPipeline
+  gpumemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileOpenCLAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/PyOpenCL.hs b/src/Futhark/CLI/PyOpenCL.hs
--- a/src/Futhark/CLI/PyOpenCL.hs
+++ b/src/Futhark/CLI/PyOpenCL.hs
@@ -12,6 +12,6 @@
   []
   "Compile PyOpenCL"
   "Generate Python + OpenCL code from optimised Futhark program."
-  gpuPipeline
+  gpumemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compilePyOpenCLAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/Python.hs b/src/Futhark/CLI/Python.hs
--- a/src/Futhark/CLI/Python.hs
+++ b/src/Futhark/CLI/Python.hs
@@ -12,6 +12,6 @@
   []
   "Compile sequential Python"
   "Generate sequential Python code from optimised Futhark program."
-  sequentialCpuPipeline
+  seqmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compilePythonAction fcfg mode outpath) prog
diff --git a/src/Futhark/CLI/REPL.hs b/src/Futhark/CLI/REPL.hs
--- a/src/Futhark/CLI/REPL.hs
+++ b/src/Futhark/CLI/REPL.hs
@@ -20,7 +20,7 @@
 import Futhark.MonadFreshNames
 import Futhark.Util (fancyTerminal)
 import Futhark.Util.Options
-import Futhark.Util.Pretty (AnsiStyle, Color (..), Doc, align, annotate, bgColorDull, bold, brackets, color, docText, docTextForHandle, hardline, oneLine, pretty, putDoc, putDocLn, unAnnotate, (<+>))
+import Futhark.Util.Pretty (AnsiStyle, Color (..), Doc, align, annotate, bgColorDull, bold, brackets, color, docText, docTextForHandle, hardline, italicized, oneLine, pretty, putDoc, putDocLn, unAnnotate, (<+>))
 import Futhark.Version
 import Language.Futhark
 import Language.Futhark.Interpreter qualified as I
@@ -220,9 +220,7 @@
               <> mconcat (intersperse ", " (map fst matches))
     _ -> do
       -- Read a declaration or expression.
-      maybe_dec_or_e <- parseDecOrExpIncrM (inputLine "  ") prompt line
-
-      case maybe_dec_or_e of
+      case parseDecOrExp prompt line of
         Left (SyntaxError _ err) -> liftIO $ T.putStrLn err
         Right (Left d) -> onDec d
         Right (Right e) -> onExp e
@@ -396,10 +394,13 @@
 
 typeCommand :: Command
 typeCommand = genTypeCommand parseExp T.checkExp $ \(ps, e) ->
-  oneLine (pretty e)
-    <> foldMap ((" " <>) . pretty) ps
-    <> " : "
-    <> oneLine (pretty (typeOf e))
+  oneLine (pretty (typeOf e))
+    <> if not (null ps)
+      then
+        annotate italicized $
+          "\n\nPolymorphic in"
+            <+> mconcat (intersperse " " $ map pretty ps) <> "."
+      else mempty
 
 mtypeCommand :: Command
 mtypeCommand = genTypeCommand parseModExp T.checkModExp $ pretty . fst
diff --git a/src/Futhark/CLI/WASM.hs b/src/Futhark/CLI/WASM.hs
--- a/src/Futhark/CLI/WASM.hs
+++ b/src/Futhark/CLI/WASM.hs
@@ -3,7 +3,7 @@
 
 import Futhark.Actions (compileCtoWASMAction)
 import Futhark.Compiler.CLI
-import Futhark.Passes (sequentialCpuPipeline)
+import Futhark.Passes (seqmemPipeline)
 
 -- | Run @futhark c@
 main :: String -> [String] -> IO ()
@@ -12,6 +12,6 @@
   []
   "Compile to WASM"
   "Generate WASM with the sequential C backend code from optimised Futhark program."
-  sequentialCpuPipeline
+  seqmemPipeline
   $ \fcfg () mode outpath prog ->
     actionProcedure (compileCtoWASMAction fcfg mode outpath) prog
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
@@ -11,6 +11,7 @@
 where
 
 import Control.Monad
+import Data.List (unzip4)
 import Data.Maybe (catMaybes)
 import Data.Text qualified as T
 import Futhark.CodeGen.Backends.CCUDA.Boilerplate
@@ -152,7 +153,7 @@
                 cuMemcpyHtoDAsync($exp:mem + $exp:idx * sizeof($ty:t),
                                   &$id:val',
                                   sizeof($ty:t),
-                                  0));
+                                  ctx->stream));
               $items:aft
              }|]
 writeCUDAScalar mem idx t "device" _ val = do
@@ -226,9 +227,9 @@
     memcpyFun _ (Space "device") (Space "device") =
       ([C.cexp|cuMemcpy($exp:dst, $exp:src, $exp:nbytes)|], copyDevToDev)
     memcpyFun GC.CopyNoBarrier DefaultSpace (Space "device") =
-      ([C.cexp|cuMemcpyDtoHAsync($exp:dst, $exp:src, $exp:nbytes, 0)|], copyDevToHost)
+      ([C.cexp|cuMemcpyDtoHAsync($exp:dst, $exp:src, $exp:nbytes, ctx->stream)|], copyDevToHost)
     memcpyFun GC.CopyNoBarrier (Space "device") DefaultSpace =
-      ([C.cexp|cuMemcpyHtoDAsync($exp:dst, $exp:src, $exp:nbytes, 0)|], copyHostToDev)
+      ([C.cexp|cuMemcpyHtoDAsync($exp:dst, $exp:src, $exp:nbytes, ctx->stream)|], copyHostToDev)
     memcpyFun _ _ _ =
       error $
         "Cannot copy to '"
@@ -274,21 +275,81 @@
   let e = kernelConstToExp $ SizeMaxConst size_class
   GC.stm [C.cstm|$id:v = $exp:e;|]
 callKernel (LaunchKernel safety kernel_name args num_blocks block_size) = do
-  args_arr <- newVName "kernel_args"
-  time_start <- newVName "time_start"
-  time_end <- newVName "time_end"
-  (args', shared_vars) <- mapAndUnzipM mkArgs args
+  (arg_params, arg_params_inits, call_args, shared_vars) <-
+    unzip4 <$> zipWithM mkArgs [(0 :: Int) ..] args
   let (shared_sizes, shared_offsets) = unzip $ catMaybes shared_vars
       shared_offsets_sc = mkOffsets shared_sizes
       shared_args = zip shared_offsets shared_offsets_sc
-      shared_tot = last shared_offsets_sc
+      shared_bytes = last shared_offsets_sc
   forM_ shared_args $ \(arg, offset) ->
     GC.decl [C.cdecl|unsigned int $id:arg = $exp:offset;|]
 
   (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_blocks
   (block_x, block_y, block_z) <- mkDims <$> mapM compileGroupDim block_size
-  let perm_args
-        | length num_blocks == 3 = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]
+
+  let need_perm = length num_blocks == 3
+  kernel_fname <- genKernelFunction kernel_name safety need_perm arg_params arg_params_inits
+
+  GC.stm
+    [C.cstm|{
+           err = $id:kernel_fname(ctx,
+                                  $exp:grid_x,$exp:grid_y,$exp:grid_z,
+                                  $exp:block_x, $exp:block_y, $exp:block_z,
+                                  $exp:shared_bytes,
+                                  $args:call_args);
+           if (err != FUTHARK_SUCCESS) { goto cleanup; }
+           }|]
+
+  when (safety >= SafetyFull) $
+    GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]
+  where
+    mkDims [] = ([C.cexp|0|], [C.cexp|0|], [C.cexp|0|])
+    mkDims [x] = (x, [C.cexp|1|], [C.cexp|1|])
+    mkDims [x, y] = (x, y, [C.cexp|1|])
+    mkDims (x : y : z : _) = (x, y, z)
+    addExp x y = [C.cexp|$exp:x + $exp:y|]
+    alignExp e = [C.cexp|$exp:e + ((8 - ($exp:e % 8)) % 8)|]
+    mkOffsets = scanl (\a b -> a `addExp` alignExp b) [C.cexp|0|]
+    mkArgs i (ValueKArg e t) = do
+      e' <- GC.compileExp e
+      pure
+        ( [C.cparam|$ty:(primStorageType t) $id:("arg" <> show i)|],
+          [C.cinit|&$id:("arg" <> show i)|],
+          toStorage t e',
+          Nothing
+        )
+    mkArgs i (MemKArg v) = do
+      v' <- GC.rawMem v
+      pure
+        ( [C.cparam|typename CUdeviceptr $id:("arg" <> show i)|],
+          [C.cinit|&$id:("arg" <> show i)|],
+          v',
+          Nothing
+        )
+    mkArgs i (SharedMemoryKArg (Count c)) = do
+      num_bytes <- GC.compileExp c
+      size <- newVName "shared_size"
+      offset <- newVName "shared_offset"
+      GC.decl [C.cdecl|unsigned int $id:size = $exp:num_bytes;|]
+      pure
+        ( [C.cparam|unsigned int $id:("arg" <> show i)|],
+          [C.cinit|&$id:("arg" <> show i)|],
+          [C.cexp|$id:offset|],
+          Just (size, offset)
+        )
+
+genKernelFunction ::
+  KernelName ->
+  KernelSafety ->
+  Bool ->
+  [C.Param] ->
+  [C.Initializer] ->
+  GC.CompilerM op s Name
+genKernelFunction kernel_name safety need_perm arg_params arg_params_inits = do
+  let kernel_fname = "gpu_kernel_" <> kernel_name
+      (bef, aft) = profilingEnclosure kernel_name
+      perm_args
+        | need_perm = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]
         | otherwise = []
       failure_args =
         take
@@ -297,92 +358,56 @@
             [C.cinit|&ctx->failure_is_an_option|],
             [C.cinit|&ctx->global_failure_args|]
           ]
-      args'' = perm_args ++ failure_args ++ [[C.cinit|&$id:a|] | a <- args']
-      sizes_nonzero =
-        expsNotZero
-          [ grid_x,
-            grid_y,
-            grid_z,
-            block_x,
-            block_y,
-            block_z
-          ]
-      (bef, aft) = profilingEnclosure kernel_name
 
-  GC.stm
-    [C.cstm|
-    if ($exp:sizes_nonzero) {
+  GC.libDecl
+    [C.cedecl|static int $id:kernel_fname(
+                struct futhark_context* ctx, unsigned int grid_x, unsigned int grid_y, unsigned int grid_z,
+                unsigned int block_x, unsigned int block_y, unsigned int block_z,
+                unsigned int shared_bytes, $params:arg_params) {
+    if (grid_x * grid_y * grid_z * block_x * block_y * block_z != 0) {
       int perm[3] = { 0, 1, 2 };
 
-      if ($exp:grid_y >= (1<<16)) {
+      if (grid_y >= (1<<16)) {
         perm[1] = perm[0];
         perm[0] = 1;
       }
 
-      if ($exp:grid_z >= (1<<16)) {
+      if (grid_z >= (1<<16)) {
         perm[2] = perm[0];
         perm[0] = 2;
       }
 
       size_t grid[3];
-      grid[perm[0]] = $exp:grid_x;
-      grid[perm[1]] = $exp:grid_y;
-      grid[perm[2]] = $exp:grid_z;
+      grid[perm[0]] = grid_x;
+      grid[perm[1]] = grid_y;
+      grid[perm[2]] = grid_z;
 
-      void *$id:args_arr[] = { $inits:args'' };
-      typename int64_t $id:time_start = 0, $id:time_end = 0;
+      void *all_args[] = { $inits:(perm_args ++ failure_args ++ arg_params_inits) };
+      typename int64_t time_start = 0, time_end = 0;
       if (ctx->debugging) {
-        fprintf(ctx->log, "Launching %s with grid size [%ld, %ld, %ld] and block size [%ld, %ld, %ld]; shared memory: %d bytes.\n",
+        fprintf(ctx->log, "Launching %s with grid size [%d, %d, %d] and block size [%d, %d, %d]; shared memory: %d bytes.\n",
                 $string:(prettyString kernel_name),
-                (long int)$exp:grid_x, (long int)$exp:grid_y, (long int)$exp:grid_z,
-                (long int)$exp:block_x, (long int)$exp:block_y, (long int)$exp:block_z,
-                (int)$exp:shared_tot);
-        $id:time_start = get_wall_time();
+                grid_x, grid_y, grid_z,
+                block_x, block_y, block_z,
+                shared_bytes);
+        time_start = get_wall_time();
       }
       $items:bef
       CUDA_SUCCEED_OR_RETURN(
         cuLaunchKernel(ctx->program->$id:kernel_name,
                        grid[0], grid[1], grid[2],
-                       $exp:block_x, $exp:block_y, $exp:block_z,
-                       $exp:shared_tot, NULL,
-                       $id:args_arr, NULL));
+                       block_x, block_y, block_z,
+                       shared_bytes, ctx->stream,
+                       all_args, NULL));
       $items:aft
       if (ctx->debugging) {
         CUDA_SUCCEED_FATAL(cuCtxSynchronize());
-        $id:time_end = get_wall_time();
+        time_end = get_wall_time();
         fprintf(ctx->log, "Kernel %s runtime: %ldus\n",
-                $string:(prettyString kernel_name), $id:time_end - $id:time_start);
+                $string:(prettyString kernel_name), time_end - time_start);
       }
-    }|]
+    }
+    return FUTHARK_SUCCESS;
+  }|]
 
-  when (safety >= SafetyFull) $
-    GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]
-  where
-    mkDims [] = ([C.cexp|0|], [C.cexp|0|], [C.cexp|0|])
-    mkDims [x] = (x, [C.cexp|1|], [C.cexp|1|])
-    mkDims [x, y] = (x, y, [C.cexp|1|])
-    mkDims (x : y : z : _) = (x, y, z)
-    addExp x y = [C.cexp|$exp:x + $exp:y|]
-    alignExp e = [C.cexp|$exp:e + ((8 - ($exp:e % 8)) % 8)|]
-    mkOffsets = scanl (\a b -> a `addExp` alignExp b) [C.cexp|0|]
-    expNotZero e = [C.cexp|$exp:e != 0|]
-    expAnd a b = [C.cexp|$exp:a && $exp:b|]
-    expsNotZero = foldl expAnd [C.cexp|1|] . map expNotZero
-    mkArgs (ValueKArg e t@(FloatType Float16)) = do
-      arg <- newVName "kernel_arg"
-      e' <- GC.compileExp e
-      GC.item [C.citem|$ty:(primStorageType t) $id:arg = $exp:(toStorage t e');|]
-      pure (arg, Nothing)
-    mkArgs (ValueKArg e t) =
-      (,Nothing) <$> GC.compileExpToName "kernel_arg" t e
-    mkArgs (MemKArg v) = do
-      v' <- GC.rawMem v
-      arg <- newVName "kernel_arg"
-      GC.decl [C.cdecl|typename CUdeviceptr $id:arg = $exp:v';|]
-      pure (arg, Nothing)
-    mkArgs (SharedMemoryKArg (Count c)) = do
-      num_bytes <- GC.compileExp c
-      size <- newVName "shared_size"
-      offset <- newVName "shared_offset"
-      GC.decl [C.cdecl|unsigned int $id:size = $exp:num_bytes;|]
-      pure (offset, Just (size, offset))
+  pure kernel_fname
diff --git a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs
@@ -37,17 +37,17 @@
 profilingEnclosure :: Name -> ([C.BlockItem], [C.BlockItem])
 profilingEnclosure name =
   ( [C.citems|
-      typename cudaEvent_t *pevents = NULL;
+      typename CUevent *pevents = NULL;
       if (ctx->profiling && !ctx->profiling_paused) {
         pevents = cuda_get_events(ctx,
                                   &ctx->program->$id:(kernelRuns name),
                                   &ctx->program->$id:(kernelRuntime name));
-        CUDA_SUCCEED_FATAL(cudaEventRecord(pevents[0], 0));
+        CUDA_SUCCEED_FATAL(cuEventRecord(pevents[0], ctx->stream));
       }
       |],
     [C.citems|
       if (pevents != NULL) {
-        CUDA_SUCCEED_FATAL(cudaEventRecord(pevents[1], 0));
+        CUDA_SUCCEED_FATAL(cuEventRecord(pevents[1], ctx->stream));
       }
       |]
   )
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
@@ -11,7 +11,6 @@
 where
 
 import Control.Monad hiding (mapM)
-import Data.List (intercalate)
 import Data.Text qualified as T
 import Futhark.CodeGen.Backends.COpenCL.Boilerplate
 import Futhark.CodeGen.Backends.GenericC qualified as GC
@@ -320,111 +319,107 @@
                                               &ctx->failure_is_an_option));
     |]
 
-  zipWithM_ setKernelArg [numFailureParams safety ..] args
+  (arg_params, arg_set, call_args) <-
+    unzip3 <$> zipWithM onArg [(0 :: Int) ..] args
+
   num_workgroups' <- mapM GC.compileExp num_workgroups
   workgroup_size' <- mapM compileGroupDim workgroup_size
   local_bytes <- foldM localBytes [C.cexp|0|] args
 
-  launchKernel name num_workgroups' workgroup_size' local_bytes
+  kernel_fname <- genKernelFunction name safety arg_params arg_set
 
+  let grid_x : grid_y : grid_z : _ = num_workgroups' ++ repeat [C.cexp|1|]
+      group_x : group_y : group_z : _ = workgroup_size' ++ repeat [C.cexp|1|]
+
+  GC.stm
+    [C.cstm|{
+           err = $id:kernel_fname(ctx,
+                                  $exp:grid_x,$exp:grid_y,$exp:grid_z,
+                                  $exp:group_x, $exp:group_y, $exp:group_z,
+                                  $exp:local_bytes,
+                                  $args:call_args);
+           if (err != FUTHARK_SUCCESS) { goto cleanup; }
+           }|]
+
   when (safety >= SafetyFull) $
     GC.stm [C.cstm|ctx->failure_is_an_option = 1;|]
   where
-    setKernelArg i (ValueKArg e pt) = do
-      v <- case pt of
-        -- We always transfer f16 values to the kernel as 16 bits, but
-        -- the actual host type may be typedef'd to a 32-bit float.
-        -- This requires some care.
-        FloatType Float16 -> do
-          v <- newVName "kernel_arg"
-          e' <- toStorage pt <$> GC.compileExp e
-          GC.decl [C.cdecl|$ty:(primStorageType pt) $id:v = $e';|]
-          pure v
-        _ -> GC.compileExpToName "kernel_arg" pt e
-      GC.stm
-        [C.cstm|
-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, sizeof($id:v), &$id:v));
-          |]
-    setKernelArg i (MemKArg v) = do
-      v' <- GC.rawMem v
-      GC.stm
-        [C.cstm|
-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, sizeof($exp:v'), &$exp:v'));
-          |]
-    setKernelArg i (SharedMemoryKArg num_bytes) = do
-      num_bytes' <- GC.compileExp $ unCount num_bytes
-      GC.stm
-        [C.cstm|
-            OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:name, $int:i, (size_t)$exp:num_bytes', NULL));
-            |]
-
     localBytes cur (SharedMemoryKArg num_bytes) = do
       num_bytes' <- GC.compileExp $ unCount num_bytes
       pure [C.cexp|$exp:cur + $exp:num_bytes'|]
     localBytes cur _ = pure cur
 
-launchKernel ::
-  C.ToExp a =>
-  KernelName ->
-  [a] ->
-  [a] ->
-  a ->
-  GC.CompilerM op s ()
-launchKernel kernel_name num_workgroups workgroup_dims local_bytes = do
-  global_work_size <- newVName "global_work_size"
-  time_start <- newVName "time_start"
-  time_end <- newVName "time_end"
-  time_diff <- newVName "time_diff"
-  local_work_size <- newVName "local_work_size"
-
-  let (debug_str, debug_args) = debugPrint global_work_size local_work_size
+    onArg i (ValueKArg e t) = do
+      let arg = "arg" <> show i
+      e' <- GC.compileExp e
+      pure
+        ( [C.cparam|$ty:(primStorageType t) $id:arg|],
+          ([C.cexp|sizeof($id:arg)|], [C.cexp|&$id:arg|]),
+          toStorage t e'
+        )
+    onArg i (MemKArg v) = do
+      let arg = "arg" <> show i
+      v' <- GC.rawMem v
+      pure
+        ( [C.cparam|typename cl_mem $id:arg|],
+          ([C.cexp|sizeof($id:arg)|], [C.cexp|&$id:arg|]),
+          v'
+        )
+    onArg i (SharedMemoryKArg (Count c)) = do
+      let arg = "arg" <> show i
+      num_bytes <- GC.compileExp c
+      pure
+        ( [C.cparam|unsigned int $id:arg|],
+          ([C.cexp|$id:arg|], [C.cexp|NULL|]),
+          num_bytes
+        )
 
-  GC.stm
-    [C.cstm|
-    if ($exp:total_elements != 0) {
-      const size_t $id:global_work_size[$int:kernel_rank] = {$inits:kernel_dims'};
-      const size_t $id:local_work_size[$int:kernel_rank] = {$inits:workgroup_dims'};
-      typename int64_t $id:time_start = 0, $id:time_end = 0;
+genKernelFunction ::
+  KernelName ->
+  KernelSafety ->
+  [C.Param] ->
+  [(C.Exp, C.Exp)] ->
+  GC.CompilerM op s Name
+genKernelFunction kernel_name safety arg_params arg_set = do
+  let kernel_fname = "gpu_kernel_" <> kernel_name
+  GC.libDecl
+    [C.cedecl|static int $id:kernel_fname(
+                struct futhark_context* ctx,
+                unsigned int grid_x, unsigned int grid_y, unsigned int grid_z,
+                unsigned int block_x, unsigned int block_y, unsigned int block_z,
+                unsigned int local_bytes, $params:arg_params) {
+    (void)local_bytes;
+    if (grid_x * grid_y * grid_z * block_x * block_y * block_z != 0) {
+      const size_t global_work_size[3] = {grid_x*block_x, grid_y*block_y, grid_z*block_z};
+      const size_t local_work_size[3] = {block_x, block_y, block_z};
+      typename int64_t time_start = 0, time_end = 0;
+      $stms:set_args
       if (ctx->debugging) {
-        fprintf(ctx->log, $string:debug_str, $args:debug_args);
-        $id:time_start = get_wall_time();
+        fprintf(ctx->log, "Launching %s with grid size [%d, %d, %d] and group size [%d, %d, %d]; local memory: %d bytes.\n",
+                $string:(prettyString kernel_name),
+                grid_x, grid_y, grid_z,
+                block_x, block_y, block_z,
+                local_bytes);
+        time_start = get_wall_time();
       }
       typename cl_event *pevent = $exp:(profilingEvent kernel_name);
       OPENCL_SUCCEED_OR_RETURN(
-        clEnqueueNDRangeKernel(ctx->queue, ctx->program->$id:kernel_name, $int:kernel_rank, NULL,
-                               $id:global_work_size, $id:local_work_size,
+        clEnqueueNDRangeKernel(ctx->queue, ctx->program->$id:kernel_name, 3, NULL,
+                               global_work_size, local_work_size,
                                0, NULL, pevent));
       if (ctx->debugging) {
         OPENCL_SUCCEED_FATAL(clFinish(ctx->queue));
-        $id:time_end = get_wall_time();
-        long int $id:time_diff = $id:time_end - $id:time_start;
+        time_end = get_wall_time();
+        long int time_diff = time_end - time_start;
         fprintf(ctx->log, "kernel %s runtime: %ldus\n",
-                $string:(prettyString kernel_name), $id:time_diff);
+                $string:(prettyString kernel_name), time_diff);
       }
-    }|]
-  where
-    kernel_rank = length kernel_dims
-    kernel_dims = zipWith multExp (map toSize num_workgroups) (map toSize workgroup_dims)
-    kernel_dims' = map toInit kernel_dims
-    workgroup_dims' = map (toInit . toSize) workgroup_dims
-    total_elements = foldl multExp [C.cexp|1|] kernel_dims
-
-    toInit e = [C.cinit|$exp:e|]
-    multExp x y = [C.cexp|$exp:x * $exp:y|]
-    toSize e = [C.cexp|(size_t)$exp:e|]
+    }
+    return FUTHARK_SUCCESS;
+  }|]
 
-    debugPrint :: VName -> VName -> (String, [C.Exp])
-    debugPrint global_work_size local_work_size =
-      ( "Launching %s with global work size "
-          ++ dims
-          ++ " and local work size "
-          ++ dims
-          ++ "; local memory: %d bytes.\n",
-        [C.cexp|$string:(prettyString kernel_name)|]
-          : map (kernelDim global_work_size) [0 .. kernel_rank - 1]
-          ++ map (kernelDim local_work_size) [0 .. kernel_rank - 1]
-          ++ [[C.cexp|(int)$exp:local_bytes|]]
-      )
-      where
-        dims = "[" ++ intercalate ", " (replicate kernel_rank "%zu") ++ "]"
-        kernelDim arr i = [C.cexp|$id:arr[$int:i]|]
+  pure kernel_fname
+  where
+    set_args = zipWith setKernelArg [numFailureParams safety ..] arg_set
+    setKernelArg i (size, e) =
+      [C.cstm|OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->program->$id:kernel_name, $int:i, $exp:size, $exp:e));|]
diff --git a/src/Futhark/CodeGen/Backends/SimpleRep.hs b/src/Futhark/CodeGen/Backends/SimpleRep.hs
--- a/src/Futhark/CodeGen/Backends/SimpleRep.hs
+++ b/src/Futhark/CodeGen/Backends/SimpleRep.hs
@@ -209,11 +209,11 @@
 instance C.ToExp FloatValue where
   toExp (Float16Value x) _
     | isInfinite x =
-        if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|]
+        if x > 0 then [C.cexp|(typename f16)INFINITY|] else [C.cexp|(typename f16)-INFINITY|]
     | isNaN x =
-        [C.cexp|NAN|]
+        [C.cexp|(typename f16)NAN|]
     | otherwise =
-        [C.cexp|$float:(fromRational (toRational x))|]
+        [C.cexp|(typename f16)$float:(fromRational (toRational x))|]
   toExp (Float32Value x) _
     | isInfinite x =
         if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|]
diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs
--- a/src/Futhark/CodeGen/ImpCode.hs
+++ b/src/Futhark/CodeGen/ImpCode.hs
@@ -602,8 +602,11 @@
         If {} -> pretty fbranch
         _ ->
           "{" </> indent 2 (pretty fbranch) </> "}"
+  pretty (Call [] fname args) =
+    "call" <+> pretty fname <> parens (commasep $ map pretty args)
   pretty (Call dests fname args) =
-    commasep (map pretty dests)
+    "call"
+      <+> commasep (map pretty dests)
       <+> "<-"
       <+> pretty fname <> parens (commasep $ map pretty args)
   pretty (Comment s code) =
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
@@ -146,6 +146,7 @@
 import Futhark.Construct hiding (ToExp (..))
 import Futhark.IR.Mem
 import Futhark.IR.Mem.IxFun qualified as IxFun
+import Futhark.IR.Mem.LMAD qualified as LMAD
 import Futhark.IR.SOACS (SOACS)
 import Futhark.Util
 import Futhark.Util.IntegralExp
@@ -207,9 +208,9 @@
 sliceMemLoc (MemLoc mem shape ixfun) slice =
   MemLoc mem shape $ IxFun.slice ixfun slice
 
-flatSliceMemLoc :: MemLoc -> FlatSlice (Imp.TExp Int64) -> MemLoc
+flatSliceMemLoc :: MemLoc -> FlatSlice (Imp.TExp Int64) -> Maybe MemLoc
 flatSliceMemLoc (MemLoc mem shape ixfun) slice =
-  MemLoc mem shape $ IxFun.flatSlice ixfun slice
+  MemLoc mem shape <$> IxFun.flatSlice ixfun slice
 
 data ArrayEntry = ArrayEntry
   { entryArrayLoc :: MemLoc,
@@ -931,7 +932,9 @@
 defCompileBasicOp (Pat [pe]) (FlatUpdate _ slice v) = do
   pe_loc <- entryArrayLoc <$> lookupArray (patElemName pe)
   v_loc <- entryArrayLoc <$> lookupArray v
-  copy (elemType (patElemType pe)) (flatSliceMemLoc pe_loc slice') v_loc
+  case flatSliceMemLoc pe_loc slice' of
+    Just pe_loc' -> copy (elemType (patElemType pe)) pe_loc' v_loc
+    Nothing -> error "defCompileBasicOp FlatUpdate"
   where
     slice' = fmap pe64 slice
 defCompileBasicOp (Pat [pe]) (Replicate shape se)
@@ -1401,8 +1404,8 @@
 copy :: CopyCompiler rep r op
 copy
   bt
-  dst@(MemLoc dst_name _ dst_ixfn@(IxFun.IxFun dst_lmads@(dst_lmad :| _) _ _))
-  src@(MemLoc src_name _ src_ixfn@(IxFun.IxFun src_lmads@(src_lmad :| _) _ _)) = do
+  dst@(MemLoc dst_name _ dst_ixfn@(IxFun.IxFun dst_lmad _))
+  src@(MemLoc src_name _ src_ixfn@(IxFun.IxFun src_lmad _)) = do
     -- If we can statically determine that the two index-functions
     -- are equivalent, don't do anything
     unless (dst_name == src_name && dst_ixfn `IxFun.equivalent` src_ixfn)
@@ -1410,7 +1413,7 @@
       -- It's also possible that we can dynamically determine that the two
       -- index-functions are equivalent.
       sUnless
-        ( fromBool (dst_name == src_name && length dst_lmads == 1 && length src_lmads == 1)
+        ( fromBool (dst_name == src_name)
             .&&. IxFun.dynamicEqualsLMAD dst_lmad src_lmad
         )
       $ do
@@ -1418,44 +1421,54 @@
         cc <- asks envCopyCompiler
         cc bt dst src
 
--- | Is this copy really a mapping with transpose?
+-- | Is this copy really a mapping with transpose?  Produce an
+-- expression that is true if so, as well as other expressions that
+-- contain information about the transpose in that case (don't trust
+-- these if the boolean is false).
 isMapTransposeCopy ::
   PrimType ->
   MemLoc ->
   MemLoc ->
   Maybe
-    ( Imp.TExp Int64,
-      Imp.TExp Int64,
-      Imp.TExp Int64,
-      Imp.TExp Int64,
-      Imp.TExp Int64
+    ( Imp.TExp Bool,
+      ( Imp.TExp Int64,
+        Imp.TExp Int64,
+        Imp.TExp Int64,
+        Imp.TExp Int64,
+        Imp.TExp Int64
+      )
     )
-isMapTransposeCopy bt (MemLoc _ _ destIxFun) (MemLoc _ _ srcIxFun)
-  | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun bt_size,
-    (perm, destshape) <- unzip perm_and_destshape,
-    Just src_offset <- IxFun.linearWithOffset srcIxFun bt_size,
+isMapTransposeCopy pt (MemLoc _ _ destIxFun) (MemLoc _ _ srcIxFun)
+  | perm <- LMAD.permutation dest_lmad,
+    LMAD.permutation src_lmad == [0 .. rank - 1],
     Just (r1, r2, _) <- isMapTranspose perm =
-      isOk destshape swap r1 r2 dest_offset src_offset
-  | Just dest_offset <- IxFun.linearWithOffset destIxFun bt_size,
-    Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun bt_size,
-    (perm, srcshape) <- unzip perm_and_srcshape,
+      isOk (IxFun.shape destIxFun) swap r1 r2
+  | perm <- LMAD.permutation src_lmad,
+    LMAD.permutation dest_lmad == [0 .. rank - 1],
     Just (r1, r2, _) <- isMapTranspose perm =
-      isOk srcshape id r1 r2 dest_offset src_offset
+      isOk (IxFun.shape srcIxFun) id r1 r2
   | otherwise =
       Nothing
   where
-    bt_size = primByteSize bt
+    rank = IxFun.rank destIxFun
     swap (x, y) = (y, x)
+    dest_lmad = IxFun.ixfunLMAD destIxFun
+    src_lmad = IxFun.ixfunLMAD srcIxFun
+    dest_offset = LMAD.offset dest_lmad
+    src_offset = LMAD.offset src_lmad
 
-    isOk shape f r1 r2 dest_offset src_offset = do
+    isOk shape f r1 r2 =
       let (num_arrays, size_x, size_y) = getSizes shape f r1 r2
-      pure
-        ( dest_offset,
-          src_offset,
-          num_arrays,
-          size_x,
-          size_y
-        )
+       in Just
+            ( LMAD.contiguous dest_lmad
+                .&&. LMAD.contiguous src_lmad,
+              ( dest_offset * primByteSize pt,
+                src_offset * primByteSize pt,
+                num_arrays,
+                size_x,
+                size_y
+              )
+            )
 
     getSizes shape f r1 r2 =
       let (mapped, notmapped) = splitAt r1 shape
@@ -1477,35 +1490,25 @@
 -- | Use 'sCopy' if possible, otherwise 'copyElementWise'.
 defaultCopy :: CopyCompiler rep r op
 defaultCopy pt dest src
-  | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <-
+  | Just (is_transpose, (destoffset, srcoffset, num_arrays, size_x, size_y)) <-
       isMapTransposeCopy pt dest src = do
       fname <- mapTransposeForType pt
-      emit
-        $ Imp.Call
-          []
-          fname
-        $ transposeArgs
-          pt
-          destmem
-          (bytes destoffset)
-          srcmem
-          (bytes srcoffset)
-          num_arrays
-          size_x
-          size_y
-  | Just destoffset <-
-      IxFun.linearWithOffset dest_ixfun pt_size,
-    Just srcoffset <-
-      IxFun.linearWithOffset src_ixfun pt_size = do
-      srcspace <- entryMemSpace <$> lookupMemory srcmem
-      destspace <- entryMemSpace <$> lookupMemory destmem
-      if isScalarSpace srcspace || isScalarSpace destspace
-        then copyElementWise pt dest src
-        else sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt
-  | otherwise =
-      copyElementWise pt dest src
+      sIf
+        is_transpose
+        ( emit . Imp.Call [] fname $
+            transposeArgs
+              pt
+              destmem
+              (bytes destoffset)
+              srcmem
+              (bytes srcoffset)
+              num_arrays
+              size_x
+              size_y
+        )
+        nontranspose
+  | otherwise = nontranspose
   where
-    pt_size = primByteSize pt
     num_elems = Imp.elements $ product $ IxFun.shape $ memLocIxFun src
 
     MemLoc destmem _ dest_ixfun = dest
@@ -1514,6 +1517,21 @@
     isScalarSpace ScalarSpace {} = True
     isScalarSpace _ = False
 
+    nontranspose = do
+      srcspace <- entryMemSpace <$> lookupMemory srcmem
+      destspace <- entryMemSpace <$> lookupMemory destmem
+      if isScalarSpace srcspace || isScalarSpace destspace
+        then copyElementWise pt dest src
+        else do
+          let dest_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD dest_ixfun
+              src_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD src_ixfun
+              destoffset = elements (LMAD.offset dest_lmad) `withElemType` pt
+              srcoffset = elements (LMAD.offset src_lmad) `withElemType` pt
+          sIf
+            (LMAD.memcpyable dest_lmad src_lmad)
+            (sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt)
+            (copyElementWise pt dest src)
+
 copyElementWise :: CopyCompiler rep r op
 copyElementWise bt dest src = do
   let bounds = IxFun.shape $ memLocIxFun src
@@ -1889,17 +1907,17 @@
 
 sCopy ::
   VName ->
-  Imp.TExp Int64 ->
+  Count Bytes (Imp.TExp Int64) ->
   Space ->
   VName ->
-  Imp.TExp Int64 ->
+  Count Bytes (Imp.TExp Int64) ->
   Space ->
   Count Elements (Imp.TExp Int64) ->
   PrimType ->
   ImpM rep r op ()
 sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt =
   if destmem == srcmem
-    then sUnless (destoffset .==. srcoffset) the_copy
+    then sUnless (Imp.unCount destoffset .==. Imp.unCount srcoffset) the_copy
     else the_copy
   where
     the_copy =
@@ -1907,10 +1925,10 @@
         $ Imp.Copy
           pt
           destmem
-          (bytes destoffset)
+          destoffset
           destspace
           srcmem
-          (bytes srcoffset)
+          srcoffset
           srcspace
         $ num_elems `withElemType` pt
 
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
@@ -12,6 +12,8 @@
 where
 
 import Control.Monad
+import Control.Monad.State
+import Data.Foldable (toList)
 import Data.List (foldl')
 import Data.Map qualified as M
 import Data.Maybe
@@ -19,6 +21,7 @@
 import Futhark.CodeGen.ImpGen hiding (compileProg)
 import Futhark.CodeGen.ImpGen qualified
 import Futhark.CodeGen.ImpGen.GPU.Base
+import Futhark.CodeGen.ImpGen.GPU.Copy
 import Futhark.CodeGen.ImpGen.GPU.SegHist
 import Futhark.CodeGen.ImpGen.GPU.SegMap
 import Futhark.CodeGen.ImpGen.GPU.SegRed
@@ -27,6 +30,7 @@
 import Futhark.Error
 import Futhark.IR.GPUMem
 import Futhark.IR.Mem.IxFun qualified as IxFun
+import Futhark.IR.Mem.LMAD qualified as LMAD
 import Futhark.MonadFreshNames
 import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem)
 import Prelude hiding (quot, rem)
@@ -231,10 +235,12 @@
   s' <- toExp s
 
   sIota (patElemName pe) (pe64 n) x' s' et
-expCompiler (Pat [pe]) (BasicOp (Replicate _ se))
+expCompiler (Pat [pe]) (BasicOp (Replicate shape se))
   | Acc {} <- patElemType pe = pure ()
   | otherwise =
-      sReplicate (patElemName pe) se
+      if shapeRank shape == 0
+        then copyDWIM (patElemName pe) [] se []
+        else sReplicate (patElemName pe) se
 -- Allocation in the "local" space is just a placeholder.
 expCompiler _ (Op (Alloc _ (Space "local"))) =
   pure ()
@@ -257,40 +263,92 @@
 expCompiler dest e =
   defCompileExp dest e
 
-callKernelCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp
-callKernelCopy bt destloc@(MemLoc destmem _ destIxFun) srcloc@(MemLoc srcmem srcshape srcIxFun)
-  | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <-
-      isMapTransposeCopy bt destloc srcloc = do
-      fname <- mapTransposeForType bt
-      emit $
-        Imp.Call
-          []
-          fname
-          [ Imp.MemArg destmem,
-            Imp.ExpArg $ untyped destoffset,
-            Imp.MemArg srcmem,
-            Imp.ExpArg $ untyped srcoffset,
-            Imp.ExpArg $ untyped num_arrays,
-            Imp.ExpArg $ untyped size_x,
-            Imp.ExpArg $ untyped size_y
-          ]
-  | bt_size <- primByteSize bt,
-    Just destoffset <- IxFun.linearWithOffset destIxFun bt_size,
-    Just srcoffset <- IxFun.linearWithOffset srcIxFun bt_size = do
-      let num_elems = Imp.elements $ product $ map pe64 srcshape
-      srcspace <- entryMemSpace <$> lookupMemory srcmem
-      destspace <- entryMemSpace <$> lookupMemory destmem
-      sCopy
-        destmem
-        (sExt64 destoffset)
-        destspace
-        srcmem
-        (sExt64 srcoffset)
-        srcspace
-        num_elems
-        bt
-  | otherwise = sCopyKernel bt destloc srcloc
+gpuCopyForType :: Rank -> PrimType -> CallKernelGen Name
+gpuCopyForType r bt = do
+  let fname = nameFromString $ "builtin#" <> gpuCopyName r bt
 
+  exists <- hasFunction fname
+  unless exists $ emitFunction fname $ gpuCopyFunction r bt
+
+  pure fname
+
+gpuCopyName :: Rank -> PrimType -> String
+gpuCopyName (Rank r) bt = "gpu_copy_" <> show r <> "d_" <> prettyString bt
+
+gpuCopyFunction :: Rank -> PrimType -> Imp.Function Imp.HostOp
+gpuCopyFunction (Rank r) pt = do
+  let tdesc = mconcat (replicate r "[]") <> prettyString pt
+  Imp.Function Nothing [] params $
+    Imp.DebugPrint ("\n# Copy " <> tdesc) Nothing
+      <> copy_code
+      <> Imp.DebugPrint "" Nothing
+  where
+    space = Space "device"
+    memparam v = Imp.MemParam v space
+    intparam v = Imp.ScalarParam v $ IntType Int64
+
+    mkIxFun desc = do
+      let new x = newVName $ desc <> "_" <> x
+          newDim i = LMAD.LMADDim <$> new "stride" <*> new "shape" <*> pure i
+      LMAD.LMAD <$> new "offset" <*> mapM newDim [0 .. r - 1]
+
+    (params, copy_code) = do
+      flip evalState blankNameSource $ do
+        dest_mem <- newVName "destmem"
+        dest_lmad <- mkIxFun "dest"
+
+        src_mem <- newVName "srcmem"
+        src_lmad <- mkIxFun "src"
+
+        group_size <- newVName "group_size"
+        num_groups <- newVName "num_groups"
+
+        let kernel =
+              copyKernel
+                pt
+                (le64 num_groups, Left $ untyped $ le64 group_size)
+                (dest_mem, le64 <$> dest_lmad)
+                (src_mem, le64 <$> src_lmad)
+
+            dest_offset =
+              Imp.elements (le64 (LMAD.offset dest_lmad)) `Imp.withElemType` pt
+
+            src_offset =
+              Imp.elements (le64 (LMAD.offset src_lmad)) `Imp.withElemType` pt
+
+            num_bytes =
+              Imp.elements (product (le64 <$> LMAD.shape src_lmad)) `Imp.withElemType` pt
+
+            do_copy =
+              Imp.Copy
+                pt
+                dest_mem
+                dest_offset
+                (Space "device")
+                src_mem
+                src_offset
+                (Space "device")
+                num_bytes
+
+        pure
+          ( [memparam dest_mem]
+              ++ map intparam (toList dest_lmad)
+              ++ [memparam src_mem]
+              ++ map intparam (toList src_lmad),
+            Imp.DeclareScalar group_size Imp.Nonvolatile int64
+              <> Imp.DeclareScalar num_groups Imp.Nonvolatile int64
+              <> Imp.Op (Imp.GetSize group_size "copy_group_size" Imp.SizeGroup)
+              <> Imp.Op (Imp.GetSize num_groups "copy_num_groups" Imp.SizeNumGroups)
+              <> Imp.If
+                (LMAD.memcpyable (le64 <$> dest_lmad) (le64 <$> src_lmad))
+                ( Imp.DebugPrint "## Simple copy" Nothing
+                    <> do_copy
+                )
+                ( Imp.DebugPrint "## Kernel copy" Nothing
+                    <> Imp.Op (Imp.CallKernel kernel)
+                )
+          )
+
 mapTransposeForType :: PrimType -> CallKernelGen Name
 mapTransposeForType bt = do
   let fname = nameFromString $ "builtin#" <> mapTransposeName bt
@@ -480,3 +538,51 @@
 -- unfortunate code bloat, and it would be preferable if we could
 -- simply optimise the 64-bit version to make this distinction
 -- unnecessary.  Fortunately these kernels are quite small.
+
+callKernelCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp
+callKernelCopy pt destloc@(MemLoc destmem _ dest_ixfun) srcloc@(MemLoc srcmem _ src_ixfun)
+  | Just (is_transpose, (destoffset, srcoffset, num_arrays, size_x, size_y)) <-
+      isMapTransposeCopy pt destloc srcloc = do
+      fname <- mapTransposeForType pt
+      sIf
+        is_transpose
+        ( emit . Imp.Call [] fname $
+            [ Imp.MemArg destmem,
+              Imp.ExpArg $ untyped destoffset,
+              Imp.MemArg srcmem,
+              Imp.ExpArg $ untyped srcoffset,
+              Imp.ExpArg $ untyped num_arrays,
+              Imp.ExpArg $ untyped size_x,
+              Imp.ExpArg $ untyped size_y
+            ]
+        )
+        nontranspose
+  | otherwise = nontranspose
+  where
+    nontranspose = do
+      fname <- gpuCopyForType (Rank (IxFun.rank dest_ixfun)) pt
+      dest_space <- entryMemSpace <$> lookupMemory destmem
+      src_space <- entryMemSpace <$> lookupMemory srcmem
+      let dest_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD dest_ixfun
+          src_lmad = LMAD.noPermutation $ IxFun.ixfunLMAD src_ixfun
+          num_elems = Imp.elements $ product $ LMAD.shape dest_lmad
+      if dest_space == Space "device" && src_space == Space "device"
+        then
+          emit . Imp.Call [] fname $
+            [Imp.MemArg destmem]
+              ++ map (Imp.ExpArg . untyped) (toList dest_lmad)
+              ++ [Imp.MemArg srcmem]
+              ++ map (Imp.ExpArg . untyped) (toList src_lmad)
+        else -- FIXME: this assumes a linear representation!
+        -- Currently we never generate code where this is not the
+        -- case, but we might in the future.
+
+          sCopy
+            destmem
+            (Imp.elements (LMAD.offset dest_lmad) `Imp.withElemType` pt)
+            dest_space
+            srcmem
+            (Imp.elements (LMAD.offset src_lmad) `Imp.withElemType` pt)
+            src_space
+            num_elems
+            pt
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs
@@ -19,7 +19,6 @@
     sKernelThread,
     KernelAttrs (..),
     defKernelAttrs,
-    sCopyKernel,
     lvlKernelAttrs,
     allocLocal,
     kernelAlloc,
@@ -1258,7 +1257,7 @@
   v_t <- subExpType v
   case v_t of
     Prim v_t'
-      | IxFun.isLinear arr_ixfun -> pure $
+      | IxFun.isDirect arr_ixfun -> pure $
           Just $ do
             fname <- replicateForType v_t'
             emit $
@@ -1360,7 +1359,7 @@
   CallKernelGen ()
 sIota arr n x s et = do
   ArrayEntry (MemLoc arr_mem _ arr_ixfun) _ <- lookupArray arr
-  if IxFun.isLinear arr_ixfun
+  if IxFun.isDirect arr_ixfun
     then do
       fname <- iotaForType et
       emit $
@@ -1369,33 +1368,6 @@
           fname
           [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s]
     else sIotaKernel arr n x s et
-
-sCopyKernel :: CopyCompiler GPUMem HostEnv Imp.HostOp
-sCopyKernel pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do
-  -- Note that the shape of the destination and the source are
-  -- necessarily the same.
-  let shape = map pe64 srcdims
-      kernel_size = product shape
-
-  (virtualise, constants) <- simpleKernelConstants kernel_size "copy"
-
-  fname <- askFunction
-  let name =
-        keyWithEntryPoint fname $
-          nameFromString $
-            "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants)
-
-  sKernelFailureTolerant True threadOperations constants name $
-    virtualise $ \gtid -> do
-      is <- dIndexSpace' "copy_i" shape gtid
-
-      (_, destspace, destidx) <- fullyIndexArray' destloc is
-      (_, srcspace, srcidx) <- fullyIndexArray' srcloc is
-
-      sWhen (gtid .<. kernel_size) $ do
-        tmp <- tvVar <$> dPrim "tmp" pt
-        emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile
-        emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt
 
 compileThreadResult ::
   SegSpace ->
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Copy.hs b/src/Futhark/CodeGen/ImpGen/GPU/Copy.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/CodeGen/ImpGen/GPU/Copy.hs
@@ -0,0 +1,81 @@
+-- | General implementation of GPU copying, using LMAD representation.
+-- That means the dynamic performance of this kernel depends crucially
+-- on the LMAD.  In most cases we should use a more specialised kernel.
+-- Written in ImpCode so we can compile it to both CUDA and OpenCL.
+module Futhark.CodeGen.ImpGen.GPU.Copy (copyKernel) where
+
+import Control.Monad
+import Control.Monad.State
+import Data.Foldable (toList)
+import Futhark.CodeGen.ImpCode.GPU
+import Futhark.IR.Mem.LMAD qualified as LMAD
+import Futhark.IR.Prop.Reshape
+import Futhark.MonadFreshNames
+import Futhark.Util (nubOrd)
+import Futhark.Util.IntegralExp (divUp)
+import Prelude hiding (quot, rem)
+
+copyKernel ::
+  PrimType ->
+  (TExp Int64, GroupDim) ->
+  (VName, LMAD.LMAD (TExp Int64)) ->
+  (VName, LMAD.LMAD (TExp Int64)) ->
+  Kernel
+copyKernel pt (num_groups, group_dim) (dest_mem, dest_lmad) (src_mem, src_lmad) =
+  Kernel
+    { kernelBody = body,
+      kernelUses =
+        let frees =
+              nubOrd
+                ( foldMap toList dest_lmad
+                    <> foldMap toList src_lmad
+                    <> toList num_groups
+                )
+         in map (`ScalarUse` IntType Int64) frees
+              ++ map MemoryUse [dest_mem, src_mem],
+      kernelNumGroups = [untyped num_groups],
+      kernelGroupSize = [group_dim],
+      kernelName = nameFromString ("copy_" <> show rank <> "d_" <> prettyString pt),
+      kernelFailureTolerant = True,
+      kernelCheckLocalMemory = False
+    }
+  where
+    shape = LMAD.shape dest_lmad
+    rank = length shape
+
+    body = flip evalState (newNameSource 1000) $ do
+      group_id <- newVName "group_id"
+      local_id <- newVName "local_id"
+      local_size <- newVName "local_size"
+      global_id <- newVName "global_id"
+      group_iter <- newVName "group_iter"
+      let global_id_e =
+            ((le64 group_id + le64 group_iter * num_groups) * le64 local_size)
+              + le64 local_id
+
+      is <- replicateM rank $ newVName "i"
+      let is_e = map untyped (unflattenIndex shape (le64 global_id))
+          in_bounds = foldl1 (.&&.) (zipWith (.<.) (map le64 is) shape)
+
+      element <- newVName "element"
+
+      let dec v = DeclareScalar v Nonvolatile $ IntType Int64
+          src_o = Count (LMAD.index src_lmad (map le64 is))
+          dest_o = Count (LMAD.index dest_lmad (map le64 is))
+          copy_elem =
+            DeclareScalar element Nonvolatile pt
+              <> Read element src_mem src_o pt (Space "device") Nonvolatile
+              <> Write dest_mem dest_o pt (Space "device") Nonvolatile (LeafExp element pt)
+      pure $
+        foldMap dec [group_id, local_id, local_size, global_id]
+          <> Op (GetLocalId local_id 0)
+          <> Op (GetLocalSize local_size 0)
+          <> Op (GetGroupId group_id 0)
+          <> For
+            group_iter
+            (untyped (product shape `divUp` (le64 local_size * num_groups)))
+            ( SetScalar global_id (untyped global_id_e)
+                <> foldMap dec is
+                <> mconcat (zipWith SetScalar is is_e)
+                <> If in_bounds copy_elem mempty
+            )
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
@@ -41,9 +41,10 @@
   ArrayEntry arr_loc pt <- lookupArray arr
   let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc)
   sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $
-    IxFun.reshape (memLocIxFun arr_loc) $
-      map pe64 $
-        shapeDims flat_shape
+    fromMaybe (error "flattenArray") $
+      IxFun.reshape (memLocIxFun arr_loc) $
+        map pe64 $
+          shapeDims flat_shape
 
 sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName
 sliceArray start size arr = do
diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs
--- a/src/Futhark/IR/Mem.hs
+++ b/src/Futhark/IR/Mem.hs
@@ -584,7 +584,7 @@
         MemMem {} -> pure ()
         MemAcc {} -> pure ()
         MemArray _ _ _ (ArrayIn _ ixfun)
-          | IxFun.isLinear ixfun ->
+          | IxFun.isDirect ixfun ->
               pure ()
           | otherwise ->
               TC.bad . TC.TypeError $
@@ -766,7 +766,7 @@
   TC.TypeM rep ()
 matchPatToExp pat e = do
   scope <- asksScope removeScopeAliases
-  rt <- runReaderT (expReturns $ removeExpAliases e) scope
+  rt <- maybe illformed pure $ runReader (expReturns $ removeExpAliases e) scope
 
   let (ctx_ids, val_ts) = unzip $ bodyReturnsFromPat $ removePatAliases pat
       (ctx_map_ids, ctx_map_exts) = getExtMaps $ zip ctx_ids [0 .. 1]
@@ -780,6 +780,13 @@
       </> "cannot match pattern type:"
       </> indent 2 (ppTupleLines' $ map pretty val_ts)
   where
+    illformed =
+      TC.bad $
+        TC.TypeError . docText $
+          "Expression"
+            </> indent 2 (pretty e)
+            </> "cannot be assigned an index function."
+
     matches _ _ (MemPrim x) (MemPrim y) = x == y
     matches _ _ (MemMem x_space) (MemMem y_space) =
       x_space == y_space
@@ -1006,49 +1013,55 @@
 
 -- | The return information of an expression.  This can be seen as the
 -- "return type with memory annotations" of the expression.
+--
+-- This can produce Nothing, which signifies that the result is an
+-- array layout that is not expressible as an index function.
 expReturns ::
   (LocalScope rep m, Mem rep inner) =>
   Exp rep ->
-  m [ExpReturns]
+  m (Maybe [ExpReturns])
 expReturns (BasicOp (SubExp se)) =
-  pure <$> subExpReturns se
+  Just . pure <$> subExpReturns se
 expReturns (BasicOp (Opaque _ (Var v))) =
-  pure <$> varReturns v
+  Just . pure <$> varReturns v
 expReturns (BasicOp (Reshape k newshape v)) = do
   (et, _, mem, ixfun) <- arrayVarReturns v
-  pure
-    [ MemArray et (fmap Free newshape) NoUniqueness $
-        Just . ReturnsInBlock mem . existentialiseIxFun [] $
-          reshaper ixfun $
-            map pe64 (shapeDims newshape)
-    ]
+  case reshaper k ixfun $ map pe64 $ shapeDims newshape of
+    Just ixfun' ->
+      pure . Just $
+        [ MemArray et (fmap Free newshape) NoUniqueness . Just $
+            ReturnsInBlock mem (existentialiseIxFun [] ixfun')
+        ]
+    Nothing -> pure Nothing
   where
-    reshaper = case k of
-      ReshapeArbitrary -> IxFun.reshape
-      ReshapeCoerce -> IxFun.coerce
+    reshaper ReshapeArbitrary ixfun =
+      IxFun.reshape ixfun
+    reshaper ReshapeCoerce ixfun =
+      Just . IxFun.coerce ixfun
 expReturns (BasicOp (Rearrange perm v)) = do
   (et, Shape dims, mem, ixfun) <- arrayVarReturns v
   let ixfun' = IxFun.permute ixfun perm
       dims' = rearrangeShape perm dims
-  pure
-    [ MemArray et (Shape $ map Free dims') NoUniqueness $
-        Just $
-          ReturnsInBlock mem $
-            existentialiseIxFun [] ixfun'
-    ]
+  pure $
+    Just
+      [ MemArray et (Shape $ map Free dims') NoUniqueness $
+          Just $
+            ReturnsInBlock mem $
+              existentialiseIxFun [] ixfun'
+      ]
 expReturns (BasicOp (Index v slice)) = do
-  pure . varInfoToExpReturns <$> sliceInfo v slice
+  Just . pure . varInfoToExpReturns <$> sliceInfo v slice
 expReturns (BasicOp (Update _ v _ _)) =
-  pure <$> varReturns v
+  Just . pure <$> varReturns v
 expReturns (BasicOp (FlatIndex v slice)) = do
-  pure . varInfoToExpReturns <$> flatSliceInfo v slice
+  fmap (pure . varInfoToExpReturns) <$> flatSliceInfo v slice
 expReturns (BasicOp (FlatUpdate v _ _)) =
-  pure <$> varReturns v
+  Just . pure <$> varReturns v
 expReturns (BasicOp op) =
-  extReturns . staticShapes <$> basicOpType op
+  Just . extReturns . staticShapes <$> basicOpType op
 expReturns e@(DoLoop merge _ _) = do
   t <- expExtType e
-  zipWithM typeWithDec t $ map fst merge
+  Just <$> zipWithM typeWithDec t (map fst merge)
   where
     typeWithDec t p =
       case (t, paramDec p) of
@@ -1073,18 +1086,20 @@
     isMergeVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars
     mergevars = map fst merge
 expReturns (Apply _ _ ret _) =
-  pure $ map (funReturnsToExpReturns . fst) ret
+  pure $ Just $ map (funReturnsToExpReturns . fst) ret
 expReturns (Match _ _ _ (MatchDec ret _)) =
-  pure $ map bodyReturnsToExpReturns ret
+  pure $ Just $ map bodyReturnsToExpReturns ret
 expReturns (Op op) =
-  opReturns op
+  Just <$> opReturns op
 expReturns (WithAcc inputs lam) =
-  (<>)
-    <$> (concat <$> mapM inputReturns inputs)
-    <*>
-    -- XXX: this is a bit dubious because it enforces extra copies.  I
-    -- think WithAcc should perhaps have a return annotation like If.
-    pure (extReturns $ staticShapes $ drop num_accs $ lambdaReturnType lam)
+  Just
+    <$> ( (<>)
+            <$> (concat <$> mapM inputReturns inputs)
+            <*>
+            -- XXX: this is a bit dubious because it enforces extra copies.  I
+            -- think WithAcc should perhaps have a return annotation like If.
+            pure (extReturns $ staticShapes $ drop num_accs $ lambdaReturnType lam)
+        )
   where
     inputReturns (_, arrs, _) = mapM varReturns arrs
     num_accs = length inputs
@@ -1107,14 +1122,13 @@
   (Monad m, HasScope rep m, Mem rep inner) =>
   VName ->
   FlatSlice SubExp ->
-  m (MemInfo SubExp NoUniqueness MemBind)
+  m (Maybe (MemInfo SubExp NoUniqueness MemBind))
 flatSliceInfo v slice@(FlatSlice offset idxs) = do
   (et, _, mem, ixfun) <- arrayVarReturns v
   map (fmap pe64) idxs
     & FlatSlice (pe64 offset)
     & IxFun.flatSlice ixfun
-    & ArrayIn mem
-    & MemArray et (Shape (flatSliceDims slice)) NoUniqueness
+    & fmap (MemArray et (Shape (flatSliceDims slice)) NoUniqueness . ArrayIn mem)
     & pure
 
 class IsOp op => OpReturns op where
diff --git a/src/Futhark/IR/Mem/IxFun.hs b/src/Futhark/IR/Mem/IxFun.hs
--- a/src/Futhark/IR/Mem/IxFun.hs
+++ b/src/Futhark/IR/Mem/IxFun.hs
@@ -7,7 +7,6 @@
     Shape,
     LMAD (..),
     LMADDim (..),
-    Monotonicity (..),
     index,
     mkExistential,
     iota,
@@ -19,20 +18,15 @@
     flatSlice,
     rebase,
     shape,
-    lmadShape,
+    permutation,
     rank,
-    linearWithOffset,
-    rearrangeWithOffset,
     isDirect,
-    isLinear,
     substituteInIxFun,
     substituteInLMAD,
     existentialize,
     closeEnough,
     equivalent,
-    hasOneLmad,
     permuteInv,
-    conservativeFlatten,
     disjoint,
     disjoint2,
     disjoint3,
@@ -43,25 +37,27 @@
 import Control.Category
 import Control.Monad
 import Control.Monad.State
-import Data.Function (on, (&))
-import Data.List (sort, sortBy, zip4, zipWith4)
-import Data.List.NonEmpty (NonEmpty (..))
-import Data.List.NonEmpty qualified as NE
+import Data.List (sort, zip4)
 import Data.Map.Strict qualified as M
-import Data.Maybe (isJust)
 import Data.Traversable
 import Futhark.Analysis.PrimExp
 import Futhark.Analysis.PrimExp.Convert
-import Futhark.IR.Mem.LMAD
+import Futhark.IR.Mem.LMAD hiding
+  ( flatSlice,
+    index,
+    iota,
+    mkExistential,
+    permutation,
+    permute,
+    reshape,
+    shape,
+    slice,
+  )
+import Futhark.IR.Mem.LMAD qualified as LMAD
 import Futhark.IR.Prop
 import Futhark.IR.Syntax
-  ( DimIndex (..),
-    FlatDimIndex (..),
-    FlatSlice (..),
+  ( FlatSlice (..),
     Slice (..),
-    dimFix,
-    flatSliceDims,
-    flatSliceStrides,
     unitSlice,
   )
 import Futhark.IR.Syntax.Core (Ext (..))
@@ -71,8 +67,6 @@
 import Futhark.Util.Pretty
 import Prelude hiding (gcd, id, mod, (.))
 
-type Indices num = [num]
-
 -- | An index function is a mapping from a multidimensional array
 -- index space (the domain) to a one-dimensional memory index space.
 -- Essentially, it explains where the element at position @[i,j,p]@ of
@@ -80,23 +74,20 @@
 -- constitutes its memory.  For example, we can use this to
 -- distinguish row-major and column-major representations.
 --
--- An index function is represented as a sequence of 'LMAD's.
+-- An index function is represented as an LMAD.
 data IxFun num = IxFun
-  { ixfunLMADs :: NonEmpty (LMAD num),
+  { ixfunLMAD :: LMAD num,
     -- | the shape of the support array, i.e., the original array
     --   that birthed (is the start point) of this index function.
-    base :: Shape num,
-    -- | ignoring permutations, is the index function contiguous?
-    contiguous :: Bool
+    base :: Shape num
   }
   deriving (Show, Eq)
 
 instance Pretty num => Pretty (IxFun num) where
-  pretty (IxFun lmads oshp cg) =
+  pretty (IxFun lmad oshp) =
     braces . semistack $
       [ "base:" <+> brackets (commasep $ map pretty oshp),
-        "contiguous:" <+> if cg then "true" else "false",
-        "LMADs:" <+> brackets (commastack $ NE.toList $ NE.map pretty lmads)
+        "LMAD:" <+> pretty lmad
       ]
 
 instance Substitute num => Substitute (IxFun num) where
@@ -117,23 +108,8 @@
 -- It is important that the traversal order here is the same as in
 -- mkExistential.
 instance Traversable IxFun where
-  traverse f (IxFun lmads oshp cg) =
-    IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg
-
-(++@) :: [a] -> NonEmpty a -> NonEmpty a
-es ++@ (ne :| nes) = case es of
-  e : es' -> e :| es' ++ [ne] ++ nes
-  [] -> ne :| nes
-
-(@++@) :: NonEmpty a -> NonEmpty a -> NonEmpty a
-(x :| xs) @++@ (y :| ys) = x :| xs ++ [y] ++ ys
-
-setLMADPermutation :: Permutation -> LMAD num -> LMAD num
-setLMADPermutation perm lmad =
-  lmad {lmadDims = zipWith (\dim p -> dim {ldPerm = p}) (lmadDims lmad) perm}
-
-setLMADShape :: Shape num -> LMAD num -> LMAD num
-setLMADShape shp lmad = lmad {lmadDims = zipWith (\dim s -> dim {ldShape = s}) (lmadDims lmad) shp}
+  traverse f (IxFun lmad oshp) =
+    IxFun <$> traverse f lmad <*> traverse f oshp
 
 -- | Substitute a name with a PrimExp in an index function.
 substituteInIxFun ::
@@ -141,44 +117,39 @@
   M.Map a (TPrimExp t a) ->
   IxFun (TPrimExp t a) ->
   IxFun (TPrimExp t a)
-substituteInIxFun tab (IxFun lmads oshp cg) =
+substituteInIxFun tab (IxFun lmad oshp) =
   IxFun
-    (NE.map (substituteInLMAD tab) lmads)
+    (substituteInLMAD tab lmad)
     (map (TPrimExp . substituteInPrimExp tab' . untyped) oshp)
-    cg
   where
     tab' = fmap untyped tab
 
 -- | Is this is a row-major array?
 isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool
-isDirect ixfun@(IxFun (LMAD offset dims :| []) oshp True) =
+isDirect (IxFun (LMAD offset dims) oshp) =
   let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp))
-   in hasContiguousPerm ixfun
-        && length oshp == length dims
+   in length oshp == length dims
         && offset == 0
         && all
-          (\(LMADDim s n p _, m, d, se) -> s == se && n == d && p == m)
+          (\(LMADDim s n p, m, d, se) -> s == se && n == d && p == m)
           (zip4 dims [0 .. length dims - 1] oshp strides_expected)
-isDirect _ = False
 
--- | Is index function "analyzable", i.e., consists of one LMAD
-hasOneLmad :: IxFun num -> Bool
-hasOneLmad (IxFun (_ :| []) _ _) = True
-hasOneLmad _ = False
-
 -- | Does the index function have an ascending permutation?
 hasContiguousPerm :: IxFun num -> Bool
-hasContiguousPerm (IxFun (lmad :| []) _ _) =
-  let perm = lmadPermutation lmad
+hasContiguousPerm (IxFun lmad _) =
+  let perm = LMAD.permutation lmad
    in perm == sort perm
-hasContiguousPerm _ = False
 
 -- | The index space of the index function.  This is the same as the
 -- shape of arrays that the index function supports.
 shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num
-shape (IxFun (lmad :| _) _ _) =
-  permuteFwd (lmadPermutation lmad) $ lmadShapeBase lmad
+shape (IxFun lmad _) =
+  permuteFwd (LMAD.permutation lmad) $ LMAD.shapeBase lmad
 
+-- | The permutation of the first LMAD of the index function.
+permutation :: IxFun num -> Permutation
+permutation = map LMAD.ldPerm . LMAD.dims . ixfunLMAD
+
 -- | Compute the flat memory index for a complete set @inds@ of array indices
 -- and a certain element size @elem_size@.
 index ::
@@ -186,52 +157,24 @@
   IxFun num ->
   Indices num ->
   num
-index = indexFromLMADs . ixfunLMADs
-  where
-    indexFromLMADs ::
-      (IntegralExp num, Eq num) =>
-      NonEmpty (LMAD num) ->
-      Indices num ->
-      num
-    indexFromLMADs (lmad :| []) inds = indexLMAD lmad inds
-    indexFromLMADs (lmad1 :| lmad2 : lmads) inds =
-      let i_flat = indexLMAD lmad1 inds
-          new_inds = unflattenIndex (permuteFwd (lmadPermutation lmad2) $ lmadShapeBase lmad2) i_flat
-       in indexFromLMADs (lmad2 :| lmads) new_inds
-    indexLMAD ::
-      (IntegralExp num, Eq num) =>
-      LMAD num ->
-      Indices num ->
-      num
-    indexLMAD lmad@(LMAD off dims) inds =
-      let prod =
-            sum $
-              zipWith
-                flatOneDim
-                (map ldStride dims)
-                (permuteInv (lmadPermutation lmad) inds)
-       in off + prod
+index = LMAD.index . ixfunLMAD
 
 -- | iota with offset.
 iotaOffset :: IntegralExp num => num -> Shape num -> IxFun num
-iotaOffset o ns = IxFun (makeRotIota Inc o ns :| []) ns True
+iotaOffset o ns = IxFun (LMAD.iota o ns) ns
 
 -- | iota.
 iota :: IntegralExp num => Shape num -> IxFun num
 iota = iotaOffset 0
 
--- | Create a contiguous single-LMAD index function that is
--- existential in everything, with the provided permutation,
--- monotonicity, and contiguousness.
-mkExistential :: Int -> [(Int, Monotonicity)] -> Bool -> Int -> IxFun (Ext a)
-mkExistential basis_rank perm contig start =
-  IxFun (NE.singleton lmad) basis contig
+-- | Create a single-LMAD index function that is
+-- existential in everything, with the provided permutation.
+mkExistential :: Int -> [Int] -> Int -> IxFun (Ext a)
+mkExistential basis_rank perm start =
+  IxFun (LMAD.mkExistential perm start) basis
   where
     basis = take basis_rank $ map Ext [start + 1 + dims_rank * 2 ..]
     dims_rank = length perm
-    lmad = LMAD (Ext start) $ zipWith onDim perm [0 ..]
-    onDim (p, mon) i =
-      LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2)) p mon
 
 -- | Permute dimensions.
 permute ::
@@ -239,114 +182,8 @@
   IxFun num ->
   Permutation ->
   IxFun num
-permute (IxFun (lmad :| lmads) oshp cg) perm_new =
-  let perm_cur = lmadPermutation lmad
-      perm = map (perm_cur !!) perm_new
-   in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg
-
--- | Handle the case where a slice can stay within a single LMAD.
-sliceOneLMAD ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  Slice num ->
-  Maybe (IxFun num)
-sliceOneLMAD (IxFun (lmad@(LMAD _ ldims) :| lmads) oshp cg) (Slice is) = do
-  let perm = lmadPermutation lmad
-      is' = permuteInv perm is
-      cg' = cg && slicePreservesContiguous lmad (Slice is')
-  let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims
-      -- need to remove the fixed dims from the permutation
-      perm' =
-        updatePerm perm $
-          map fst $
-            filter (isJust . dimFix . snd) $
-              zip [0 .. length is' - 1] is'
-
-  pure $ IxFun (setLMADPermutation perm' lmad' :| lmads) oshp cg'
-  where
-    updatePerm ps inds = concatMap decrease ps
-      where
-        decrease p =
-          let f n i
-                | i == p = -1
-                | i > p = n
-                | n /= -1 = n + 1
-                | otherwise = n
-              d = foldl f 0 inds
-           in [p - d | d /= -1]
-
-    -- XXX: TODO: what happens to r on a negative-stride slice; is there
-    -- such a case?
-    sliceOne ::
-      (Eq num, IntegralExp num) =>
-      LMAD num ->
-      (DimIndex num, LMADDim num) ->
-      LMAD num
-    sliceOne (LMAD off dims) (DimFix i, LMADDim s _x _ _) =
-      LMAD (off + flatOneDim s i) dims
-    sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ p _) =
-      LMAD off (dims ++ [LMADDim 0 ne p Unknown])
-    sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ n _ _))
-      | dmind == unitSlice 0 n = LMAD off (dims ++ [dim])
-    sliceOne (LMAD off dims) (dmind, LMADDim s n p m)
-      | dmind == DimSlice (n - 1) n (-1) =
-          let off' = off + flatOneDim s (n - 1)
-           in LMAD off' (dims ++ [LMADDim (s * (-1)) n p (invertMonotonicity m)])
-    sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s _ p _) =
-      LMAD (off + flatOneDim s b) (dims ++ [LMADDim 0 ne p Unknown])
-    sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s _ p m) =
-      let m' = case sgn ss of
-            Just 1 -> m
-            Just (-1) -> invertMonotonicity m
-            _ -> Unknown
-       in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) ns p m'])
-
-    slicePreservesContiguous ::
-      (Eq num, IntegralExp num) =>
-      LMAD num ->
-      Slice num ->
-      Bool
-    slicePreservesContiguous (LMAD _ dims) (Slice slc) =
-      -- remove from the slice the LMAD dimensions that have stride 0.
-      -- If the LMAD was contiguous in mem, then these dims will not
-      -- influence the contiguousness of the result.
-      -- Also normalize the input slice, i.e., 0-stride and size-1
-      -- slices are rewritten as DimFixed.
-      let (dims', slc') =
-            unzip $
-              filter ((/= 0) . ldStride . fst) $
-                zip dims $
-                  map normIndex slc
-          -- Check that:
-          -- 1. a clean split point exists between Fixed and Sliced dims
-          -- 2. the outermost sliced dim has +/- 1 stride.
-          -- 3. the rest of inner sliced dims are full.
-          (_, success) =
-            foldl
-              ( \(found, res) (slcdim, LMADDim _ n _ _) ->
-                  case (slcdim, found) of
-                    (DimFix {}, True) -> (found, False)
-                    (DimFix {}, False) -> (found, res)
-                    (DimSlice _ _ ds, False) ->
-                      -- outermost sliced dim: +/-1 stride
-                      let res' = (ds == 1 || ds == -1)
-                       in (True, res && res')
-                    (DimSlice _ ne ds, True) ->
-                      -- inner sliced dim: needs to be full
-                      let res' = (n == ne) && (ds == 1 || ds == -1)
-                       in (found, res && res')
-              )
-              (False, True)
-              $ zip slc' dims'
-       in success
-
-    normIndex ::
-      (Eq num, IntegralExp num) =>
-      DimIndex num ->
-      DimIndex num
-    normIndex (DimSlice b 1 _) = DimFix b
-    normIndex (DimSlice b _ 0) = DimFix b
-    normIndex d = d
+permute (IxFun lmad oshp) perm_new =
+  IxFun (LMAD.permute lmad perm_new) oshp
 
 -- | Slice an index function.
 slice ::
@@ -354,46 +191,23 @@
   IxFun num ->
   Slice num ->
   IxFun num
-slice ixfun@(IxFun (lmad@(LMAD _ _) :| lmads) oshp cg) dim_slices
+slice ixfun@(IxFun lmad@(LMAD _ _) oshp) (Slice is)
   -- Avoid identity slicing.
-  | unSlice dim_slices == map (unitSlice 0) (shape ixfun) = ixfun
-  | Just ixfun' <- sliceOneLMAD ixfun dim_slices = ixfun'
+  | is == map (unitSlice 0) (shape ixfun) = ixfun
   | otherwise =
-      case sliceOneLMAD (iota (lmadShape lmad)) dim_slices of
-        Just (IxFun (lmad' :| []) _ cg') ->
-          IxFun (lmad' :| lmad : lmads) oshp (cg && cg')
-        _ -> error "slice: reached impossible case"
+      IxFun (LMAD.slice lmad (Slice is)) oshp
 
 -- | Flat-slice an index function.
 flatSlice ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   FlatSlice num ->
-  IxFun num
-flatSlice ixfun@(IxFun (LMAD offset (dim : dims) :| lmads) oshp cg) (FlatSlice new_offset is)
-  | hasContiguousPerm ixfun =
-      let lmad =
-            LMAD
-              (offset + new_offset * ldStride dim)
-              (map (helper $ ldStride dim) is <> dims)
-              & setLMADPermutation [0 ..]
-       in IxFun (lmad :| lmads) oshp cg
-  where
-    helper s0 (FlatDimIndex n s) =
-      let new_mon = if s0 * s == 1 then Inc else Unknown
-       in LMADDim (s0 * s) n 0 new_mon
-flatSlice (IxFun (lmad :| lmads) oshp cg) s@(FlatSlice new_offset _) =
-  IxFun (LMAD (new_offset * base_stride) (new_dims <> tail_dims) :| lmad : lmads) oshp cg
-  where
-    tail_shapes = tail $ lmadShape lmad
-    base_stride = product tail_shapes
-    tail_strides = tail $ scanr (*) 1 tail_shapes
-    tail_dims = zipWith4 LMADDim tail_strides tail_shapes [length new_shapes ..] (repeat Inc)
-    new_shapes = flatSliceDims s
-    new_strides = map (* base_stride) $ flatSliceStrides s
-    new_dims = zipWith4 LMADDim new_strides new_shapes [0 ..] (repeat Inc)
+  Maybe (IxFun num)
+flatSlice (IxFun lmad oshp) s = do
+  lmad' <- LMAD.flatSlice lmad s
+  Just $ IxFun lmad' oshp
 
--- | Handle the case where a reshape operation can stay inside a single LMAD.
+-- | Reshape an index function.
 --
 -- There are four conditions that all must hold for the result of a reshape
 -- operation to remain in the one-LMAD domain:
@@ -402,82 +216,18 @@
 --       the LMAD dimensions that were *not* reshape coercions.
 --   (2) the repetition of dimensions of the underlying LMAD must
 --       refer only to the coerced-dimensions of the reshape operation.
---   (3) finally, the underlying memory is contiguous (and monotonous).
 --
--- If any of these conditions do not hold, then the reshape operation will
--- conservatively add a new LMAD to the list, leading to a representation that
--- provides less opportunities for further analysis.
-reshapeOneLMAD ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  Shape num ->
-  Maybe (IxFun num)
-reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) oldbase cg) newshape = do
-  let perm = lmadPermutation lmad
-      dims_perm = permuteFwd perm dims
-      mid_dims = take (length dims) dims_perm
-      mon = ixfunMonotonicity ixfun
-
-  guard $
-    -- checking conditions (2)
-    all (\(LMADDim s _ _ _) -> s /= 0) mid_dims
-      &&
-      -- checking condition (1)
-      consecutive 0 (map ldPerm mid_dims)
-      &&
-      -- checking condition (3)
-      hasContiguousPerm ixfun
-      && cg
-      && (mon == Inc || mon == Dec)
-
-  -- make new permutation
-  let rsh_len = length newshape
-      diff = length newshape - length dims
-      iota_shape = [0 .. length newshape - 1]
-      perm' =
-        map
-          ( \i ->
-              let ind = i - diff
-               in if (i >= 0) && (i < rsh_len)
-                    then i -- already checked mid_dims not affected
-                    else ldPerm (dims !! ind) + diff
-          )
-          iota_shape
-      -- split the dimensions
-      (support_inds, repeat_inds) =
-        foldl
-          (\(sup, rpt) (shpdim, ip) -> ((ip, shpdim) : sup, rpt))
-          ([], [])
-          $ reverse
-          $ zip newshape perm'
-
-      (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds
-      (rpt_inds, repeats) = unzip repeat_inds
-      LMAD off' dims_sup = makeRotIota mon off support
-      repeats' = map (\n -> LMADDim 0 n 0 Unknown) repeats
-      dims' =
-        map snd $
-          sortBy (compare `on` fst) $
-            zip sup_inds dims_sup ++ zip rpt_inds repeats'
-      lmad' = LMAD off' dims'
-  pure $ IxFun (setLMADPermutation perm' lmad' :| lmads) oldbase cg
-  where
-    consecutive _ [] = True
-    consecutive i [p] = i == p
-    consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..]
-
--- | Reshape an index function.
+-- If any of these conditions do not hold, then the reshape operation
+-- will conservatively add a new LMAD to the list, leading to a
+-- representation that provides less opportunities for further
+-- analysis
 reshape ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   Shape num ->
-  IxFun num
-reshape ixfun new_shape
-  | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun'
-reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape =
-  case iota new_shape of
-    IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg
-    _ -> error "reshape: reached impossible case"
+  Maybe (IxFun num)
+reshape (IxFun lmad _) new_shape =
+  IxFun <$> LMAD.reshape lmad new_shape <*> pure new_shape
 
 -- | Coerce an index function to look like it has a new shape.
 -- Dynamically the shape must be the same.
@@ -486,18 +236,15 @@
   IxFun num ->
   Shape num ->
   IxFun num
-coerce (IxFun (lmad :| lmads) oshp cg) new_shape =
-  IxFun (onLMAD lmad :| lmads) oshp cg
+coerce (IxFun lmad _) new_shape =
+  IxFun (onLMAD lmad) new_shape
   where
     onLMAD (LMAD offset dims) = LMAD offset $ zipWith onDim dims new_shape
     onDim ld d = ld {ldShape = d}
 
 -- | The number of dimensions in the domain of the input function.
-rank ::
-  IntegralExp num =>
-  IxFun num ->
-  Int
-rank (IxFun (LMAD _ sss :| _) _ _) = length sss
+rank :: IntegralExp num => IxFun num -> Int
+rank (IxFun (LMAD _ sss) _) = length sss
 
 -- | Essentially @rebase new_base ixfun = ixfun o new_base@
 -- Core soundness condition: @base ixfun == shape new_base@
@@ -509,15 +256,6 @@
 -- We can often stay in that domain if the original ixfun is essentially a
 -- slice, e.g. `x[i, (k1,m,s1), (k2,n,s2)] = orig`.
 --
--- XXX: TODO: handle repetitions in both lmads.
---
--- How to handle repeated dimensions in the original?
---
---   (a) Shave them off of the last lmad of original
---   (b) Compose the result from (a) with the first
---       lmad of the new base
---   (c) apply a repeat operation on the result of (b).
---
 -- However, I strongly suspect that for in-place update what we need is actually
 -- the INVERSE of the rebase function, i.e., given an index function new-base
 -- and another one orig, compute the index function ixfun0 such that:
@@ -527,177 +265,67 @@
 --
 -- because then I can go bottom up and compose with ixfun0 all the index
 -- functions corresponding to the memory block associated with ixfun.
-rebaseNice ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  IxFun num ->
-  Maybe (IxFun num)
-rebaseNice
-  new_base@(IxFun (lmad_base :| lmads_base) _ cg_base)
-  ixfun@(IxFun lmads shp cg) = do
-    let (lmad :| lmads') = NE.reverse lmads
-        dims = lmadDims lmad
-        perm = lmadPermutation lmad
-        perm_base = lmadPermutation lmad_base
-
-    guard $
-      -- Core rebase condition.
-      base ixfun == shape new_base
-        -- Conservative safety conditions: ixfun is contiguous and has known
-        -- monotonicity for all dimensions.
-        && cg
-        && all ((/= Unknown) . ldMon) dims
-        -- XXX: We should be able to handle some basic cases where both index
-        -- functions have non-trivial permutations.
-        && (hasContiguousPerm ixfun || hasContiguousPerm new_base)
-        -- We need the permutations to be of the same size if we want to compose
-        -- them.  They don't have to be of the same size if the ixfun has a trivial
-        -- permutation.  Supporting this latter case allows us to rebase when ixfun
-        -- has been created by slicing with fixed dimensions.
-        && (length perm == length perm_base || hasContiguousPerm ixfun)
-        -- To not have to worry about ixfun having non-1 strides, we also check that
-        -- it is a row-major array (modulo permutation, which is handled
-        -- separately).  Accept a non-full innermost dimension.  XXX: Maybe this can
-        -- be less conservative?
-        && and
-          ( zipWith3
-              (\sn ld inner -> sn == ldShape ld || (inner && ldStride ld == 1))
-              shp
-              dims
-              (replicate (length dims - 1) False ++ [True])
-          )
-
-    -- Compose permutations, reverse strides and adjust offset if necessary.
-    let perm_base' =
-          if hasContiguousPerm ixfun
-            then perm_base
-            else map (perm !!) perm_base
-        lmad_base' = setLMADPermutation perm_base' lmad_base
-        dims_base = lmadDims lmad_base'
-        n_fewer_dims = length dims_base - length dims
-        (dims_base', offs_contrib) =
-          unzip $
-            zipWith
-              ( \(LMADDim s1 n1 p1 _) (LMADDim _ _ _ m2) ->
-                  let (s', off')
-                        | m2 == Inc = (s1, 0)
-                        | otherwise = (s1 * (-1), s1 * (n1 - 1))
-                   in (LMADDim s' n1 (p1 - n_fewer_dims) Inc, off')
-              )
-              -- If @dims@ is morally a slice, it might have fewer dimensions than
-              -- @dims_base@.  Drop extraneous outer dimensions.
-              (drop n_fewer_dims dims_base)
-              dims
-        off_base = lmadOffset lmad_base' + sum offs_contrib
-        lmad_base''
-          | lmadOffset lmad == 0 = LMAD off_base dims_base'
-          | otherwise =
-              -- If the innermost dimension of the ixfun was not full (but still
-              -- had a stride of 1), add its offset relative to the new base.
-              setLMADShape
-                (lmadShape lmad)
-                ( LMAD
-                    (off_base + ldStride (last dims_base) * lmadOffset lmad)
-                    dims_base'
-                )
-        new_base' = IxFun (lmad_base'' :| lmads_base) shp cg_base
-        IxFun lmads_base' _ _ = new_base'
-        lmads'' = lmads' ++@ lmads_base'
-    pure $ IxFun lmads'' shp (cg && cg_base)
-
--- | Rebase an index function on top of a new base.
 rebase ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   IxFun num ->
-  IxFun num
-rebase new_base@(IxFun lmads_base shp_base cg_base) ixfun@(IxFun lmads shp cg)
-  | Just ixfun' <- rebaseNice new_base ixfun = ixfun'
-  -- In the general case just concatenate LMADs since this refers to index
-  -- function composition, which is always safe.
-  | otherwise =
-      let (lmads_base', shp_base') =
-            if base ixfun == shape new_base
-              then (lmads_base, shp_base)
-              else
-                let IxFun lmads' shp_base'' _ = reshape new_base shp
-                 in (lmads', shp_base'')
-       in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base)
-
--- | If the memory support of the index function is contiguous and row-major
--- (i.e., no transpositions, repetitions, rotates, etc.), then this should
--- return the offset from which the memory-support of this index function
--- starts.
-linearWithOffset ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  num ->
-  Maybe num
-linearWithOffset ixfun@(IxFun (lmad :| []) _ cg) elem_size
-  | hasContiguousPerm ixfun && cg && ixfunMonotonicity ixfun == Inc =
-      Just $ lmadOffset lmad * elem_size
-linearWithOffset _ _ = Nothing
-
--- | Similar restrictions to @linearWithOffset@ except for transpositions, which
--- are returned together with the offset.
-rearrangeWithOffset ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  num ->
-  Maybe (num, [(Int, num)])
-rearrangeWithOffset (IxFun (lmad :| []) oshp cg) elem_size = do
-  -- Note that @cg@ describes whether the index function is
-  -- contiguous, *ignoring permutations*.  This function requires that
-  -- functionality.
-  let perm = lmadPermutation lmad
-      perm_contig = [0 .. length perm - 1]
-  offset <-
-    linearWithOffset
-      (IxFun (setLMADPermutation perm_contig lmad :| []) oshp cg)
-      elem_size
-  pure (offset, zip perm (permuteFwd perm (lmadShapeBase lmad)))
-rearrangeWithOffset _ _ = Nothing
-
--- | Is this a row-major array starting at offset zero?
-isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool
-isLinear = (== Just 0) . flip linearWithOffset 1
-
-flatOneDim ::
-  (Eq num, IntegralExp num) =>
-  num ->
-  num ->
-  num
-flatOneDim s i
-  | s == 0 = 0
-  | otherwise = i * s
+  Maybe (IxFun num)
+rebase new_base@(IxFun lmad_base _) ixfun@(IxFun lmad shp) = do
+  let dims = LMAD.dims lmad
+      perm = LMAD.permutation lmad
+      perm_base = LMAD.permutation lmad_base
 
--- | Check monotonicity of an index function.
-ixfunMonotonicity ::
-  (Eq num, IntegralExp num) =>
-  IxFun num ->
-  Monotonicity
-ixfunMonotonicity (IxFun (lmad :| lmads) _ _) =
-  let mon0 = lmadMonotonicityRots lmad
-   in if all ((== mon0) . lmadMonotonicityRots) lmads
-        then mon0
-        else Unknown
-  where
-    lmadMonotonicityRots ::
-      (Eq num, IntegralExp num) =>
-      LMAD num ->
-      Monotonicity
-    lmadMonotonicityRots (LMAD _ dims)
-      | all (isMonDim Inc) dims = Inc
-      | all (isMonDim Dec) dims = Dec
-      | otherwise = Unknown
+  guard $
+    -- Core rebase condition.
+    base ixfun == shape new_base
+      -- XXX: We should be able to handle some basic cases where both index
+      -- functions have non-trivial permutations.
+      && (hasContiguousPerm ixfun || hasContiguousPerm new_base)
+      -- We need the permutations to be of the same size if we want to compose
+      -- them.  They don't have to be of the same size if the ixfun has a trivial
+      -- permutation.  Supporting this latter case allows us to rebase when ixfun
+      -- has been created by slicing with fixed dimensions.
+      && (length perm == length perm_base || hasContiguousPerm ixfun)
+      -- To not have to worry about ixfun having non-1 strides, we also check that
+      -- it is a row-major array (modulo permutation, which is handled
+      -- separately).  Accept a non-full outermost dimension.  XXX: Maybe this can
+      -- be less conservative?
+      && and
+        ( zipWith3
+            (\sn ld inner -> inner || sn == ldShape ld)
+            shp
+            dims
+            (True : replicate (length dims - 1) False)
+        )
 
-    isMonDim ::
-      (Eq num, IntegralExp num) =>
-      Monotonicity ->
-      LMADDim num ->
-      Bool
-    isMonDim mon (LMADDim s _ _ ldmon) =
-      s == 0 || mon == ldmon
+  -- Compose permutations, reverse strides and adjust offset if necessary.
+  let perm_base' =
+        if hasContiguousPerm ixfun
+          then perm_base
+          else map (perm !!) perm_base
+      lmad_base' = LMAD.setPermutation perm_base' lmad_base
+      dims_base = LMAD.dims lmad_base'
+      n_fewer_dims = length dims_base - length dims
+      (dims_base', offs_contrib) =
+        unzip $
+          zipWith
+            ( \(LMADDim s1 n1 p1) (LMADDim {}) ->
+                let (s', off') = (s1, 0)
+                 in (LMADDim s' n1 (p1 - n_fewer_dims), off')
+            )
+            -- If @dims@ is morally a slice, it might have fewer dimensions than
+            -- @dims_base@.  Drop extraneous outer dimensions.
+            (drop n_fewer_dims dims_base)
+            dims
+      off_base = LMAD.offset lmad_base' + sum offs_contrib
+      lmad_base'' =
+        LMAD.setShape
+          (LMAD.shape lmad)
+          ( LMAD
+              (off_base + ldStride (last dims_base) * LMAD.offset lmad)
+              dims_base'
+          )
+  pure $ IxFun lmad_base'' shp
 
 -- | Turn all the leaves of the index function into 'Ext's.  We
 --  require that there's only one LMAD, that the index function is
@@ -725,31 +353,22 @@
 closeEnough :: IxFun num -> IxFun num -> Bool
 closeEnough ixf1 ixf2 =
   (length (base ixf1) == length (base ixf2))
-    && (NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2))
-    && all closeEnoughLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2))
-    -- This treats ixf1 as the "declared type" that we are matching against.
-    && (contiguous ixf1 <= contiguous ixf2)
+    && closeEnoughLMADs (ixfunLMAD ixf1) (ixfunLMAD ixf2)
   where
-    closeEnoughLMADs :: (LMAD num, LMAD num) -> Bool
-    closeEnoughLMADs (lmad1, lmad2) =
-      length (lmadDims lmad1) == length (lmadDims lmad2)
-        && map ldPerm (lmadDims lmad1)
-          == map ldPerm (lmadDims lmad2)
+    closeEnoughLMADs lmad1 lmad2 =
+      length (LMAD.dims lmad1) == length (LMAD.dims lmad2)
+        && map ldPerm (LMAD.dims lmad1) == map ldPerm (LMAD.dims lmad2)
 
 -- | Returns true if two 'IxFun's are equivalent.
 --
 -- Equivalence in this case is defined as having the same number of LMADs, with
--- each pair of LMADs matching in permutation, offsets, strides and rotations.
+-- each pair of LMADs matching in permutation, offsets, and strides.
 equivalent :: Eq num => IxFun num -> IxFun num -> Bool
 equivalent ixf1 ixf2 =
-  NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2)
-    && all equivalentLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2))
+  equivalentLMADs (ixfunLMAD ixf1) (ixfunLMAD ixf2)
   where
-    equivalentLMADs (lmad1, lmad2) =
-      length (lmadDims lmad1) == length (lmadDims lmad2)
-        && map ldPerm (lmadDims lmad1)
-          == map ldPerm (lmadDims lmad2)
-        && lmadOffset lmad1
-          == lmadOffset lmad2
-        && map ldStride (lmadDims lmad1)
-          == map ldStride (lmadDims lmad2)
+    equivalentLMADs lmad1 lmad2 =
+      length (LMAD.dims lmad1) == length (LMAD.dims lmad2)
+        && map ldPerm (LMAD.dims lmad1) == map ldPerm (LMAD.dims lmad2)
+        && LMAD.offset lmad1 == LMAD.offset lmad2
+        && map ldStride (LMAD.dims lmad1) == map ldStride (LMAD.dims lmad2)
diff --git a/src/Futhark/IR/Mem/LMAD.hs b/src/Futhark/IR/Mem/LMAD.hs
--- a/src/Futhark/IR/Mem/LMAD.hs
+++ b/src/Futhark/IR/Mem/LMAD.hs
@@ -1,40 +1,62 @@
 -- | This module contains a representation of linear-memory accessor
 -- descriptors (LMAD); see work by Zhu, Hoeflinger and David.
+--
+-- This module is designed to be used as a qualified import, as the
+-- exported names are quite generic.
 module Futhark.IR.Mem.LMAD
   ( Shape,
+    Indices,
     LMAD (..),
     LMADDim (..),
-    Monotonicity (..),
     Permutation,
-    lmadShape,
-    lmadShapeBase,
+    index,
+    slice,
+    flatSlice,
+    reshape,
+    permute,
+    shape,
+    permutation,
+    shapeBase,
+    setPermutation,
+    setShape,
     substituteInLMAD,
     permuteInv,
     permuteFwd,
-    conservativeFlatten,
     disjoint,
     disjoint2,
     disjoint3,
     dynamicEqualsLMAD,
-    lmadPermutation,
-    makeRotIota,
-    invertMonotonicity,
+    contiguous,
+    memcpyable,
+    noPermutation,
+    iota,
+    mkExistential,
   )
 where
 
 import Control.Category
 import Control.Monad
+import Data.Foldable (toList)
 import Data.Function (on, (&))
-import Data.List (elemIndex, partition, sortBy, zipWith4)
+import Data.List (elemIndex, partition, sort, sortBy)
 import Data.Map.Strict qualified as M
-import Data.Maybe (fromJust, isNothing)
+import Data.Maybe (fromJust, isJust, isNothing)
 import Data.Traversable
 import Futhark.Analysis.AlgSimplify qualified as AlgSimplify
 import Futhark.Analysis.PrimExp
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.IR.Mem.Interval
 import Futhark.IR.Prop
-import Futhark.IR.Syntax (Type)
+import Futhark.IR.Syntax
+  ( DimIndex (..),
+    Ext (..),
+    FlatDimIndex (..),
+    FlatSlice (..),
+    Slice (..),
+    Type,
+    dimFix,
+    unitSlice,
+  )
 import Futhark.IR.Syntax.Core (VName (..))
 import Futhark.Transform.Rename
 import Futhark.Transform.Substitute
@@ -46,51 +68,26 @@
 -- | The shape of an index function.
 type Shape num = [num]
 
-type Permutation = [Int]
+-- | Indices passed to an LMAD.  Must always match the rank of the LMAD.
+type Indices num = [num]
 
--- | The physical element ordering alongside a dimension, i.e. the
--- sign of the stride.
-data Monotonicity
-  = -- | Increasing.
-    Inc
-  | -- | Decreasing.
-    Dec
-  | -- | Unknown.
-    Unknown
-  deriving (Show, Eq)
+-- | A complete permutation.
+type Permutation = [Int]
 
 -- | A single dimension in an 'LMAD'.
 data LMADDim num = LMADDim
   { ldStride :: num,
     ldShape :: num,
-    ldPerm :: Int,
-    ldMon :: Monotonicity
+    ldPerm :: Int
   }
-  deriving (Show, Eq)
-
-instance Ord Monotonicity where
-  (<=) _ Inc = True
-  (<=) Unknown _ = True
-  (<=) _ Unknown = False
-  (<=) Inc Dec = False
-  (<=) _ Dec = True
-
-instance Ord num => Ord (LMADDim num) where
-  (LMADDim s1 q1 p1 m1) <= (LMADDim s2 q2 p2 m2) =
-    ([q1, s1] < [q2, s2])
-      || ( ([q1, s1] == [q2, s2])
-             && ( (p1 < p2)
-                    || ( (p1 == p2)
-                           && (m1 <= m2)
-                       )
-                )
-         )
+  deriving (Show, Eq, Ord)
 
--- | LMAD's representation consists of a general offset and for each dimension a
--- stride, number of elements (or shape), permutation, and
--- monotonicity. Note that the permutation is not strictly necessary in that the
--- permutation can be performed directly on LMAD dimensions, but then it is
--- difficult to extract the permutation back from an LMAD.
+-- | LMAD's representation consists of a general offset and for each
+-- dimension a stride, number of elements (or shape), and
+-- permutation. Note that the permutation is not strictly necessary in
+-- that the permutation can be performed directly on LMAD dimensions,
+-- but then it is difficult to extract the permutation back from an
+-- LMAD.
 --
 -- LMAD algebra is closed under composition w.r.t. operators such as
 -- permute, index and slice.  However, other operations, such as
@@ -116,22 +113,18 @@
 --    \{ ~ \sigma + i_1 * s_1 + \ldots + i_m * s_m ~ | ~ 0 \leq i_1 < n_1, \ldots, 0 \leq i_m < n_m ~ \}
 -- \]
 data LMAD num = LMAD
-  { lmadOffset :: num,
-    lmadDims :: [LMADDim num]
+  { offset :: num,
+    dims :: [LMADDim num]
   }
   deriving (Show, Eq, Ord)
 
-instance Pretty Monotonicity where
-  pretty = pretty . show
-
 instance Pretty num => Pretty (LMAD num) where
   pretty (LMAD offset dims) =
     braces . semistack $
       [ "offset:" <+> group (pretty offset),
         "strides:" <+> p ldStride,
         "shape:" <+> p ldShape,
-        "permutation:" <+> p ldPerm,
-        "monotonicity:" <+> p ldMon
+        "permutation:" <+> p ldPerm
       ]
     where
       p f = group $ brackets $ align $ commasep $ map (pretty . f) dims
@@ -146,7 +139,7 @@
   freeIn' = foldMap freeIn'
 
 instance FreeIn num => FreeIn (LMADDim num) where
-  freeIn' (LMADDim s n _ _) = freeIn' s <> freeIn' n
+  freeIn' (LMADDim s n _) = freeIn' s <> freeIn' n
 
 instance Functor LMAD where
   fmap = fmapDefault
@@ -158,16 +151,188 @@
   traverse f (LMAD offset dims) =
     LMAD <$> f offset <*> traverse f' dims
     where
-      f' (LMADDim s n p m) = LMADDim <$> f s <*> f n <*> pure p <*> pure m
+      f' (LMADDim s n p) = LMADDim <$> f s <*> f n <*> pure p
 
-invertMonotonicity :: Monotonicity -> Monotonicity
-invertMonotonicity Inc = Dec
-invertMonotonicity Dec = Inc
-invertMonotonicity Unknown = Unknown
+flatOneDim ::
+  (Eq num, IntegralExp num) =>
+  num ->
+  num ->
+  num
+flatOneDim s i
+  | s == 0 = 0
+  | otherwise = i * s
 
-lmadPermutation :: LMAD num -> Permutation
-lmadPermutation = map ldPerm . lmadDims
+index :: (IntegralExp num, Eq num) => LMAD num -> Indices num -> num
+index lmad@(LMAD off dims) inds =
+  off + sum prods
+  where
+    prods =
+      zipWith
+        flatOneDim
+        (map ldStride dims)
+        (permuteInv (permutation lmad) inds)
 
+setLMADPermutation :: Permutation -> LMAD num -> LMAD num
+setLMADPermutation perm lmad =
+  lmad {dims = zipWith (\dim p -> dim {ldPerm = p}) (dims lmad) perm}
+
+-- | Handle the case where a slice can stay within a single LMAD.
+slice ::
+  (Eq num, IntegralExp num) =>
+  LMAD num ->
+  Slice num ->
+  LMAD num
+slice lmad@(LMAD _ ldims) (Slice is) =
+  let perm = permutation lmad
+      is' = permuteInv perm is
+      lmad' = foldl sliceOne (LMAD (offset lmad) []) $ zip is' ldims
+      -- need to remove the fixed dims from the permutation
+      perm' =
+        updatePerm perm $
+          map fst $
+            filter (isJust . dimFix . snd) $
+              zip [0 .. length is' - 1] is'
+   in setLMADPermutation perm' lmad'
+  where
+    updatePerm ps inds = concatMap decrease ps
+      where
+        decrease p =
+          let f n i
+                | i == p = -1
+                | i > p = n
+                | n /= -1 = n + 1
+                | otherwise = n
+              d = foldl f 0 inds
+           in [p - d | d /= -1]
+
+    sliceOne ::
+      (Eq num, IntegralExp num) =>
+      LMAD num ->
+      (DimIndex num, LMADDim num) ->
+      LMAD num
+    sliceOne (LMAD off dims) (DimFix i, LMADDim s _x _) =
+      LMAD (off + flatOneDim s i) dims
+    sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ p) =
+      LMAD off (dims ++ [LMADDim 0 ne p])
+    sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ n _))
+      | dmind == unitSlice 0 n = LMAD off (dims ++ [dim])
+    sliceOne (LMAD off dims) (dmind, LMADDim s n p)
+      | dmind == DimSlice (n - 1) n (-1) =
+          let off' = off + flatOneDim s (n - 1)
+           in LMAD off' (dims ++ [LMADDim (s * (-1)) n p])
+    sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s _ p) =
+      LMAD (off + flatOneDim s b) (dims ++ [LMADDim 0 ne p])
+    sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s _ p) =
+      LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) ns p])
+
+hasContiguousPerm :: LMAD num -> Bool
+hasContiguousPerm lmad = perm == sort perm
+  where
+    perm = permutation lmad
+
+-- | Flat-slice an LMAD.
+flatSlice ::
+  (IntegralExp num) =>
+  LMAD num ->
+  FlatSlice num ->
+  Maybe (LMAD num)
+flatSlice lmad@(LMAD offset (dim : dims)) (FlatSlice new_offset is)
+  | hasContiguousPerm lmad =
+      Just $
+        LMAD
+          (offset + new_offset * ldStride dim)
+          (map (helper $ ldStride dim) is <> dims)
+          & setLMADPermutation [0 ..]
+  where
+    helper s0 (FlatDimIndex n s) = LMADDim (s0 * s) n 0
+flatSlice _ _ = Nothing
+
+-- | Handle the case where a reshape operation can stay inside a
+-- single LMAD.  See "Futhark.IR.Mem.IxFun.reshape" for
+-- conditions.
+reshape ::
+  (Eq num, IntegralExp num) => LMAD num -> Shape num -> Maybe (LMAD num)
+--
+-- First a special case for when we are merely injecting unit
+-- dimensions into a non-permuted LMAD.
+reshape lmad@(LMAD off dims) newshape
+  | sort (permutation lmad) == permutation lmad,
+    Just dims' <- addingVacuous 0 newshape dims =
+      Just $ LMAD off dims'
+  where
+    addingVacuous i (dnew : dnews) (dold : dolds)
+      | dnew == ldShape dold =
+          (dold {ldPerm = i} :) <$> addingVacuous (i + 1) dnews dolds
+    addingVacuous i (1 : dnews) dolds =
+      (LMADDim 0 1 i :) <$> addingVacuous (i + 1) dnews dolds
+    addingVacuous _ [] [] = Just []
+    addingVacuous _ _ _ = Nothing
+
+-- Then the general case.
+reshape lmad@(LMAD off dims) newshape = do
+  let perm = permutation lmad
+      dims_perm = permuteFwd perm dims
+      mid_dims = take (length dims) dims_perm
+
+  guard $
+    -- checking conditions (2)
+    all (\(LMADDim s _ _) -> s /= 0) mid_dims
+      &&
+      -- checking condition (1)
+      consecutive 0 (map ldPerm mid_dims)
+      &&
+      -- checking condition (3)
+      hasContiguousPerm lmad
+      && all
+        (\(ld, se) -> ldStride ld == se)
+        (zip dims (reverse $ scanl (*) 1 (reverse (tail (shape lmad)))))
+
+  -- make new permutation
+  let rsh_len = length newshape
+      diff = length newshape - length dims
+      iota_shape = [0 .. length newshape - 1]
+      perm' =
+        map
+          ( \i ->
+              let ind = i - diff
+               in if (i >= 0) && (i < rsh_len)
+                    then i -- already checked mid_dims not affected
+                    else ldPerm (dims !! ind) + diff
+          )
+          iota_shape
+      -- split the dimensions
+      (support_inds, repeat_inds) =
+        foldl
+          (\(sup, rpt) (shpdim, ip) -> ((ip, shpdim) : sup, rpt))
+          ([], [])
+          $ reverse
+          $ zip newshape perm'
+
+      (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds
+      (rpt_inds, repeats) = unzip repeat_inds
+      LMAD off' dims_sup = iota off support
+      repeats' = map (\n -> LMADDim 0 n 0) repeats
+      dims' =
+        map snd $
+          sortBy (compare `on` fst) $
+            zip sup_inds dims_sup ++ zip rpt_inds repeats'
+      lmad' = LMAD off' dims'
+  Just $ setLMADPermutation perm' lmad'
+  where
+    consecutive _ [] = True
+    consecutive i [p] = i == p
+    consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..]
+
+permutation :: LMAD num -> Permutation
+permutation = map ldPerm . dims
+
+setPermutation :: Permutation -> LMAD num -> LMAD num
+setPermutation perm lmad =
+  lmad {dims = zipWith (\dim p -> dim {ldPerm = p}) (dims lmad) perm}
+
+setShape :: Shape num -> LMAD num -> LMAD num
+setShape shp lmad = lmad {dims = zipWith (\dim s -> dim {ldShape = s}) (dims lmad) shp}
+
 -- | Substitute a name with a PrimExp in an LMAD.
 substituteInLMAD ::
   Ord a =>
@@ -175,29 +340,19 @@
   LMAD (TPrimExp t a) ->
   LMAD (TPrimExp t a)
 substituteInLMAD tab (LMAD offset dims) =
-  let offset' = sub offset
-      dims' =
-        map
-          ( \(LMADDim s n p m) ->
-              LMADDim
-                (sub s)
-                (sub n)
-                p
-                m
-          )
-          dims
-   in LMAD offset' dims'
+  LMAD (sub offset) $
+    map (\(LMADDim s n p) -> LMADDim (sub s) (sub n) p) dims
   where
     tab' = fmap untyped tab
     sub = TPrimExp . substituteInPrimExp tab' . untyped
 
 -- | Shape of an LMAD.
-lmadShape :: LMAD num -> Shape num
-lmadShape lmad = permuteInv (lmadPermutation lmad) $ lmadShapeBase lmad
+shape :: LMAD num -> Shape num
+shape lmad = permuteInv (permutation lmad) $ shapeBase lmad
 
 -- | Shape of an LMAD, ignoring permutations.
-lmadShapeBase :: LMAD num -> Shape num
-lmadShapeBase = map ldShape . lmadDims
+shapeBase :: LMAD num -> Shape num
+shapeBase = map ldShape . dims
 
 permuteFwd :: Permutation -> [a] -> [a]
 permuteFwd ps elems = map (elems !!) ps
@@ -205,28 +360,37 @@
 permuteInv :: Permutation -> [a] -> [a]
 permuteInv ps elems = map snd $ sortBy (compare `on` fst) $ zip ps elems
 
--- | Generalised iota with user-specified offset and rotates.
-makeRotIota ::
+-- | Generalised iota with user-specified offset.
+iota ::
   IntegralExp num =>
-  Monotonicity ->
   -- | Offset
   num ->
   -- | Shape
   [num] ->
   LMAD num
-makeRotIota mon off ns
-  | mon == Inc || mon == Dec =
-      let rk = length ns
-          ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns
-          ss =
-            if mon == Inc
-              then ss0
-              else map (* (-1)) ss0
-          ps = map fromIntegral [0 .. rk - 1]
-          fi = replicate rk mon
-       in LMAD off $ zipWith4 LMADDim ss ns ps fi
-  | otherwise = error "makeRotIota: requires Inc or Dec"
+iota off ns =
+  let rk = length ns
+      ss = reverse $ take rk $ scanl (*) 1 $ reverse ns
+      ps = map fromIntegral [0 .. rk - 1]
+   in LMAD off $ zipWith3 LMADDim ss ns ps
 
+-- | Create an LMAD that is existential in everything, with the
+-- provided permutation.
+mkExistential :: [Int] -> Int -> LMAD (Ext a)
+mkExistential perm start =
+  lmad
+  where
+    lmad = LMAD (Ext start) $ zipWith onDim perm [0 ..]
+    onDim p i =
+      LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2)) p
+
+-- | Permute dimensions.
+permute :: LMAD num -> Permutation -> LMAD num
+permute lmad perm_new =
+  let perm_cur = permutation lmad
+      perm = map (perm_cur !!) perm_new
+   in setPermutation perm lmad
+
 -- | Computes the maximum span of an 'LMAD'. The result is the lowest and
 -- highest flat values representable by that 'LMAD'.
 flatSpan :: LMAD (TPrimExp Int64 VName) -> TPrimExp Int64 VName
@@ -248,7 +412,7 @@
 -- conservativeFlatten :: (IntegralExp e, Ord e, Pretty e) => LMAD e -> LMAD e
 conservativeFlatten :: LMAD (TPrimExp Int64 VName) -> Maybe (LMAD (TPrimExp Int64 VName))
 conservativeFlatten (LMAD offset []) =
-  pure $ LMAD offset [LMADDim 1 1 0 Inc]
+  pure $ LMAD offset [LMADDim 1 1 0]
 conservativeFlatten l@(LMAD _ [_]) =
   pure l
 conservativeFlatten l@(LMAD offset dims) = do
@@ -257,7 +421,7 @@
       gcd
       (ldStride $ head dims)
       $ map ldStride dims
-  pure $ LMAD offset [LMADDim strd (shp + 1) 0 Unknown]
+  pure $ LMAD offset [LMADDim strd (shp + 1) 0]
   where
     shp = flatSpan l
 
@@ -440,12 +604,12 @@
 lmadToIntervals :: LMAD (TPrimExp Int64 VName) -> (AlgSimplify.SofP, [Interval])
 lmadToIntervals (LMAD offset []) = (AlgSimplify.simplify0 $ untyped offset, [Interval 0 1 1])
 lmadToIntervals lmad@(LMAD offset dims0) =
-  (offset', map helper $ permuteInv (lmadPermutation lmad) dims0)
+  (offset', map helper $ permuteInv (permutation lmad) dims0)
   where
     offset' = AlgSimplify.simplify0 $ untyped offset
 
     helper :: LMADDim (TPrimExp Int64 VName) -> Interval
-    helper (LMADDim strd shp _ _) = do
+    helper (LMADDim strd shp _) = do
       Interval 0 (AlgSimplify.simplify' shp) (AlgSimplify.simplify' strd)
 
 -- | Dynamically determine if two 'LMADDim' are equal.
@@ -456,15 +620,52 @@
   ldStride dim1 .==. ldStride dim2
     .&&. ldShape dim1 .==. ldShape dim2
     .&&. fromBool (ldPerm dim1 == ldPerm dim2)
-    .&&. fromBool (ldMon dim1 == ldMon dim2)
 
 -- | Dynamically determine if two 'LMAD' are equal.
 --
 -- True if offset and constituent 'LMADDim' are equal.
 dynamicEqualsLMAD :: Eq num => LMAD (TPrimExp t num) -> LMAD (TPrimExp t num) -> TPrimExp Bool num
 dynamicEqualsLMAD lmad1 lmad2 =
-  lmadOffset lmad1 .==. lmadOffset lmad2
+  offset lmad1 .==. offset lmad2
     .&&. foldr
       ((.&&.) . uncurry dynamicEqualsLMADDim)
       true
-      (zip (lmadDims lmad1) (lmadDims lmad2))
+      (zip (dims lmad1) (dims lmad2))
+
+-- | True if these LMADs represent the same function (ignoring
+-- offset).
+compatible ::
+  Eq num =>
+  LMAD (TPrimExp Int64 num) ->
+  LMAD (TPrimExp Int64 num) ->
+  TPrimExp Bool num
+compatible x y =
+  foldl1 (.&&.) $ zipWith dynamicEqualsLMADDim (dims x) (dims y)
+
+-- | True if this LMAD corresponds to an array without "holes".  This
+-- implies it can be copied with a memcpy()-like operation.
+contiguous ::
+  (Pretty num, Eq num) =>
+  LMAD (TPrimExp Int64 num) ->
+  TPrimExp Bool num
+contiguous lmad =
+  foldl1 (.&&.) $ zipWith (.==.) (toList lmad) lmad'
+  where
+    lmad' = toList (iota (offset lmad) $ map ldShape $ dims lmad)
+
+-- | True if these LMADs have the same contiguous representation, such
+-- that one can be copied to the other with a @memcpy()@-like
+-- operation.
+memcpyable ::
+  (Pretty num, Eq num) =>
+  LMAD (TPrimExp Int64 num) ->
+  LMAD (TPrimExp Int64 num) ->
+  TPrimExp Bool num
+memcpyable dest_lmad src_lmad =
+  contiguous dest_lmad .&&. compatible dest_lmad src_lmad
+
+-- | Remove the permutation of an LMAD by actually applying it to the
+-- dimensions.
+noPermutation :: LMAD t -> LMAD t
+noPermutation lmad =
+  lmad {dims = rearrangeShape (permutation lmad) $ dims lmad}
diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs
--- a/src/Futhark/IR/Parse.hs
+++ b/src/Futhark/IR/Parse.hs
@@ -17,9 +17,8 @@
 
 import Data.Char (isAlpha)
 import Data.Functor
-import Data.List (singleton, zipWith4)
+import Data.List (singleton)
 import Data.List.NonEmpty (NonEmpty (..))
-import Data.List.NonEmpty qualified as NE
 import Data.Maybe
 import Data.Set qualified as S
 import Data.Text qualified as T
@@ -68,9 +67,6 @@
     pTag = "_" *> L.decimal <* notFollowedBy (satisfy constituent)
     exprBox = ("<{" <>) . (<> "}>") <$> (chunk "<{" *> manyTill anySingle (chunk "}>"))
 
-pBool :: Parser Bool
-pBool = choice [keyword "true" $> True, keyword "false" $> False]
-
 pInt :: Parser Int
 pInt = lexeme L.decimal
 
@@ -975,24 +971,16 @@
 pIxFunBase pNum =
   braces $ do
     base <- pLab "base" $ brackets (pNum `sepBy` pComma) <* pSemi
-    ct <- pLab "contiguous" $ pBool <* pSemi
-    lmads <- pLab "LMADs" $ brackets (pLMAD `sepBy1` pComma)
-    pure $ IxFun.IxFun (NE.fromList lmads) base ct
+    lmad <- pLab "LMAD" pLMAD
+    pure $ IxFun.IxFun lmad base
   where
     pLab s m = keyword s *> pColon *> m
-    pMon =
-      choice
-        [ "Inc" $> IxFun.Inc,
-          "Dec" $> IxFun.Dec,
-          "Unknown" $> IxFun.Unknown
-        ]
     pLMAD = braces $ do
       offset <- pLab "offset" pNum <* pSemi
       strides <- pLab "strides" $ brackets (pNum `sepBy` pComma) <* pSemi
       shape <- pLab "shape" $ brackets (pNum `sepBy` pComma) <* pSemi
-      perm <- pLab "permutation" $ brackets (pInt `sepBy` pComma) <* pSemi
-      mon <- pLab "monotonicity" $ brackets (pMon `sepBy` pComma)
-      pure $ IxFun.LMAD offset $ zipWith4 IxFun.LMADDim strides shape perm mon
+      perm <- pLab "permutation" $ brackets (pInt `sepBy` pComma)
+      pure $ IxFun.LMAD offset $ zipWith3 IxFun.LMADDim strides shape perm
 
 pPrimExpLeaf :: Parser VName
 pPrimExpLeaf = pVName
diff --git a/src/Futhark/Internalise/ReplaceRecords.hs b/src/Futhark/Internalise/ReplaceRecords.hs
--- a/src/Futhark/Internalise/ReplaceRecords.hs
+++ b/src/Futhark/Internalise/ReplaceRecords.hs
@@ -4,6 +4,7 @@
 -- around huge records of which only a tiny part is needed.
 module Futhark.Internalise.ReplaceRecords (transformProg) where
 
+import Control.Monad
 import Control.Monad.Reader
 import Control.Monad.State
 import Data.Map.Strict qualified as M
diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs
--- a/src/Futhark/Internalise/TypesValues.hs
+++ b/src/Futhark/Internalise/TypesValues.hs
@@ -21,6 +21,7 @@
   )
 where
 
+import Control.Monad
 import Control.Monad.Free (Free (..))
 import Control.Monad.State
 import Data.Bitraversable (bitraverse)
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
@@ -89,8 +89,8 @@
 -- | Given an 'IORef' tracking the state, produce a set of handlers.
 -- When we want to add more features to the language server, this is
 -- the thing to change.
-handlers :: IORef State -> Handlers (LspM ())
-handlers state_mvar =
+handlers :: IORef State -> ClientCapabilities -> Handlers (LspM ())
+handlers state_mvar _ =
   mconcat
     [ onInitializeHandler,
       onDocumentOpenHandler state_mvar,
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs
@@ -17,7 +17,6 @@
 import Data.Function ((&))
 import Data.List qualified as L
 import Data.List.NonEmpty (NonEmpty (..))
-import Data.List.NonEmpty qualified as NE
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Data.Sequence (Seq (..))
@@ -421,8 +420,10 @@
                 (aggSummaryMapPartial (scalarTable td_env) $ unSegSpace space0)
                 (S.toList s)
             Undeterminable -> pure Undeterminable
-        let res = noMemOverlap td_env destination_uses source_writes
-        if res
+        -- Do not allow short-circuiting from a segop-local memory
+        -- block (not in the topdown scope) to an outer memory block.
+        if dstmem entry `M.member` scope td_env
+          && noMemOverlap td_env destination_uses source_writes
           then pure bu_env_f
           else do
             let (ac, inh) = markFailedCoal (activeCoals bu_env_f, inhibit bu_env_f) k
@@ -484,8 +485,8 @@
                     Just (Coalesced knd mbd@(MemBlock _ _ _ ixfn) _) -> pure $
                       case freeVarSubstitutions (scope td_env) (scalarTable td_env) ixfn of
                         Just fv_subst ->
-                          if ixfunPermutation ixfn
-                            == ixfunPermutation (ixfun $ fromJust $ getScopeMemInfo (patElemName p) $ scope td_env)
+                          if IxFun.permutation ixfn
+                            == IxFun.permutation (ixfun $ fromJust $ getScopeMemInfo (patElemName p) $ scope td_env)
                             then
                               let entry =
                                     coal_entry
@@ -506,9 +507,6 @@
 
   foldM checkMergeeOverlap bu_env''' mergee_writes
 
-ixfunPermutation :: IxFun -> [Int]
-ixfunPermutation = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs
-
 -- | Given a pattern element and the corresponding kernel result, try to put the
 -- kernel result directly in the memory block of pattern element
 makeSegMapCoals ::
@@ -534,8 +532,7 @@
       case M.lookup pat_mem active of
         Nothing ->
           -- We are not in a transitive case
-          case ( IxFun.hasOneLmad pat_ixf
-                   && maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),
+          case ( maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),
                  Coalesced
                    InPlaceCoal
                    (MemBlock tp return_shp pat_mem $ resultSlice pat_ixf)
@@ -1259,8 +1256,7 @@
                         _ -> (failed, s_acc) -- fail!
 
 ixfunToAccessSummary :: IxFun.IxFun (TPrimExp Int64 VName) -> AccessSummary
-ixfunToAccessSummary (IxFun.IxFun (lmad NE.:| []) _ _) = Set $ S.singleton lmad
-ixfunToAccessSummary _ = Undeterminable
+ixfunToAccessSummary (IxFun.IxFun lmad _) = Set $ S.singleton lmad
 
 -- | Check safety conditions 2 and 5 and update new substitutions:
 -- called on the pat-elements of loop and if-then-else expressions.
@@ -1365,9 +1361,8 @@
                       Nothing ->
                         ind_y
                  in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0, certs <> certs'')
-          success0 = IxFun.hasOneLmad ind_yx
           m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx
-       in if success0 && not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)
+       in if not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)
       -- Finally update the @activeCoals@ table with a fresh
       --   binding for @m_b@; if such one exists then overwrite.
       -- Also, add all variables from the alias chain of @b@ to
@@ -1562,7 +1557,10 @@
   where
     updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun
     updateIndFunSlice ind_fun (FlatSlice offset dims) =
-      IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims
+      fromMaybe (error "updateIndFunSlice") $
+        IxFun.flatSlice ind_fun $
+          FlatSlice (pe64 offset) $
+            map (fmap pe64) dims
 
 -- CASE b) @let x = concat(a, b^{lu})@
 genCoalStmtInfo lutab _ scopetab (Let pat aux (BasicOp (Concat concat_dim (b0 :| bs) _)))
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs
@@ -134,12 +134,12 @@
         Constant _ -> Just (ws, ws)
         Var x -> Just (ws, ws ++ mapMaybe (getDirAliasedIxfn td_env coal_tab) [x])
 getUseSumFromStm td_env coal_tab (Let (Pat [x]) _ (BasicOp (FlatUpdate _ (FlatSlice offset slc) v)))
-  | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) =
-      let x_ixfn_slc = IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc
-          r1 = (m_b, m_x, x_ixfn_slc)
-       in case getDirAliasedIxfn td_env coal_tab v of
-            Nothing -> Just ([r1], [r1])
-            Just r2 -> Just ([r1], [r1, r2])
+  | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) = do
+      x_ixfn_slc <- IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc
+      let r1 = (m_b, m_x, x_ixfn_slc)
+      case getDirAliasedIxfn td_env coal_tab v of
+        Nothing -> Just ([r1], [r1])
+        Just r2 -> Just ([r1], [r1, r2])
 -- getUseSumFromStm td_env coal_tab (Let (Pat ys) _ (BasicOp bop)) =
 --   let wrt = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) ys
 --    in trace ("getUseBla: " <> show bop) $ pure (wrt, wrt)
@@ -251,7 +251,7 @@
         <> fromMaybe mempty (M.lookup m (m_alias td_env))
     mbLmad indfun
       | Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun,
-        (IxFun.IxFun (lmad :| []) _ _) <- IxFun.substituteInIxFun subs indfun =
+        (IxFun.IxFun lmad _) <- IxFun.substituteInIxFun subs indfun =
           Just lmad
     mbLmad _ = Nothing
     addLmads wrts uses etry =
@@ -458,7 +458,7 @@
           new_offset = replaceIteratorWith lower_bound offset0
           new_lmad =
             IxFun.LMAD new_offset $
-              IxFun.LMADDim new_stride spn 0 IxFun.Inc : map incPerm dims0
+              IxFun.LMADDim new_stride spn 0 : map incPerm dims0
       if new_var `nameIn` freeIn new_lmad
         then pure Undeterminable
         else pure $ Set $ S.singleton new_lmad
diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
--- a/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
+++ b/src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs
@@ -27,7 +27,7 @@
 import Futhark.IR.Mem.IxFun qualified as IxFun
 import Futhark.Optimise.ArrayShortCircuiting.DataStructs
 
-type DirAlias = IxFun -> IxFun
+type DirAlias = IxFun -> Maybe IxFun
 -- ^ A direct aliasing transformation
 
 type InvAlias = Maybe (IxFun -> IxFun)
@@ -69,23 +69,23 @@
 
 -- | Get alias and (direct) index function mapping from expression
 getDirAliasFromExp :: Exp (Aliases rep) -> Maybe (VName, DirAlias)
-getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, id)
-getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, id)
+getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, Just)
+getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, Just)
 getDirAliasFromExp (BasicOp (Reshape ReshapeCoerce shp x)) =
-  Just (x, (`IxFun.coerce` shapeDims (fmap pe64 shp)))
+  Just (x, Just . (`IxFun.coerce` shapeDims (fmap pe64 shp)))
 getDirAliasFromExp (BasicOp (Reshape ReshapeArbitrary shp x)) =
   Just (x, (`IxFun.reshape` shapeDims (fmap pe64 shp)))
 getDirAliasFromExp (BasicOp (Rearrange _ _)) =
   Nothing
 getDirAliasFromExp (BasicOp (Index x slc)) =
-  Just (x, (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))
-getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, id)
+  Just (x, Just . (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))
+getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, Just)
 getDirAliasFromExp (BasicOp (FlatIndex x (FlatSlice offset idxs))) =
   Just
     ( x,
       (`IxFun.flatSlice` FlatSlice (pe64 offset) (map (fmap pe64) idxs))
     )
-getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, id)
+getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, Just)
 getDirAliasFromExp _ = Nothing
 
 -- | This was former @createsAliasedArrOK@ from DataStructs
@@ -268,7 +268,8 @@
 walkAliasTab alias_tab vtab x
   | Just (x0, alias0, _) <- M.lookup x alias_tab = do
       Coalesced knd (MemBlock pt shp vname ixf) substs <- walkAliasTab alias_tab vtab x0
-      pure $ Coalesced knd (MemBlock pt shp vname $ alias0 ixf) substs
+      ixf' <- alias0 ixf
+      pure $ Coalesced knd (MemBlock pt shp vname ixf') substs
 walkAliasTab _ _ _ = Nothing
 
 -- | We assume @x@ is in @vartab@ and we add the variables that @x@ aliases
diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs
--- a/src/Futhark/Optimise/BlkRegTiling.hs
+++ b/src/Futhark/Optimise/BlkRegTiling.hs
@@ -20,12 +20,12 @@
 
 import Control.Monad
 import Data.List qualified as L
-import Data.List.NonEmpty qualified as NE
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Data.Sequence qualified as Seq
 import Futhark.IR.GPU
 import Futhark.IR.Mem.IxFun qualified as IxFun
+import Futhark.IR.Mem.LMAD qualified as LMAD
 import Futhark.MonadFreshNames
 import Futhark.Optimise.TileLoops.Shared
 import Futhark.Tools
@@ -143,12 +143,12 @@
         | [slc_X'] <- patNames pat,
           slc_X == slc_X',
           Just ixf_fn <- M.lookup x ixfn_env,
-          (IxFun.IxFun lmads _ _) <- ixf_fn =
-            all innerHasStride1 $ NE.toList lmads
+          (IxFun.IxFun lmad _) <- ixf_fn =
+            innerHasStride1 lmad
       isInnerCoal _ _ _ =
         error "kkLoopBody.isInnerCoal: not an error, but I would like to know why!"
       innerHasStride1 lmad =
-        let lmad_dims = IxFun.lmadDims lmad
+        let lmad_dims = LMAD.dims lmad
             q = length lmad_dims
             last_perm = IxFun.ldPerm $ last lmad_dims
             stride = IxFun.ldStride $ last lmad_dims
diff --git a/src/Futhark/Optimise/Simplify/Rules/Simple.hs b/src/Futhark/Optimise/Simplify/Rules/Simple.hs
--- a/src/Futhark/Optimise/Simplify/Rules/Simple.hs
+++ b/src/Futhark/Optimise/Simplify/Rules/Simple.hs
@@ -68,7 +68,7 @@
     op1 == op2,
     Just res <- doBinOp op1 x1 x2 =
       Just (BinOp op1 (Constant res) y2, cs)
-simplifyBinOp look _ (BinOp Add {} e1 e2)
+simplifyBinOp look _ (BinOp (Add it ovf) e1 e2)
   | isCt0 e1 = resIsSubExp e2
   | isCt0 e2 = resIsSubExp e1
   -- x+(y-x) => y
@@ -76,13 +76,20 @@
     Just (BasicOp (BinOp Sub {} e2_a e2_b), cs) <- look v2,
     e2_b == e1 =
       Just (SubExp e2_a, cs)
+  -- x+(-1*y) => x-y
+  | Var v2 <- e2,
+    Just (BasicOp (BinOp Mul {} (Constant (IntValue x)) e3), cs) <- look v2,
+    valueIntegral x == (-1 :: Int) =
+      Just (BinOp (Sub it ovf) e1 e3, cs)
 simplifyBinOp _ _ (BinOp FAdd {} e1 e2)
   | isCt0 e1 = resIsSubExp e2
   | isCt0 e2 = resIsSubExp e1
 simplifyBinOp look _ (BinOp sub@(Sub t _) e1 e2)
   | isCt0 e2 = resIsSubExp e1
-  -- Cases for simplifying (a+b)-b and permutations.
-
+  | e1 == e2 = Just (SubExp (intConst t 0), mempty)
+  --
+  -- Below are cases for simplifying (a+b)-b and permutations.
+  --
   -- (e1_a+e1_b)-e1_a == e1_b
   | Var v1 <- e1,
     Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1,
diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs
--- a/src/Futhark/Optimise/TileLoops/Shared.hs
+++ b/src/Futhark/Optimise/TileLoops/Shared.hs
@@ -39,11 +39,7 @@
 index :: MonadBuilder m => String -> VName -> [VName] -> m VName
 index se_desc arr outer_indices = do
   arr_t <- lookupType arr
-  let shape = arrayShape arr_t
-      inner_dims = shapeDims $ stripDims (length outer_indices) shape
-      untouched d = DimSlice (intConst Int64 0) d (intConst Int64 1)
-      inner_slices = map untouched inner_dims
-      slice = Slice $ map (DimFix . Var) outer_indices ++ inner_slices
+  let slice = fullSlice arr_t $ map (DimFix . Var) outer_indices
   letExp se_desc $ BasicOp $ Index arr slice
 
 update :: MonadBuilder m => String -> VName -> [VName] -> SubExp -> m VName
@@ -295,23 +291,23 @@
        in (lam_op, ne)
 changeWithEnv with_env _ = pure with_env
 
-composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> IxFun) -> TileM IxFnEnv
+composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> Maybe IxFun) -> TileM IxFnEnv
 composeIxfuns env y x ixf_fun =
-  case M.lookup x env of
-    Just ixf -> pure $ M.insert y (ixf_fun ixf) env
+  case ixf_fun =<< M.lookup x env of
+    Just ixf -> pure $ M.insert y ixf env
     Nothing -> do
       tp <- lookupType x
-      case tp of
-        Array _ptp shp _u -> do
-          let shp' = map ExpMem.pe64 (shapeDims shp)
-          pure $ M.insert y (ixf_fun $ IxFun.iota shp') env
-        _ -> pure env
+      pure $ case tp of
+        Array _ptp shp _u
+          | Just ixf <- ixf_fun $ IxFun.iota $ map ExpMem.pe64 (shapeDims shp) ->
+              M.insert y ixf env
+        _ -> env
 
 changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv
 changeIxFnEnv env y (BasicOp (Reshape ReshapeArbitrary shp_chg x)) =
   composeIxfuns env y x (`IxFun.reshape` fmap ExpMem.pe64 (shapeDims shp_chg))
 changeIxFnEnv env y (BasicOp (Reshape ReshapeCoerce shp_chg x)) =
-  composeIxfuns env y x (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg))
+  composeIxfuns env y x (Just . (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg)))
 changeIxFnEnv env y (BasicOp (Manifest perm x)) = do
   tp <- lookupType x
   case tp of
@@ -321,9 +317,9 @@
       pure $ M.insert y ixfn env
     _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!"
 changeIxFnEnv env y (BasicOp (Rearrange perm x)) =
-  composeIxfuns env y x (`IxFun.permute` perm)
+  composeIxfuns env y x (Just . (`IxFun.permute` perm))
 changeIxFnEnv env y (BasicOp (Index x slc)) =
-  composeIxfuns env y x (`IxFun.slice` (Slice $ map (fmap ExpMem.pe64) $ unSlice slc))
+  composeIxfuns env y x (Just . (`IxFun.slice` Slice (map (fmap ExpMem.pe64) $ unSlice slc)))
 changeIxFnEnv env y (BasicOp (Opaque _ (Var x))) =
-  composeIxfuns env y x id
+  composeIxfuns env y x Just
 changeIxFnEnv env _ _ = pure env
diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs
--- a/src/Futhark/Pass/ExpandAllocations.hs
+++ b/src/Futhark/Pass/ExpandAllocations.hs
@@ -30,7 +30,6 @@
 import Futhark.Transform.Rename (renameStm)
 import Futhark.Transform.Substitute
 import Futhark.Util (mapAccumLM)
-import Futhark.Util.IntegralExp
 import Prelude hiding (quot)
 
 -- | The memory expansion pass definition.
@@ -538,28 +537,26 @@
 
   pure (slice_stms' <> stmsFromList alloc_stms, mconcat rebases)
   where
-    expand (mem, (offset, total_size, space)) = do
+    expand (mem, (_offset, total_size, space)) = do
       let allocpat = Pat [PatElem mem $ MemMem space]
       pure
         ( Let allocpat (defAux ()) $ Op $ Alloc total_size space,
-          M.singleton mem $ newBase offset
+          M.singleton mem newBase
         )
 
     num_threads' = pe64 num_threads
     gtid = le64 $ segFlat kspace
 
+    untouched d = DimSlice 0 d 1
+
     -- For the variant allocations, we add an inner dimension,
     -- which is then offset by a thread-specific amount.
-    newBase size_per_thread (old_shape, pt) =
-      let elems_per_thread =
-            pe64 size_per_thread `quot` primByteSize pt
-          root_ixfun = IxFun.iota [elems_per_thread, num_threads']
+    newBase (old_shape, _pt) =
+      let root_ixfun = IxFun.iota $ old_shape ++ [num_threads']
           offset_ixfun =
             IxFun.slice root_ixfun . Slice $
-              [DimSlice 0 num_threads' 1, DimFix gtid]
-       in if length old_shape == 1
-            then IxFun.coerce offset_ixfun old_shape
-            else IxFun.reshape offset_ixfun old_shape
+              map untouched old_shape ++ [DimFix gtid]
+       in offset_ixfun
 
 -- | A map from memory block names to new index function bases.
 type RebaseMap = M.Map VName (([TPrimExp Int64 VName], PrimType) -> IxFun)
@@ -622,11 +619,16 @@
 offsetMemoryInStm :: Stm GPUMem -> OffsetM (Scope GPUMem, Stm GPUMem)
 offsetMemoryInStm (Let pat dec e) = do
   e' <- offsetMemoryInExp e
-  pat' <- offsetMemoryInPat pat =<< expReturns e'
+  pat' <-
+    offsetMemoryInPat pat
+      =<< maybe (throwError "offsetMemoryInStm: ill-typed") pure
+      =<< expReturns e'
   scope <- askScope
   -- Try to recompute the index function.  Fall back to creating rebase
   -- operations with the RebaseMap.
-  rts <- runReaderT (expReturns e') scope
+  rts <-
+    maybe (throwError "offsetMemoryInStm: ill-typed") pure $
+      runReader (expReturns e') scope
   let pat'' = Pat $ zipWith pick (patElems pat') rts
       stm = Let pat'' dec e'
   let scope' = scopeOf stm <> scope
@@ -659,32 +661,48 @@
         pure . PatElem name . MemArray pt shape u . ArrayIn mem $
           fmap (fmap unExt) ixfun
     onPE pe _ = do
-      new_dec <- offsetMemoryInMemBound $ patElemDec pe
+      new_dec <- offsetMemoryInMemBound (patElemName pe) $ patElemDec pe
       pure pe {patElemDec = new_dec}
     unExt (Ext i) = patElemName (pes !! i)
     unExt (Free v) = v
 
 offsetMemoryInParam :: Param (MemBound u) -> OffsetM (Param (MemBound u))
 offsetMemoryInParam fparam = do
-  fparam' <- offsetMemoryInMemBound $ paramDec fparam
+  fparam' <- offsetMemoryInMemBound (paramName fparam) $ paramDec fparam
   pure fparam {paramDec = fparam'}
 
-offsetMemoryInMemBound :: MemBound u -> OffsetM (MemBound u)
-offsetMemoryInMemBound summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do
+offsetMemoryInMemBound :: VName -> MemBound u -> OffsetM (MemBound u)
+offsetMemoryInMemBound v summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do
   new_base <- lookupNewBase mem (IxFun.base ixfun, pt)
-  pure . fromMaybe summary $ do
-    new_base' <- new_base
-    pure $ MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base' ixfun
-offsetMemoryInMemBound summary = pure summary
+  case new_base of
+    Nothing -> pure summary
+    Just new_base' -> do
+      let problem =
+            throwError . unlines $
+              [ "offsetMemoryInMemBound",
+                prettyString v,
+                prettyString new_base',
+                prettyString ixfun
+              ]
+      ixfun' <- maybe problem pure $ IxFun.rebase new_base' ixfun
+      pure $ MemArray pt shape u $ ArrayIn mem ixfun'
+offsetMemoryInMemBound _ summary = pure summary
 
 offsetMemoryInBodyReturns :: BodyReturns -> OffsetM BodyReturns
 offsetMemoryInBodyReturns br@(MemArray pt shape u (ReturnsInBlock mem ixfun))
   | Just ixfun' <- isStaticIxFun ixfun = do
       new_base <- lookupNewBase mem (IxFun.base ixfun', pt)
-      pure . fromMaybe br $ do
-        new_base' <- new_base
-        pure . MemArray pt shape u . ReturnsInBlock mem $
-          IxFun.rebase (fmap (fmap Free) new_base') ixfun
+      case new_base of
+        Nothing -> pure br
+        Just new_base' -> do
+          let problem =
+                throwError . unlines $
+                  [ "offsetMemoryInBodyReturns",
+                    prettyString new_base',
+                    prettyString ixfun
+                  ]
+          ixfun'' <- maybe problem pure $ IxFun.rebase (fmap (fmap Free) new_base') ixfun
+          pure $ MemArray pt shape u $ ReturnsInBlock mem ixfun''
 offsetMemoryInBodyReturns br = pure br
 
 offsetMemoryInLambda :: Lambda GPUMem -> OffsetM (Lambda GPUMem)
diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs
--- a/src/Futhark/Pass/ExplicitAllocations.hs
+++ b/src/Futhark/Pass/ExplicitAllocations.hs
@@ -43,7 +43,6 @@
 import Data.Either (partitionEithers)
 import Data.Foldable (toList)
 import Data.List (foldl', transpose, zip4)
-import Data.List.NonEmpty qualified as NE
 import Data.Map.Strict qualified as M
 import Data.Maybe
 import Data.Set qualified as S
@@ -186,6 +185,40 @@
 allocForArray t space = do
   allocForArray' t space
 
+-- | Repair an expression that cannot be assigned an index function.
+-- There is a simple remedy for this: normalise the input arrays and
+-- try again.
+repairExpression ::
+  (Allocable fromrep torep inner) =>
+  Exp torep ->
+  AllocM fromrep torep (Exp torep)
+repairExpression (BasicOp (Reshape k shape v)) = do
+  v_mem <- fst <$> lookupArraySummary v
+  space <- lookupMemSpace v_mem
+  v' <- snd <$> ensureDirectArray (Just space) v
+  pure $ BasicOp $ Reshape k shape v'
+repairExpression e =
+  error $ "repairExpression:\n" <> prettyString e
+
+expReturns' ::
+  (Allocable fromrep torep inner) =>
+  Exp torep ->
+  AllocM fromrep torep ([ExpReturns], Exp torep)
+expReturns' e = do
+  maybe_rts <- expReturns e
+  case maybe_rts of
+    Just rts -> pure (rts, e)
+    Nothing -> do
+      e' <- repairExpression e
+      let bad =
+            error . unlines $
+              [ "expReturns': impossible index transformation",
+                prettyString e,
+                prettyString e'
+              ]
+      rts <- fromMaybe bad <$> expReturns e'
+      pure (rts, e')
+
 allocsForStm ::
   (Allocable fromrep torep inner) =>
   [Ident] ->
@@ -194,10 +227,10 @@
 allocsForStm idents e = do
   def_space <- askDefaultSpace
   hints <- expHints e
-  rts <- expReturns e
+  (rts, e') <- expReturns' e
   pes <- allocsForPat def_space idents rts hints
-  dec <- mkExpDecM (Pat pes) e
-  pure $ Let (Pat pes) (defAux dec) e
+  dec <- mkExpDecM (Pat pes) e'
+  pure $ Let (Pat pes) (defAux dec) e'
 
 patWithAllocations ::
   (MonadBuilder m, Mem (Rep m) inner) =>
@@ -209,7 +242,7 @@
 patWithAllocations def_space names e hints = do
   ts' <- instantiateShapes' names <$> expExtType e
   let idents = zipWith Ident names ts'
-  rts <- expReturns e
+  rts <- fromMaybe (error "patWithAllocations: ill-typed") <$> expReturns e
   Pat <$> allocsForPat def_space idents rts hints
 
 mkMissingIdents :: MonadFreshNames m => [Ident] -> [ExpReturns] -> m [Ident]
@@ -344,11 +377,9 @@
   mem_space <- lookupMemSpace mem
   default_space <- askDefaultSpace
   let space = fromMaybe default_space space_ok
-  if numLMADs ixfun == 1
-    && ixFunPerm ixfun == [0 .. IxFun.rank ixfun - 1]
+  if IxFun.permutation ixfun == [0 .. IxFun.rank ixfun - 1]
     && length (IxFun.base ixfun) == IxFun.rank ixfun
     && maybe True (== mem_space) space_ok
-    && IxFun.contiguous ixfun
     then pure (mem, v)
     else allocLinearArray space (baseString v) v
 
@@ -527,6 +558,22 @@
     _ ->
       error $ "allocPermArray: " ++ prettyString t
 
+ensurePermArray ::
+  (Allocable fromrep torep inner) =>
+  Maybe Space ->
+  [Int] ->
+  VName ->
+  AllocM fromrep torep (VName, VName)
+ensurePermArray space_ok perm v = do
+  (mem, ixfun) <- lookupArraySummary v
+  mem_space <- lookupMemSpace mem
+  default_space <- askDefaultSpace
+  if length (IxFun.base ixfun) == length (IxFun.shape ixfun)
+    && IxFun.permutation ixfun == perm
+    && maybe True (== mem_space) space_ok
+    then pure (mem, v)
+    else allocPermArray (fromMaybe default_space space_ok) perm (baseString v) v
+
 allocLinearArray ::
   (Allocable fromrep torep inner) =>
   Space ->
@@ -710,25 +757,16 @@
 allocInLambda params body =
   mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body
 
-numLMADs :: IxFun -> Int
-numLMADs = length . IxFun.ixfunLMADs
-
-ixFunPerm :: IxFun -> [Int]
-ixFunPerm = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs
-
-ixFunMon :: IxFun -> [IxFun.Monotonicity]
-ixFunMon = map IxFun.ldMon . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs
-
 data MemReq
-  = MemReq Space [Int] [IxFun.Monotonicity] Rank Bool
-  | NeedsLinearisation Space
+  = MemReq Space [Int] Rank
+  | NeedsNormalisation Space
   deriving (Eq, Show)
 
 combMemReqs :: MemReq -> MemReq -> MemReq
-combMemReqs x@NeedsLinearisation {} _ = x
-combMemReqs _ y@NeedsLinearisation {} = y
-combMemReqs x@(MemReq x_space _ _ _ _) y@MemReq {} =
-  if x == y then x else NeedsLinearisation x_space
+combMemReqs x@NeedsNormalisation {} _ = x
+combMemReqs _ y@NeedsNormalisation {} = y
+combMemReqs x@(MemReq x_space _ _) y@MemReq {} =
+  if x == y then x else NeedsNormalisation x_space
 
 type MemReqType = MemInfo (Ext SubExp) NoUniqueness MemReq
 
@@ -738,13 +776,13 @@
 combMemReqTypes x _ = x
 
 contextRets :: MemReqType -> [MemInfo d u r]
-contextRets (MemArray _ shape _ (MemReq space _ _ (Rank base_rank) _)) =
+contextRets (MemArray _ shape _ (MemReq space _ (Rank base_rank))) =
   -- Memory + offset + base_rank + (stride,size)*rank.
   MemMem space
     : MemPrim int64
     : replicate base_rank (MemPrim int64)
     ++ replicate (2 * shapeRank shape) (MemPrim int64)
-contextRets (MemArray _ shape _ (NeedsLinearisation space)) =
+contextRets (MemArray _ shape _ (NeedsNormalisation space)) =
   -- Memory + offset + (base,stride,size)*rank.
   MemMem space
     : MemPrim int64
@@ -771,15 +809,10 @@
         (Array pt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) -> do
           space <- lookupMemSpace mem
           pure . MemArray pt shape u $
-            if numLMADs ixfun == 1
-              then
-                MemReq
-                  space
-                  (ixFunPerm ixfun)
-                  (ixFunMon ixfun)
-                  (Rank $ length $ IxFun.base ixfun)
-                  (IxFun.contiguous ixfun)
-              else NeedsLinearisation space
+            MemReq
+              space
+              (IxFun.permutation ixfun)
+              (Rank $ length $ IxFun.base ixfun)
         (_, MemMem space) -> pure $ MemMem space
         (_, MemPrim pt) -> pure $ MemPrim pt
         (_, MemAcc acc ispace ts u) -> pure $ MemAcc acc ispace ts u
@@ -800,17 +833,16 @@
         res_rets_acc ++ [inspect ctx_offset req]
       )
 
-    arrayInfo rank (NeedsLinearisation space) =
-      (space, [0 .. rank - 1], repeat IxFun.Inc, rank, True)
-    arrayInfo _ (MemReq space perm mon (Rank base_rank) contig) =
-      (space, perm, mon, base_rank, contig)
+    arrayInfo rank (NeedsNormalisation space) =
+      (space, [0 .. rank - 1], rank)
+    arrayInfo _ (MemReq space perm (Rank base_rank)) =
+      (space, perm, base_rank)
 
     inspect ctx_offset (MemArray pt shape u req) =
       let shape' = fmap (adjustExt num_new_ctx) shape
-          (space, perm, mon, base_rank, contig) = arrayInfo (shapeRank shape) req
+          (space, perm, base_rank) = arrayInfo (shapeRank shape) req
        in MemArray pt shape' u . ReturnsNewBlock space ctx_offset $
-            convert
-              <$> IxFun.mkExistential base_rank (zip perm mon) contig (ctx_offset + 1)
+            convert <$> IxFun.mkExistential base_rank perm (ctx_offset + 1)
     inspect _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u
     inspect _ (MemPrim pt) = MemPrim pt
     inspect _ (MemMem space) = MemMem space
@@ -828,13 +860,14 @@
   Body torep ->
   AllocM fromrep torep (Body torep)
 addCtxToMatchBody reqs body = buildBody_ $ do
-  res <- zipWithM linearIfNeeded reqs =<< bodyBind body
+  res <- zipWithM normaliseIfNeeded reqs =<< bodyBind body
   ctx <- concat <$> mapM resCtx res
   pure $ ctx ++ res
   where
-    linearIfNeeded (MemArray _ _ _ (NeedsLinearisation space)) (SubExpRes cs (Var v)) =
-      SubExpRes cs . Var . snd <$> ensureRowMajorArray (Just space) v
-    linearIfNeeded _ res =
+    normaliseIfNeeded (MemArray _ shape _ (NeedsNormalisation space)) (SubExpRes cs (Var v)) =
+      SubExpRes cs . Var . snd
+        <$> ensurePermArray (Just space) [0 .. shapeRank shape - 1] v
+    normaliseIfNeeded _ res =
       pure res
 
     resCtx (SubExpRes _ Constant {}) =
@@ -1136,7 +1169,7 @@
       innerUsage inner
 
     simplifyPat (Pat pes) e = do
-      rets <- expReturns e
+      rets <- fromMaybe (error "simplifyPat: ill-typed") <$> expReturns e
       Pat <$> zipWithM update pes rets
       where
         names = map patElemName pes
diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs
--- a/src/Futhark/Passes.hs
+++ b/src/Futhark/Passes.hs
@@ -1,12 +1,12 @@
 -- | Optimisation pipelines.
 module Futhark.Passes
   ( standardPipeline,
-    sequentialPipeline,
-    kernelsPipeline,
-    sequentialCpuPipeline,
+    seqPipeline,
     gpuPipeline,
+    seqmemPipeline,
+    gpumemPipeline,
     mcPipeline,
-    multicorePipeline,
+    mcmemPipeline,
   )
 where
 
@@ -82,8 +82,8 @@
 
 -- | The pipeline used by the CUDA and OpenCL backends, but before
 -- adding memory information.  Includes 'standardPipeline'.
-kernelsPipeline :: Pipeline SOACS GPU
-kernelsPipeline =
+gpuPipeline :: Pipeline SOACS GPU
+gpuPipeline =
   standardPipeline
     >>> onePass extractKernels
     >>> passes
@@ -110,8 +110,8 @@
 
 -- | The pipeline used by the sequential backends.  Turns all
 -- parallelism into sequential loops.  Includes 'standardPipeline'.
-sequentialPipeline :: Pipeline SOACS Seq
-sequentialPipeline =
+seqPipeline :: Pipeline SOACS Seq
+seqPipeline =
   standardPipeline
     >>> onePass firstOrderTransform
     >>> passes
@@ -119,11 +119,11 @@
         inPlaceLoweringSeq
       ]
 
--- | Run 'sequentialPipeline', then add memory information (and
+-- | Run 'seqPipeline', then add memory information (and
 -- optimise it slightly).
-sequentialCpuPipeline :: Pipeline SOACS SeqMem
-sequentialCpuPipeline =
-  sequentialPipeline
+seqmemPipeline :: Pipeline SOACS SeqMem
+seqmemPipeline =
+  seqPipeline
     >>> onePass Seq.explicitAllocations
     >>> passes
       [ performCSE False,
@@ -140,11 +140,11 @@
         simplifySeqMem
       ]
 
--- | Run 'kernelsPipeline', then add memory information (and optimise
+-- | Run 'gpuPipeline', then add memory information (and optimise
 -- it a lot).
-gpuPipeline :: Pipeline SOACS GPUMem
-gpuPipeline =
-  kernelsPipeline
+gpumemPipeline :: Pipeline SOACS GPUMem
+gpumemPipeline =
+  gpuPipeline
     >>> onePass GPU.explicitAllocations
     >>> passes
       [ simplifyGPUMem,
@@ -185,8 +185,8 @@
       ]
 
 -- | Run 'mcPipeline' and then add memory information.
-multicorePipeline :: Pipeline SOACS MCMem
-multicorePipeline =
+mcmemPipeline :: Pipeline SOACS MCMem
+mcmemPipeline =
   mcPipeline
     >>> onePass MC.explicitAllocations
     >>> passes
diff --git a/src/Futhark/Util/Pretty.hs b/src/Futhark/Util/Pretty.hs
--- a/src/Futhark/Util/Pretty.hs
+++ b/src/Futhark/Util/Pretty.hs
@@ -46,7 +46,7 @@
 import Data.Text qualified as T
 import Numeric.Half
 import Prettyprinter
-import Prettyprinter.Render.Terminal (AnsiStyle, Color (..), bgColor, bgColorDull, bold, color, colorDull)
+import Prettyprinter.Render.Terminal (AnsiStyle, Color (..), bgColor, bgColorDull, bold, color, colorDull, italicized, underlined)
 import Prettyprinter.Render.Terminal qualified
 import Prettyprinter.Render.Text qualified
 import Prettyprinter.Symbols.Ascii
diff --git a/src/Language/Futhark/Parser.hs b/src/Language/Futhark/Parser.hs
--- a/src/Language/Futhark/Parser.hs
+++ b/src/Language/Futhark/Parser.hs
@@ -5,7 +5,7 @@
     parseExp,
     parseModExp,
     parseType,
-    parseDecOrExpIncrM,
+    parseDecOrExp,
     SyntaxError (..),
     Comment (..),
   )
@@ -57,29 +57,13 @@
   Either SyntaxError UncheckedTypeExp
 parseType = parse futharkType
 
--- | Parse an Futhark expression incrementally from monadic actions, using the
--- 'FilePath' as the source name for error messages.
-parseExpIncrM ::
-  Monad m =>
-  m T.Text ->
-  FilePath ->
-  T.Text ->
-  m (Either SyntaxError UncheckedExp)
-parseExpIncrM fetch file program =
-  getLinesFromM fetch $ fmap fst <$> parseInMonad expression file program
-
--- | Parse either an expression or a declaration incrementally;
--- favouring declarations in case of ambiguity.
-parseDecOrExpIncrM ::
-  Monad m =>
-  m T.Text ->
+-- | Parse either an expression or a declaration; favouring
+-- declarations in case of ambiguity.
+parseDecOrExp ::
   FilePath ->
   T.Text ->
-  m (Either SyntaxError (Either UncheckedDec UncheckedExp))
-parseDecOrExpIncrM fetch file input =
-  case parseInMonad declaration file input of
-    Value Left {} -> fmap Right <$> parseExpIncrM fetch file input
-    Value (Right (d, _)) -> pure $ Right $ Left d
-    GetLine _ -> do
-      l <- fetch
-      parseDecOrExpIncrM fetch file $ input <> "\n" <> l
+  Either SyntaxError (Either UncheckedDec UncheckedExp)
+parseDecOrExp file input =
+  case parse declaration file input of
+    Left {} -> Right <$> parseExp file input
+    Right d -> Right $ Left d
diff --git a/src/Language/Futhark/Parser/Lexer.x b/src/Language/Futhark/Parser/Lexer.x
--- a/src/Language/Futhark/Parser/Lexer.x
+++ b/src/Language/Futhark/Parser/Lexer.x
@@ -3,7 +3,7 @@
 -- | The Futhark lexer.  Takes a string, produces a list of tokens with position information.
 module Language.Futhark.Parser.Lexer
   ( Token(..)
-  , scanTokens
+  , getToken
   , scanTokensText
   ) where
 
@@ -12,7 +12,7 @@
 import qualified Data.Text as T
 import qualified Data.Text.Encoding as T
 import qualified Data.Text.Read as T
-import Data.Char (ord, toLower, digitToInt)
+import Data.Char (chr, ord, toLower)
 import Data.Int (Int8, Int16, Int32, Int64)
 import Data.Word (Word8)
 import Data.Loc (Loc (..), L(..), Pos(..))
@@ -20,7 +20,7 @@
 
 import Language.Futhark.Core (Int8, Int16, Int32, Int64,
                               Word8, Word16, Word32, Word64,
-                              Name, nameFromText, nameToText)
+                              Name, nameFromText, nameToText, nameFromString)
 import Language.Futhark.Prop (leadingOperator)
 import Language.Futhark.Syntax (BinOp(..))
 import Language.Futhark.Parser.Lexer.Wrapper
@@ -34,7 +34,6 @@
 @declit = [0-9][0-9_]*
 @binlit = 0[bB][01][01_]*
 @romlit = 0[rR][IVXLCDM][IVXLCDM_]*
-@intlit = @hexlit|@binlit|@declit|@romlit
 @reallit = (([0-9][0-9_]*("."[0-9][0-9_]*)?))([eE][\+\-]?[0-9]+)?
 @hexreallit = 0[xX][0-9a-fA-F][0-9a-fA-F_]*"."[0-9a-fA-F][0-9a-fA-F_]*([pP][\+\-]?[0-9_]+)
 
@@ -54,7 +53,7 @@
 tokens :-
 
   $white+                               ;
-  @doc                     { tokenM $ pure . DOC . T.intercalate "\n" .
+  @doc                     { tokenS $ DOC . T.intercalate "\n" .
                                       map (T.drop 3 . T.stripStart) .
                                            T.split (== '\n') . ("--"<>) .
                                            T.drop 4 }
@@ -91,62 +90,123 @@
   "$"                      { tokenC DOLLAR }
   "???"                    { tokenC HOLE }
 
-  @declit                  { tokenS $ \s -> NATLIT (nameFromText s) (readIntegral s) }
-  @intlit i8               { tokenM $ pure . I8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }
-  @intlit i16              { tokenM $ pure . I16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }
-  @intlit i32              { tokenM $ pure . I32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }
-  @intlit i64              { tokenM $ pure . I64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') }
-  @intlit u8               { tokenM $ pure . U8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }
-  @intlit u16              { tokenM $ pure . U16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }
-  @intlit u32              { tokenM $ pure . U32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }
-  @intlit u64              { tokenM $ pure . U64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') }
-  @intlit                  { tokenM $ pure . INTLIT . readIntegral . T.filter (/= '_') }
+  @declit                  { \s -> decToken (NATLIT (nameFromBS s)) s }
 
-  [\n[^\.]] ^ @reallit f16     { tokenM $ fmap F16LIT . tryRead "f16" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
-  [\n[^\.]] ^ @reallit f32     { tokenM $ fmap F32LIT . tryRead "f32" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
-  [\n[^\.]] ^ @reallit f64     { tokenM $ fmap F64LIT . tryRead "f64" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
-  [\n[^\.]] ^ @reallit         { tokenM $ fmap FLOATLIT . tryRead "f64" . suffZero . T.filter (/= '_') }
-  @hexreallit f16          { tokenM $ fmap F16LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
-  @hexreallit f32          { tokenM $ fmap F32LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
-  @hexreallit f64          { tokenM $ fmap F64LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
-  @hexreallit              { tokenM $ fmap FLOATLIT . readHexRealLit . T.filter (/= '_') }
-  "'" @charlit "'"         { tokenM $ fmap CHARLIT . tryRead "char" }
-  \" @stringcharlit* \"    { tokenM $ fmap (STRINGLIT . T.pack) . tryRead "string"  }
+  @declit i8               { decToken I8LIT . BS.dropEnd 2 }
+  @binlit i8               { binToken I8LIT . BS.drop 2 . BS.dropEnd 2 }
+  @hexlit i8               { hexToken I8LIT . BS.drop 2 . BS.dropEnd 2 }
+  @romlit i8               { romToken I8LIT . BS.drop 2 . BS.dropEnd 2 }
 
-  @identifier              { tokenS keyword }
+  @declit i16              { decToken I16LIT . BS.dropEnd 3 }
+  @binlit i16              { binToken I16LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit i16              { hexToken I16LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit i16              { romToken I16LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit i32              { decToken I32LIT . BS.dropEnd 3 }
+  @binlit i32              { binToken I32LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit i32              { hexToken I32LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit i32              { romToken I32LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit i64              { decToken I64LIT . BS.dropEnd 3 }
+  @binlit i64              { binToken I64LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit i64              { hexToken I64LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit i64              { romToken I64LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit u8               { decToken U8LIT . BS.dropEnd 2 }
+  @binlit u8               { binToken U8LIT . BS.drop 2 . BS.dropEnd 2 }
+  @hexlit u8               { hexToken U8LIT . BS.drop 2 . BS.dropEnd 2 }
+  @romlit u8               { romToken U8LIT . BS.drop 2 . BS.dropEnd 2 }
+
+  @declit u16              { decToken U16LIT . BS.dropEnd 3 }
+  @binlit u16              { binToken U16LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit u16              { hexToken U16LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit u16              { romToken U16LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit u32              { decToken U32LIT . BS.dropEnd 3 }
+  @binlit u32              { binToken U32LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit u32              { hexToken U32LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit u32              { romToken U32LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit u64              { decToken U64LIT . BS.dropEnd 3 }
+  @binlit u64              { binToken U64LIT . BS.drop 2 . BS.dropEnd 3 }
+  @hexlit u64              { hexToken U64LIT . BS.drop 2 . BS.dropEnd 3 }
+  @romlit u64              { romToken U64LIT . BS.drop 2 . BS.dropEnd 3 }
+
+  @declit                  { decToken INTLIT }
+  @binlit                  { binToken INTLIT . BS.drop 2 }
+  @hexlit                  { hexToken INTLIT . BS.drop 2 }
+  @romlit                  { romToken INTLIT . BS.drop 2 }
+
+  [\n[^\.]] ^ @reallit f16 { tokenS $ F16LIT . tryRead "f16" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
+  [\n[^\.]] ^ @reallit f32 { tokenS $ F32LIT . tryRead "f32" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
+  [\n[^\.]] ^ @reallit f64 { tokenS $ F64LIT . tryRead "f64" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') }
+  [\n[^\.]] ^ @reallit     { tokenS $ FLOATLIT . tryRead "f64" . suffZero . T.filter (/= '_') }
+  @hexreallit f16          { tokenS $ F16LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
+  @hexreallit f32          { tokenS $ F32LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
+  @hexreallit f64          { tokenS $ F64LIT . readHexRealLit . T.filter (/= '_') . T.dropEnd 3 }
+  @hexreallit              { tokenS $ FLOATLIT . readHexRealLit . T.filter (/= '_') }
+  "'" @charlit "'"         { tokenS $ CHARLIT . tryRead "char" }
+  \" @stringcharlit* \"    { tokenS $ STRINGLIT . T.pack . tryRead "string"  }
+
+  "true"                   { tokenC TRUE }
+  "false"                  { tokenC FALSE }
+  "if"                     { tokenC IF }
+  "then"                   { tokenC THEN }
+  "else"                   { tokenC ELSE }
+  "def"                    { tokenC DEF }
+  "let"                    { tokenC LET }
+  "loop"                   { tokenC LOOP }
+  "in"                     { tokenC IN }
+  "val"                    { tokenC VAL }
+  "for"                    { tokenC FOR }
+  "do"                     { tokenC DO }
+  "with"                   { tokenC WITH }
+  "local"                  { tokenC LOCAL }
+  "open"                   { tokenC OPEN }
+  "include"                { tokenC INCLUDE }
+  "import"                 { tokenC IMPORT }
+  "type"                   { tokenC TYPE }
+  "entry"                  { tokenC ENTRY }
+  "module"                 { tokenC MODULE }
+  "while"                  { tokenC WHILE }
+  "assert"                 { tokenC ASSERT }
+  "match"                  { tokenC MATCH }
+  "case"                   { tokenC CASE }
+
+  @identifier              { tokenS $ ID . nameFromText }
+
   "#" @identifier          { tokenS $ CONSTRUCTOR . nameFromText . T.drop 1 }
 
-  @binop                   { tokenM $ pure . symbol [] . nameFromText }
-  @qualbinop               { tokenM $ \s -> do (qs,k) <- mkQualId s; pure (symbol qs k) }
+  @binop                   { tokenS $ symbol [] . nameFromText }
+  @qualbinop               { tokenS $ uncurry symbol . mkQualId }
 {
 
-getToken :: Alex (Lexeme Token)
-getToken = do
-  inp@(_,_,_,n) <- alexGetInput
-  sc <- alexGetStartCode
-  case alexScan inp sc of
-    AlexEOF -> do pos <- alexGetPos
-                  pure (pos, pos, EOF)
+nameFromBS :: BS.ByteString -> Name
+nameFromBS = nameFromString . map (chr . fromIntegral) . BS.unpack
+
+getToken :: AlexInput -> Either LexerError (AlexInput, (Pos, Pos, Token))
+getToken state@(pos,c,s,n) =
+  case alexScan state 0 of
+    AlexEOF -> Right (state, (pos, pos, EOF))
     AlexError (pos,_,_,_) ->
-      alexError (Loc pos pos) "Invalid lexical syntax."
-    AlexSkip  inp' _len -> do
-      alexSetInput inp'
-      getToken
-    AlexToken inp'@(_,_,_,n') _ action -> let len = n'-n in do
-      alexSetInput inp'
-      action inp len
+      Left $ LexerError (Loc pos pos) "Invalid lexical syntax."
+    AlexSkip state' _len ->
+      getToken state'
+    AlexToken state'@(pos',_,_,n') _ action -> do
+      let x = action (BS.take (n'-n) s)
+      x `seq` Right (state', (pos, pos', x))
 
 scanTokens :: Pos -> BS.ByteString -> Either LexerError ([L Token], Pos)
-scanTokens pos str =
-  runAlex pos str $ do
-  fix $ \loop -> do
-    tok <- getToken
-    case tok of
-      (start, end, EOF) ->
-        pure ([], end)
-      (start, end, t) -> do
-        (rest, endpos) <- loop
-        pure (L (Loc start end) t : rest, endpos)
+scanTokens pos str = loop $ initialLexerState pos str
+  where
+   loop s = do
+     (s', tok) <- getToken s
+     case tok of
+       (start, end, EOF) ->
+         pure ([], end)
+       (start, end, t) -> do
+         (rest, endpos) <- loop s'
+         pure (L (Loc start end) t : rest, endpos)
 
 -- | Given a starting position, produce tokens from the given text (or
 -- a lexer error).  Returns the final position.
diff --git a/src/Language/Futhark/Parser/Lexer/Tokens.hs b/src/Language/Futhark/Parser/Lexer/Tokens.hs
--- a/src/Language/Futhark/Parser/Lexer/Tokens.hs
+++ b/src/Language/Futhark/Parser/Lexer/Tokens.hs
@@ -5,28 +5,27 @@
 -- Also defines other useful building blocks for constructing tokens.
 module Language.Futhark.Parser.Lexer.Tokens
   ( Token (..),
-    Lexeme,
     fromRoman,
     symbol,
     mkQualId,
-    tokenPosM,
-    tokenM,
     tokenC,
-    keyword,
     tokenS,
-    indexing,
     suffZero,
     tryRead,
-    readIntegral,
+    decToken,
+    binToken,
+    hexToken,
+    romToken,
+    advance,
     readHexRealLit,
   )
 where
 
 import Data.ByteString.Lazy qualified as BS
-import Data.Char (digitToInt, ord)
+import Data.Char (ord)
 import Data.Either
-import Data.List (find, foldl')
-import Data.Loc (Loc (..), Pos (..))
+import Data.List (find)
+import Data.Loc (Pos (..))
 import Data.Text qualified as T
 import Data.Text.Encoding qualified as T
 import Data.Text.Read qualified as T
@@ -41,7 +40,6 @@
     Word64,
     Word8,
   )
-import Language.Futhark.Parser.Lexer.Wrapper
 import Language.Futhark.Prop (leadingOperator)
 import Language.Futhark.Syntax (BinOp, nameFromText, nameToText)
 import Numeric.Half
@@ -133,94 +131,68 @@
   | HOLE
   deriving (Show, Eq, Ord)
 
-keyword :: T.Text -> Token
-keyword s =
-  case s of
-    "true" -> TRUE
-    "false" -> FALSE
-    "if" -> IF
-    "then" -> THEN
-    "else" -> ELSE
-    "def" -> DEF
-    "let" -> LET
-    "loop" -> LOOP
-    "in" -> IN
-    "val" -> VAL
-    "for" -> FOR
-    "do" -> DO
-    "with" -> WITH
-    "local" -> LOCAL
-    "open" -> OPEN
-    "include" -> INCLUDE
-    "import" -> IMPORT
-    "type" -> TYPE
-    "entry" -> ENTRY
-    "module" -> MODULE
-    "while" -> WHILE
-    "assert" -> ASSERT
-    "match" -> MATCH
-    "case" -> CASE
-    _ -> ID $ nameFromText s
-
-indexing :: (Loc, T.Text) -> Alex Name
-indexing (loc, s) = case keyword s of
-  ID v -> pure v
-  _ -> alexError loc $ "Cannot index keyword '" <> s <> "'."
-
-mkQualId :: T.Text -> Alex ([Name], Name)
+mkQualId :: T.Text -> ([Name], Name)
 mkQualId s = case reverse $ T.splitOn "." s of
   [] -> error "mkQualId: no components"
-  k : qs -> pure (map nameFromText (reverse qs), nameFromText k)
+  k : qs -> (map nameFromText (reverse qs), nameFromText k)
 
 -- | Suffix a zero if the last character is dot.
 suffZero :: T.Text -> T.Text
 suffZero s = if T.last s == '.' then s <> "0" else s
 
-tryRead :: Read a => String -> T.Text -> Alex a
+tryRead :: Read a => String -> T.Text -> a
 tryRead desc s = case reads s' of
-  [(x, "")] -> pure x
+  [(x, "")] -> x
   _ -> error $ "Invalid " ++ desc ++ " literal: `" ++ T.unpack s ++ "'."
   where
     s' = T.unpack s
 
-readIntegral :: Integral a => T.Text -> a
-readIntegral s
-  | "0x" `T.isPrefixOf` s || "0X" `T.isPrefixOf` s = parseBase 16 (T.drop 2 s)
-  | "0b" `T.isPrefixOf` s || "0B" `T.isPrefixOf` s = parseBase 2 (T.drop 2 s)
-  | "0r" `T.isPrefixOf` s || "0R" `T.isPrefixOf` s = fromRoman (T.drop 2 s)
-  | otherwise = parseBase 10 s
-  where
-    parseBase base = T.foldl (\acc c -> acc * base + fromIntegral (digitToInt c)) 0
+{-# INLINE tokenC #-}
+tokenC :: a -> BS.ByteString -> a
+tokenC v _ = v
 
-tokenC :: a -> (Pos, Char, BS.ByteString, Int64) -> Int64 -> Alex (Lexeme a)
-tokenC v = tokenS $ const v
+{-# INLINE decToken #-}
+decToken :: Integral a => (a -> Token) -> BS.ByteString -> Token
+decToken f = f . BS.foldl' digit 0
+  where
+    digit x c =
+      if c >= 48 && c <= 57
+        then x * 10 + fromIntegral (c - 48)
+        else x
 
-tokenS :: (T.Text -> a) -> (Pos, Char, BS.ByteString, Int64) -> Int64 -> Alex (Lexeme a)
-tokenS f = tokenM $ pure . f
+{-# INLINE binToken #-}
+binToken :: Integral a => (a -> Token) -> BS.ByteString -> Token
+binToken f = f . BS.foldl' digit 0
+  where
+    digit x c =
+      if c >= 48 && c <= 49
+        then x * 2 + fromIntegral (c - 48)
+        else x
 
-type Lexeme a = (Pos, Pos, a)
+{-# INLINE hexToken #-}
+hexToken :: Integral a => (a -> Token) -> BS.ByteString -> Token
+hexToken f = f . BS.foldl' digit 0
+  where
+    digit x c
+      | c >= 48 && c <= 57 =
+          x * 16 + fromIntegral (c - 48)
+      | c >= 65 && c <= 70 =
+          x * 16 + fromIntegral (10 + c - 65)
+      | c >= 97 && c <= 102 =
+          x * 16 + fromIntegral (10 + c - 97)
+      | otherwise =
+          x
 
-tokenM ::
-  (T.Text -> Alex a) ->
-  (Pos, Char, BS.ByteString, Int64) ->
-  Int64 ->
-  Alex (Lexeme a)
-tokenM f = tokenPosM (f . snd)
+{-# INLINE romToken #-}
+romToken :: Integral a => (a -> Token) -> BS.ByteString -> Token
+romToken f = tokenS $ f . fromRoman
 
-tokenPosM ::
-  ((Loc, T.Text) -> Alex a) ->
-  (Pos, Char, BS.ByteString, Int64) ->
-  Int64 ->
-  Alex (Lexeme a)
-tokenPosM f (pos, _, s, _) len = do
-  x <- f (Loc pos pos', T.decodeUtf8 $ BS.toStrict s')
-  pure (pos, pos', x)
-  where
-    pos' = advance pos s'
-    s' = BS.take len s
+{-# INLINE tokenS #-}
+tokenS :: (T.Text -> a) -> BS.ByteString -> a
+tokenS f = f . T.decodeUtf8 . BS.toStrict
 
 advance :: Pos -> BS.ByteString -> Pos
-advance orig_pos = foldl' advance' orig_pos . init . BS.unpack
+advance orig_pos = BS.foldl' advance' orig_pos . BS.init
   where
     advance' (Pos f !line !col !addr) c
       | c == nl = Pos f (line + 1) 1 (addr + 1)
@@ -261,7 +233,7 @@
     Nothing -> 0
     Just (d, n) -> n + fromRoman (T.drop (T.length d) s)
 
-readHexRealLit :: RealFloat a => T.Text -> Alex a
+readHexRealLit :: RealFloat a => T.Text -> a
 readHexRealLit s =
   let num = T.drop 2 s
    in -- extract number into integer, fractional and (optional) exponent
@@ -276,5 +248,5 @@
                   fracLen = fromIntegral $ T.length f
                   fracVal = fracPart / (16.0 ** fracLen)
                   totalVal = (intPart + fracVal) * (2.0 ** exponent)
-               in pure totalVal
+               in totalVal
             _ -> error "bad hex real literal"
diff --git a/src/Language/Futhark/Parser/Lexer/Wrapper.hs b/src/Language/Futhark/Parser/Lexer/Wrapper.hs
--- a/src/Language/Futhark/Parser/Lexer/Wrapper.hs
+++ b/src/Language/Futhark/Parser/Lexer/Wrapper.hs
@@ -1,27 +1,18 @@
 {-# OPTIONS_GHC -funbox-strict-fields #-}
 
 -- | Utility definitions used by the lexer.  None of the default Alex
--- "wrappers" are precisely what we need.  The code here is based on
--- the "monad-bytestring" wrapper.  The code here is completely
--- Futhark-agnostic, and perhaps it can even serve as inspiration for
--- other Alex lexer wrappers.
+-- "wrappers" are precisely what we need.  The code here is highly
+-- minimalistic.  Lexers should not be complicated!
 module Language.Futhark.Parser.Lexer.Wrapper
-  ( runAlex,
-    Alex,
+  ( initialLexerState,
     AlexInput,
     alexInputPrevChar,
-    Byte,
     LexerError (..),
-    alexSetInput,
-    alexGetInput,
     alexGetByte,
-    alexGetStartCode,
-    alexError,
     alexGetPos,
   )
 where
 
-import Control.Applicative (liftA)
 import Data.ByteString.Internal qualified as BS (w2c)
 import Data.ByteString.Lazy qualified as BS
 import Data.Int (Int64)
@@ -61,6 +52,9 @@
           n' = n + 1
        in p' `seq` cs' `seq` n' `seq` Just (b, (p', c, cs', n'))
 
+alexGetPos :: AlexInput -> Pos
+alexGetPos (pos, _, _, _) = pos
+
 tabSize :: Int
 tabSize = 8
 
@@ -70,71 +64,11 @@
 alexMove (Pos !f !l _ !a) '\n' = Pos f (l + 1) 1 (a + 1)
 alexMove (Pos !f !l !c !a) _ = Pos f l (c + 1) (a + 1)
 
-data AlexState = AlexState
-  { alex_pos :: !Pos, -- position at current input location
-    alex_bpos :: !Int64, -- bytes consumed so far
-    alex_inp :: BS.ByteString, -- the current input
-    alex_chr :: !Char, -- the character before the input
-    alex_scd :: !Int -- the current startcode
-  }
-
-runAlex :: Pos -> BS.ByteString -> Alex a -> Either LexerError a
-runAlex start_pos input (Alex f) =
-  case f
-    ( AlexState
-        { alex_pos = start_pos,
-          alex_bpos = 0,
-          alex_inp = input,
-          alex_chr = '\n',
-          alex_scd = 0
-        }
-    ) of
-    Left msg -> Left msg
-    Right (_, a) -> Right a
-
-newtype Alex a = Alex {unAlex :: AlexState -> Either LexerError (AlexState, a)}
+initialLexerState :: Pos -> BS.ByteString -> AlexInput
+initialLexerState start_pos input =
+  (start_pos, '\n', input, 0)
 
 data LexerError = LexerError Loc T.Text
 
 instance Show LexerError where
   show (LexerError _ s) = T.unpack s
-
-instance Functor Alex where
-  fmap = liftA
-
-instance Applicative Alex where
-  pure a = Alex $ \s -> Right (s, a)
-  fa <*> a = Alex $ \s -> case unAlex fa s of
-    Left msg -> Left msg
-    Right (s', f) -> case unAlex a s' of
-      Left msg -> Left msg
-      Right (s'', b) -> Right (s'', f b)
-
-instance Monad Alex where
-  m >>= k = Alex $ \s -> case unAlex m s of
-    Left msg -> Left msg
-    Right (s', a) -> unAlex (k a) s'
-
-alexGetInput :: Alex AlexInput
-alexGetInput =
-  Alex $ \s@AlexState {alex_pos = pos, alex_bpos = bpos, alex_chr = c, alex_inp = inp} ->
-    Right (s, (pos, c, inp, bpos))
-
-alexSetInput :: AlexInput -> Alex ()
-alexSetInput (pos, c, inp, bpos) =
-  Alex $ \s -> case s
-    { alex_pos = pos,
-      alex_bpos = bpos,
-      alex_chr = c,
-      alex_inp = inp
-    } of
-    state@AlexState {} -> Right (state, ())
-
-alexError :: Loc -> T.Text -> Alex a
-alexError loc message = Alex $ const $ Left $ LexerError loc message
-
-alexGetStartCode :: Alex Int
-alexGetStartCode = Alex $ \s@AlexState {alex_scd = sc} -> Right (s, sc)
-
-alexGetPos :: Alex Pos
-alexGetPos = Alex $ \s -> Right (s, alex_pos s)
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
@@ -6,12 +6,9 @@
 module Language.Futhark.Parser.Monad
   ( ParserMonad,
     ParserState,
-    ReadLineMonad (..),
     Comment (..),
-    parseInMonad,
     parse,
     parseWithComments,
-    getLinesFromM,
     lexer,
     mustBeEmpty,
     arrayFromList,
@@ -38,19 +35,20 @@
   )
 where
 
-import Control.Applicative (liftA)
 import Control.Monad
 import Control.Monad.Except (ExceptT, MonadError (..), runExceptT)
 import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.State
 import Data.Array hiding (index)
+import Data.ByteString.Lazy qualified as BS
 import Data.List.NonEmpty qualified as NE
 import Data.Monoid
 import Data.Text qualified as T
+import Data.Text.Encoding qualified as T
 import Futhark.Util.Loc
 import Futhark.Util.Pretty hiding (line, line')
 import Language.Futhark.Parser.Lexer
-import Language.Futhark.Parser.Lexer.Wrapper (LexerError (..))
+import Language.Futhark.Parser.Lexer.Wrapper (AlexInput, LexerError (..), initialLexerState)
 import Language.Futhark.Pretty ()
 import Language.Futhark.Prop
 import Language.Futhark.Syntax
@@ -105,38 +103,10 @@
     parserInput :: T.Text,
     -- | Note: reverse order.
     parserComments :: [Comment],
-    parserLexical :: ([L Token], Pos)
+    parserLexerState :: AlexInput
   }
 
-type ParserMonad = ExceptT SyntaxError (StateT ParserState ReadLineMonad)
-
-data ReadLineMonad a
-  = Value a
-  | GetLine (Maybe T.Text -> ReadLineMonad a)
-
-readLineFromMonad :: ReadLineMonad (Maybe T.Text)
-readLineFromMonad = GetLine Value
-
-instance Monad ReadLineMonad where
-  Value x >>= f = f x
-  GetLine g >>= f = GetLine $ g >=> f
-
-instance Functor ReadLineMonad where
-  fmap = liftA
-
-instance Applicative ReadLineMonad where
-  pure = Value
-  (<*>) = ap
-
-getLinesFromM :: Monad m => m T.Text -> ReadLineMonad a -> m a
-getLinesFromM _ (Value x) = pure x
-getLinesFromM fetch (GetLine f) = do
-  s <- fetch
-  getLinesFromM fetch $ f $ Just s
-
-getNoLines :: ReadLineMonad a -> Either SyntaxError a
-getNoLines (Value x) = Right x
-getNoLines (GetLine f) = getNoLines $ f Nothing
+type ParserMonad = ExceptT SyntaxError (State ParserState)
 
 arrayFromList :: [a] -> Array Int a
 arrayFromList l = listArray (0, length l - 1) l
@@ -169,9 +139,6 @@
   where
     field (name, pat) = RecordFieldExplicit name <$> patternExp pat <*> pure loc
 
-eof :: Pos -> L Token
-eof pos = L (Loc pos pos) EOF
-
 binOpName :: L Token -> (QualName Name, Loc)
 binOpName (L loc (SYMBOL _ qs op)) = (QualName qs op, loc)
 binOpName t = error $ "binOpName: unexpected " ++ show t
@@ -181,12 +148,6 @@
   AppExp (BinOp (QualName qs op, srclocOf loc) NoInfo (x, NoInfo) (y, NoInfo) (srcspan x y)) NoInfo
 binOp _ t _ = error $ "binOp: unexpected " ++ show t
 
-getTokens :: ParserMonad ([L Token], Pos)
-getTokens = lift $ gets parserLexical
-
-putTokens :: ([L Token], Pos) -> ParserMonad ()
-putTokens l = lift $ modify $ \env -> env {parserLexical = l}
-
 putComment :: Comment -> ParserMonad ()
 putComment c = lift $ modify $ \env ->
   env {parserComments = c : parserComments env}
@@ -208,42 +169,21 @@
 primNegate (UnsignedValue v) = UnsignedValue $ intNegate v
 primNegate (BoolValue v) = BoolValue $ not v
 
-readLine :: ParserMonad (Maybe T.Text)
-readLine = do
-  s <- lift $ lift readLineFromMonad
-  case s of
-    Just s' ->
-      lift $ modify $ \env -> env {parserInput = parserInput env <> "\n" <> s'}
-    Nothing -> pure ()
-  pure s
-
 lexer :: (L Token -> ParserMonad a) -> ParserMonad a
 lexer cont = do
-  (ts, pos) <- getTokens
-  case ts of
-    [] -> do
-      ended <- lift $ runExceptT $ cont $ eof pos
-      case ended of
-        Right x -> pure x
-        Left parse_e -> do
-          line <- readLine
-          ts' <-
-            case line of
-              Nothing -> throwError parse_e
-              Just line' -> pure $ scanTokensText (advancePos pos '\n') line'
-          (ts'', pos') <- either (throwError . lexerErrToParseErr) pure ts'
-          case ts'' of
-            [] -> cont $ eof pos
-            xs -> do
-              putTokens (xs, pos')
-              lexer cont
-    (L loc (COMMENT text) : xs) -> do
-      putComment $ Comment loc text
-      putTokens (xs, pos)
-      lexer cont
-    (x : xs) -> do
-      putTokens (xs, pos)
-      cont x
+  ls <- lift $ gets parserLexerState
+  case getToken ls of
+    Left e ->
+      throwError $ lexerErrToParseErr e
+    Right (ls', (start, end, tok)) -> do
+      let loc = Loc start end
+      lift $ modify $ \s -> s {parserLexerState = ls'}
+      case tok of
+        COMMENT text -> do
+          putComment $ Comment loc text
+          lexer cont
+        _ ->
+          cont $ L loc tok
 
 parseError :: (L Token, [String]) -> ParserMonad a
 parseError (L loc EOF, expected) =
@@ -257,7 +197,7 @@
 parseError (L loc _, expected) = do
   input <- lift $ gets parserInput
   let ~(Loc (Pos _ _ _ beg) (Pos _ _ _ end)) = locOf loc
-      tok_src = T.take (end - beg + 1) $ T.drop beg input
+      tok_src = T.take (end - beg) $ T.drop beg input
   parseErrorAt loc . Just . T.unlines $
     [ "Unexpected token: '" <> tok_src <> "'",
       "Expected one of the following: " <> T.unwords (map T.pack expected)
@@ -288,23 +228,23 @@
 lexerErrToParseErr :: LexerError -> SyntaxError
 lexerErrToParseErr (LexerError loc msg) = SyntaxError loc msg
 
-parseInMonad ::
+parseWithComments ::
   ParserMonad a ->
   FilePath ->
   T.Text ->
-  ReadLineMonad (Either SyntaxError (a, [Comment]))
-parseInMonad p file program =
-  either
-    (pure . Left . lexerErrToParseErr)
-    (fmap onRes . runStateT (runExceptT p) . env)
-    (scanTokensText (Pos file 1 1 0) program)
+  Either SyntaxError (a, [Comment])
+parseWithComments p file program =
+  onRes $ runState (runExceptT p) env
   where
-    env = ParserState file program []
+    env =
+      ParserState
+        file
+        program
+        []
+        (initialLexerState start $ BS.fromStrict . T.encodeUtf8 $ program)
+    start = Pos file 1 1 0
     onRes (Left err, _) = Left err
     onRes (Right x, s) = Right (x, reverse $ parserComments s)
-
-parseWithComments :: ParserMonad a -> FilePath -> T.Text -> Either SyntaxError (a, [Comment])
-parseWithComments p file program = join $ getNoLines $ parseInMonad p file program
 
 parse :: ParserMonad a -> FilePath -> T.Text -> Either SyntaxError a
 parse p file program = fst <$> parseWithComments p file program
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
@@ -8,9 +8,6 @@
   , futharkType
   , parse
   , parseWithComments
-  , ReadLineMonad (..)
-  , getLinesFromM
-  , parseInMonad
   , SyntaxError(..)
   , Comment(..)
   )
@@ -212,9 +209,12 @@
     | Doc Dec_          { addDoc $1 $2 }
 
 Decs :: { [UncheckedDec] }
-      :          { [] }
-      | Dec Decs { $1 : $2 }
+      : Decs_     { reverse $1 }
 
+Decs_ :: { [UncheckedDec] }
+      :           { [] }
+      | Decs_ Dec { $2 : $1 }
+
 Dec_ :: { UncheckedDec }
     : Val               { ValDec $1 }
     | TypeAbbr          { TypeDec $1 }
@@ -320,9 +320,12 @@
         { addAttrSpec $2 $4 }
 
 Specs :: { [SpecBase NoInfo Name] }
-       : Spec Specs { $1 : $2 }
-       |            { [] }
+       : Specs_      { reverse $1 }
 
+Specs_ :: { [SpecBase NoInfo Name] }
+       : Specs_ Spec { $2 : $1 }
+       |             { [] }
+
 SizeBinder :: { SizeBinder Name }
             : '[' id ']'  { let L _ (ID name) = $2 in SizeBinder name (srcspan $1 $>) }
             | '...[' id ']' { let L _ (ID name) = $2 in SizeBinder name (srcspan $1 $>) }
@@ -600,9 +603,9 @@
                         StringLit (BS.unpack (T.encodeUtf8 s)) (srclocOf loc) }
      | hole           { Hole NoInfo (srclocOf $1) }
      | '(' Exp ')'            { Parens $2 (srcspan $1 $>) }
-     | '(' Exp ',' Exps1 ')'  { TupLit ($2 : fst $4 : snd $4) (srcspan $1 $>) }
+     | '(' Exp ',' Exps1 ')'  { TupLit ($2 : $4) (srcspan $1 $>) }
      | '('      ')'           { TupLit [] (srcspan $1 $>) }
-     | '[' Exps1 ']'          { ArrayLit (fst $2:snd $2) NoInfo (srcspan $1 $>) }
+     | '[' Exps1 ']'          { ArrayLit $2 NoInfo (srcspan $1 $>) }
      | '['       ']'          { ArrayLit [] NoInfo (srcspan $1 $>) }
 
      | id { let L loc (ID v)  = $1 in Var (QualName [] v) NoInfo (srclocOf loc) }
@@ -651,14 +654,12 @@
         | false  { (BoolValue False, $1) }
         | NumLit { $1 }
 
-Exps1 :: { (UncheckedExp, [UncheckedExp]) }
-       : Exps1_ { case reverse (snd $1 : fst $1) of
-                    []   -> (snd $1, [])
-                    y:ys -> (y, ys) }
+Exps1 :: { [UncheckedExp] }
+       : Exps1_ { reverse $1 }
 
-Exps1_ :: { ([UncheckedExp], UncheckedExp) }
-        : Exps1_ ',' Exp { (snd $1 : fst $1, $3) }
-        | Exp            { ([], $1) }
+Exps1_ :: { [UncheckedExp] }
+        : Exps1_ ',' Exp { $3 : $1 }
+        | Exp            { [$1] }
 
 FieldAccesses :: { [(Name, Loc)] }
                : '.' FieldId FieldAccesses { $2 : $3 }
diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs
--- a/unittests/Futhark/IR/Mem/IxFunTests.hs
+++ b/unittests/Futhark/IR/Mem/IxFunTests.hs
@@ -8,6 +8,7 @@
 import Data.Function ((&))
 import Data.List qualified as L
 import Data.Map qualified as M
+import Data.Text qualified as T
 import Futhark.Analysis.PrimExp.Convert
 import Futhark.IR.Mem.IxFun qualified as IxFunLMAD
 import Futhark.IR.Mem.IxFun.Alg qualified as IxFunAlg
@@ -17,7 +18,7 @@
 import Futhark.IR.Syntax
 import Futhark.IR.Syntax.Core ()
 import Futhark.Util.IntegralExp qualified as IE
-import Futhark.Util.Pretty qualified as PR
+import Futhark.Util.Pretty
 import Test.Tasty
 import Test.Tasty.HUnit
 import Prelude hiding (span)
@@ -47,42 +48,42 @@
           ([], x)
           strides
 
-compareIxFuns :: IxFunLMAD.IxFun Int -> IxFunAlg.IxFun Int -> Assertion
-compareIxFuns ixfunLMAD ixfunAlg =
+compareIxFuns :: Maybe (IxFunLMAD.IxFun Int) -> IxFunAlg.IxFun Int -> Assertion
+compareIxFuns (Just ixfunLMAD) ixfunAlg =
   let lmadShape = IxFunLMAD.shape ixfunLMAD
       algShape = IxFunAlg.shape ixfunAlg
       points = allPoints lmadShape
       resLMAD = map (IxFunLMAD.index ixfunLMAD) points
       resAlg = map (IxFunAlg.index ixfunAlg) points
       errorMessage =
-        "lmad ixfun:  "
-          ++ PR.prettyString ixfunLMAD
-          ++ "\n"
-          ++ "alg ixfun:   "
-          ++ PR.prettyString ixfunAlg
-          ++ "\n"
-          ++ "lmad shape:  "
-          ++ show lmadShape
-          ++ "\n"
-          ++ "alg shape:   "
-          ++ show algShape
-          ++ "\n"
-          ++ "lmad points length: "
-          ++ show (length resLMAD)
-          ++ "\n"
-          ++ "alg points length:  "
-          ++ show (length resAlg)
-          ++ "\n"
-          ++ "lmad points: "
-          ++ show resLMAD
-          ++ "\n"
-          ++ "alg points:  "
-          ++ show resAlg
+        T.unpack . docText $
+          "lmad ixfun:  " <> pretty ixfunLMAD
+            </> "alg ixfun:   " <> pretty ixfunAlg
+            </> "lmad shape:  " <> pretty lmadShape
+            </> "alg shape:   " <> pretty algShape
+            </> "lmad points length: " <> pretty (length resLMAD)
+            </> "alg points length:  " <> pretty (length resAlg)
+            </> "lmad points: " <> pretty resLMAD
+            </> "alg points:  " <> pretty resAlg
    in (lmadShape == algShape && resLMAD == resAlg) @? errorMessage
+compareIxFuns Nothing ixfunAlg =
+  assertFailure $
+    unlines
+      [ "lmad ixfun: Nothing",
+        "alg ixfun:  " <> prettyString ixfunAlg
+      ]
 
 compareOps :: IxFunWrap.IxFun Int -> Assertion
 compareOps (ixfunLMAD, ixfunAlg) = compareIxFuns ixfunLMAD ixfunAlg
 
+compareOpsFailure :: IxFunWrap.IxFun Int -> Assertion
+compareOpsFailure (Nothing, _) = pure ()
+compareOpsFailure (Just ixfunLMAD, ixfunAlg) =
+  assertFailure . T.unpack . docText $
+    "Not supposed to be representable as LMAD."
+      </> "lmad ixfun: " <> pretty ixfunLMAD
+      </> "alg ixfun:  " <> pretty ixfunAlg
+
 -- XXX: Clean this up.
 n :: Int
 n = 19
@@ -102,10 +103,11 @@
     concat
       [ test_iota,
         test_slice_iota,
-        test_reshape_slice_iota1,
+        test_slice_reshape_iota1,
         test_permute_slice_iota,
+        test_reshape_iota,
         test_reshape_permute_iota,
-        test_reshape_slice_iota2,
+        test_slice_reshape_iota2,
         test_reshape_slice_iota3,
         test_complex1,
         test_complex2,
@@ -113,6 +115,8 @@
         test_rebase2,
         test_rebase3,
         test_rebase4_5,
+        test_rebase6,
+        test_rebase7,
         test_flatSlice_iota,
         test_slice_flatSlice_iota,
         test_flatSlice_flatSlice_iota,
@@ -129,194 +133,161 @@
 
 test_iota :: [TestTree]
 test_iota =
-  singleton $
-    testCase "iota" $
-      compareOps $
-        iota [n]
+  singleton . testCase "iota" . compareOps $
+    iota [n]
 
 test_slice_iota :: [TestTree]
 test_slice_iota =
-  singleton $
-    testCase "slice . iota" $
-      compareOps $
-        slice (iota [n, n, n]) slice3
+  singleton . testCase "slice . iota" . compareOps $
+    slice (iota [n, n, n]) slice3
 
-test_reshape_slice_iota1 :: [TestTree]
-test_reshape_slice_iota1 =
-  singleton $
-    testCase "reshape . slice . iota 1" $
-      compareOps $
-        reshape
-          (slice (iota [n, n, n]) slice3)
-          [n `P.div` 2, n `P.div` 3]
+test_slice_reshape_iota1 :: [TestTree]
+test_slice_reshape_iota1 =
+  singleton . testCase "slice . reshape . iota 1" . compareOps $
+    slice (reshape (iota [n, n, n]) [n `P.div` 2, n `P.div` 3, 1]) slice3
 
 test_permute_slice_iota :: [TestTree]
 test_permute_slice_iota =
-  singleton $
-    testCase "permute . slice . iota" $
-      compareOps $
-        permute (slice (iota [n, n, n]) slice3) [1, 0]
+  singleton . testCase "permute . slice . iota" . compareOps $
+    permute (slice (iota [n, n, n]) slice3) [1, 0]
 
+test_reshape_iota :: [TestTree]
+test_reshape_iota =
+  -- This tests a pattern that occurs with ScalarSpace.
+  singleton . testCase "reshape . zeroslice . iota" . compareOps $
+    let s = Slice [DimSlice 0 n 0, DimSlice 0 n 1]
+     in reshape (slice (iota [n, n]) s) [1, n, 1, n]
+
 test_reshape_permute_iota :: [TestTree]
 test_reshape_permute_iota =
   -- negative reshape test
-  singleton $
-    testCase "reshape . permute . iota" $
-      compareOps $
-        let newdims = [n * n, n]
-         in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims
+  singleton . testCase "reshape . permute . iota" . compareOpsFailure $
+    let newdims = [n * n, n]
+     in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims
 
-test_reshape_slice_iota2 :: [TestTree]
-test_reshape_slice_iota2 =
-  -- negative reshape test
-  singleton $
-    testCase "reshape . slice . iota 2" $
-      compareOps $
-        let newdims = [n * n, n]
-            slc =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice 0 n 1,
-                  DimSlice (n - 1) n (-1)
-                ]
-         in reshape (slice (iota [n, n, n, n]) slc) newdims
+test_slice_reshape_iota2 :: [TestTree]
+test_slice_reshape_iota2 =
+  singleton . testCase "slice . reshape . iota 2" . compareOps $
+    let newdims = [n * n, n]
+        slc =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice 0 n 1
+            ]
+     in slice (reshape (iota [n, n, n, n]) newdims) slc
 
 test_reshape_slice_iota3 :: [TestTree]
 test_reshape_slice_iota3 =
   -- negative reshape test
-  singleton $
-    testCase "reshape . slice . iota 3" $
-      compareOps $
-        let newdims = [n * n, n]
-            slc =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice 0 n 1,
-                  DimSlice 0 (n `P.div` 2) 1,
-                  DimSlice 0 n 1
-                ]
-         in reshape (slice (iota [n, n, n, n]) slc) newdims
+  singleton . testCase "reshape . slice . iota 3" . compareOpsFailure $
+    let newdims = [n * n, n]
+        slc =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice 0 n 1,
+              DimSlice 0 (n `P.div` 2) 1,
+              DimSlice 0 n 1
+            ]
+     in reshape (slice (iota [n, n, n, n]) slc) newdims
 
 test_complex1 :: [TestTree]
 test_complex1 =
-  singleton $
-    testCase "reshape . permute . slice . permute . slice . iota 1" $
-      compareOps $
-        let newdims =
-              [ n,
-                n,
-                n,
-                (n `P.div` 3) - 2
-              ]
-            slice33 =
-              Slice
-                [ DimSlice (n - 1) (n `P.div` 3) (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice 0 n 1
-                ]
-            ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0]
-            m = n `P.div` 3
-            slice1 =
-              Slice
-                [ DimSlice 0 n 1,
-                  DimSlice (n - 1) n (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice 1 (m - 2) (-1)
-                ]
-            ixfun' = reshape (slice ixfun slice1) newdims
-         in ixfun'
+  singleton . testCase "permute . slice . permute . slice . iota 1" . compareOps $
+    let slice33 =
+          Slice
+            [ DimSlice (n - 1) (n `P.div` 3) (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice 0 n 1
+            ]
+        ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0]
+        m = n `P.div` 3
+        slice1 =
+          Slice
+            [ DimSlice 0 n 1,
+              DimSlice (n - 1) n (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice 1 (m - 2) (-1)
+            ]
+        ixfun' = slice ixfun slice1
+     in ixfun'
 
 test_complex2 :: [TestTree]
 test_complex2 =
-  singleton $
-    testCase "reshape . permute . slice . permute . slice . iota 2" $
-      compareOps $
-        let newdims =
-              [ n,
-                n * n,
-                (n `P.div` 3) - 2
-              ]
-            slc2 =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice (n - 1) (n `P.div` 3) (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice 0 n 1
-                ]
-            ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0]
-            m = n `P.div` 3
-            slice1 =
-              Slice
-                [ DimSlice 0 n 1,
-                  DimSlice (n - 1) n (-1),
-                  DimSlice (n - 1) n (-1),
-                  DimSlice 1 (m - 2) (-1)
-                ]
-            ixfun' = reshape (slice ixfun slice1) newdims
-         in ixfun'
+  singleton . testCase "permute . slice . permute . slice . iota 2" . compareOps $
+    let slc2 =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice (n - 1) (n `P.div` 3) (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice 0 n 1
+            ]
+        ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0]
+        m = n `P.div` 3
+        slice1 =
+          Slice
+            [ DimSlice 0 n 1,
+              DimSlice (n - 1) n (-1),
+              DimSlice (n - 1) n (-1),
+              DimSlice 1 (m - 2) (-1)
+            ]
+        ixfun' = slice ixfun slice1
+     in ixfun'
 
 test_rebase1 :: [TestTree]
 test_rebase1 =
-  singleton $
-    testCase "rebase 1" $
-      compareOps $
-        let slice_base =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice 2 (n - 2) 1,
-                  DimSlice 3 (n - 3) 1
-                ]
-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
-            ixfn_orig = permute (iota [n - 3, n - 2]) [1, 0]
-            ixfn_rebase = rebase ixfn_base ixfn_orig
-         in ixfn_rebase
+  singleton . testCase "rebase 1" . compareOpsFailure $
+    let slice_base =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice 2 (n - 2) 1,
+              DimSlice 3 (n - 3) 1
+            ]
+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
+        ixfn_orig = permute (iota [n - 3, n - 2]) [1, 0]
+        ixfn_rebase = rebase ixfn_base ixfn_orig
+     in ixfn_rebase
 
 test_rebase2 :: [TestTree]
 test_rebase2 =
-  singleton $
-    testCase "rebase 2" $
-      compareOps $
-        let slice_base =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice (n - 1) (n - 2) (-1),
-                  DimSlice (n - 1) (n - 3) (-1)
-                ]
-            slice_orig =
-              Slice
-                [ DimSlice (n - 4) (n - 3) (-1),
-                  DimSlice (n - 3) (n - 2) (-1)
-                ]
-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
-            ixfn_orig = permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0]
-            ixfn_rebase = rebase ixfn_base ixfn_orig
-         in ixfn_rebase
+  singleton . testCase "rebase 2" . compareOpsFailure $
+    let slice_base =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice (n - 1) (n - 2) (-1),
+              DimSlice (n - 1) (n - 3) (-1)
+            ]
+        slice_orig =
+          Slice
+            [ DimSlice (n - 4) (n - 3) (-1),
+              DimSlice (n - 3) (n - 2) (-1)
+            ]
+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
+        ixfn_orig = permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0]
+        ixfn_rebase = rebase ixfn_base ixfn_orig
+     in ixfn_rebase
 
 test_rebase3 :: [TestTree]
 test_rebase3 =
-  singleton $
-    testCase "rebase full orig but not monotonic" $
-      compareOps $
-        let n2 = (n - 2) `P.div` 3
-            n3 = (n - 3) `P.div` 2
-            slice_base =
-              Slice
-                [ DimFix (n `P.div` 2),
-                  DimSlice (n - 1) n2 (-3),
-                  DimSlice (n - 1) n3 (-2)
-                ]
-            slice_orig =
-              Slice
-                [ DimSlice (n3 - 1) n3 (-1),
-                  DimSlice (n2 - 1) n2 (-1)
-                ]
-            ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
-            ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]
-            ixfn_rebase = rebase ixfn_base ixfn_orig
-         in ixfn_rebase
+  singleton . testCase "rebase full orig but not monotonic" . compareOpsFailure $
+    let n2 = (n - 2) `P.div` 3
+        n3 = (n - 3) `P.div` 2
+        slice_base =
+          Slice
+            [ DimFix (n `P.div` 2),
+              DimSlice (n - 1) n2 (-3),
+              DimSlice (n - 1) n3 (-2)
+            ]
+        slice_orig =
+          Slice
+            [ DimSlice (n3 - 1) n3 (-1),
+              DimSlice (n2 - 1) n2 (-1)
+            ]
+        ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
+        ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]
+        ixfn_rebase = rebase ixfn_base ixfn_orig
+     in ixfn_rebase
 
 test_rebase4_5 :: [TestTree]
 test_rebase4_5 =
@@ -335,54 +306,67 @@
           ]
       ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0]
       ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0]
-   in [ testCase "rebase mixed monotonicities" $
-          compareOps $
-            rebase ixfn_base ixfn_orig
+   in [ testCase "rebase mixed monotonicities" . compareOpsFailure $
+          rebase ixfn_base ixfn_orig
       ]
 
+-- Imitates a case from memory expansion.
+test_rebase6 :: [TestTree]
+test_rebase6 =
+  [ testCase "rebase . slice1 . iota" . compareOps $
+      rebase
+        (slice (iota [n, n, n]) (Slice [DimSlice 0 n 1, DimSlice 0 n 1]))
+        ( slice
+            (iota [n, n])
+            (Slice [DimSlice 1 (n - 1) 1, DimSlice 0 n 1])
+        )
+  ]
+
+-- Imitates another case from memory expansion.
+test_rebase7 :: [TestTree]
+test_rebase7 =
+  [ testCase "rebase . slice2 . iota" . compareOps $
+      rebase
+        (slice (iota [n, n, n]) (Slice [DimSlice 0 n 1, DimSlice 0 n 1]))
+        ( slice
+            (iota [n, n])
+            (Slice [DimSlice 0 (n - 1) 1, DimSlice 0 n 1])
+        )
+  ]
+
 test_flatSlice_iota :: [TestTree]
 test_flatSlice_iota =
-  singleton $
-    testCase "flatSlice . iota" $
-      compareOps $
-        flatSlice (iota [n * n * n * n]) $
-          FlatSlice 2 [FlatDimIndex (n * 2) 4, FlatDimIndex n 3, FlatDimIndex 1 2]
+  singleton . testCase "flatSlice . iota" . compareOps $
+    flatSlice (iota [n * n * n * n]) $
+      FlatSlice 2 [FlatDimIndex (n * 2) 4, FlatDimIndex n 3, FlatDimIndex 1 2]
 
 test_slice_flatSlice_iota :: [TestTree]
 test_slice_flatSlice_iota =
-  singleton $
-    testCase "slice . flatSlice . iota " $
-      compareOps $
-        slice (flatSlice (iota [2 + n * n * n]) flat_slice) $
-          Slice [DimFix 2, DimSlice 0 n 1, DimFix 0]
+  singleton . testCase "slice . flatSlice . iota " . compareOps $
+    slice (flatSlice (iota [2 + n * n * n]) flat_slice) $
+      Slice [DimFix 2, DimSlice 0 n 1, DimFix 0]
   where
     flat_slice = FlatSlice 2 [FlatDimIndex (n * n) 1, FlatDimIndex n 1, FlatDimIndex 1 1]
 
 test_flatSlice_flatSlice_iota :: [TestTree]
 test_flatSlice_flatSlice_iota =
-  singleton $
-    testCase "flatSlice . flatSlice . iota " $
-      compareOps $
-        flatSlice (flatSlice (iota [10 * 10]) flat_slice_1) flat_slice_2
+  singleton . testCase "flatSlice . flatSlice . iota " . compareOps $
+    flatSlice (flatSlice (iota [10 * 10]) flat_slice_1) flat_slice_2
   where
     flat_slice_1 = FlatSlice 17 [FlatDimIndex 3 27, FlatDimIndex 3 10, FlatDimIndex 3 1]
     flat_slice_2 = FlatSlice 2 [FlatDimIndex 2 (-2)]
 
 test_flatSlice_slice_iota :: [TestTree]
 test_flatSlice_slice_iota =
-  singleton $
-    testCase "flatSlice . slice . iota " $
-      compareOps $
-        flatSlice (slice (iota [210, 100]) $ Slice [DimSlice 10 100 2, DimFix 10]) flat_slice_1
+  singleton . testCase "flatSlice . slice . iota " . compareOps $
+    flatSlice (slice (iota [210, 100]) $ Slice [DimSlice 10 100 2, DimFix 10]) flat_slice_1
   where
     flat_slice_1 = FlatSlice 17 [FlatDimIndex 3 27, FlatDimIndex 3 10, FlatDimIndex 3 1]
 
 test_flatSlice_transpose_slice_iota :: [TestTree]
 test_flatSlice_transpose_slice_iota =
-  singleton $
-    testCase "flatSlice . transpose . slice . iota " $
-      compareOps $
-        flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1
+  singleton . testCase "flatSlice . transpose . slice . iota " . compareOpsFailure $
+    flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1
   where
     flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2]
 
@@ -419,15 +403,15 @@
 --             lm1 =
 --               IxFunLMAD.LMAD
 --                 256
---                 [ IxFunLMAD.LMADDim 256 0 (sub64 (num_blocks_8284) 1) 0 IxFunLMAD.Inc,
---                   IxFunLMAD.LMADDim 1 0 16 1 IxFunLMAD.Inc,
---                   IxFunLMAD.LMADDim 16 0 16 2 IxFunLMAD.Inc
+--                 [ IxFunLMAD.LMADDim 256 0 (sub64 (num_blocks_8284) 1) 0 ,
+--                   IxFunLMAD.LMADDim 1 0 16 1 ,
+--                   IxFunLMAD.LMADDim 16 0 16 2
 --                 ]
 --             lm2 :: IxFunLMAD.LMAD (TPrimExp Int64 VName)
 --             lm2 =
 --               IxFunLMAD.LMAD
 --                 (add_nw64 (add_nw64 (add_nw64 (add_nw64 (mul_nw64 (256) (num_blocks_8284)) (256)) (mul_nw64 (gtid_8472) (mul_nw64 (256) (num_blocks_8284)))) (mul_nw64 (gtid_8473) (256))) (mul_nw64 (gtid_8474) (16)))
---                 [IxFunLMAD.LMADDim 1 0 16 0 IxFunLMAD.Inc]
+--                 [IxFunLMAD.LMADDim 1 0 16 0 ]
 --          in testCase (pretty lm1 <> " and " <> pretty lm2) $ IxFunLMAD.disjoint2 lessthans nonnegs lm1 lm2 @? "Failed"
 --       ]
 
@@ -480,37 +464,37 @@
               lm1 =
                 IxFunLMAD.LMAD
                   (add_nw64 (mul64 block_size_12121 i_12214) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (add64 1 i_12214) gtid_12553) 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (add64 1 i_12214) gtid_12553) 1) 0,
+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1
                   ]
 
               lm2 =
                 IxFunLMAD.LMAD
                   (block_size_12121 * i_12214)
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,
+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1
                   ]
 
               lm_w =
                 IxFunLMAD.LMAD
                   (add_nw64 (add64 (add64 1 n_blab) (mul64 block_size_12121 i_12214)) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))
-                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0,
+                    IxFunLMAD.LMADDim 1 block_size_12121 1
                   ]
 
               lm_blocks =
                 IxFunLMAD.LMAD
                   (block_size_12121 * i_12214 + n_blab + 1)
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 block_size_12121 2 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0,
+                    IxFunLMAD.LMADDim n_blab block_size_12121 1,
+                    IxFunLMAD.LMADDim 1 block_size_12121 2
                   ]
 
               lm_lower_per =
                 IxFunLMAD.LMAD
                   (block_size_12121 * i_12214)
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (i_12214 + 1) 0,
+                    IxFunLMAD.LMADDim 1 (block_size_12121 + 1) 1
                   ]
 
               res1 = disjointTester asserts lessthans lm1 lm_w
@@ -539,36 +523,36 @@
               lm1 =
                 IxFunLMAD.LMAD
                   (add_nw64 (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0,
+                    IxFunLMAD.LMADDim n_blab block_size_12121 1
                   ]
 
               lm2 =
                 IxFunLMAD.LMAD
                   (add_nw64 (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1) (mul_nw64 (add_nw64 gtid_12553 1) (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) (sub_nw64 (sub_nw64 (sub64 (sub64 (sdiv64 (sub64 n_blab 1) block_size_12121) i_12214) 1) gtid_12553) 1) 0,
+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1
                   ]
 
               lm3 =
                 IxFunLMAD.LMAD
                   (add64 n_blab (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1))
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim n_blab block_size_12121 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,
+                    IxFunLMAD.LMADDim n_blab block_size_12121 1
                   ]
 
               lm4 =
                 IxFunLMAD.LMAD
                   (sub64 (sub64 (mul64 n_blab (add64 1 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) 1)
-                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (add_nw64 (mul_nw64 block_size_12121 n_blab) (mul_nw64 (-1) block_size_12121)) gtid_12553 0,
+                    IxFunLMAD.LMADDim 1 (1 + block_size_12121) 1
                   ]
 
               lm_w =
                 IxFunLMAD.LMAD
                   (add_nw64 (sub64 (mul64 n_blab (add64 2 (mul64 block_size_12121 (add64 1 i_12214)))) block_size_12121) (mul_nw64 gtid_12553 (sub64 (mul64 block_size_12121 n_blab) block_size_12121)))
-                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 block_size_12121 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim n_blab block_size_12121 0,
+                    IxFunLMAD.LMADDim 1 block_size_12121 1
                   ]
 
               res1 = disjointTester asserts lessthans lm1 lm_w
@@ -589,29 +573,29 @@
               lm1 =
                 IxFunLMAD.LMAD
                   (1024 * num_blocks * (1 + step) + 1024 * step)
-                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 32 32 1 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 32 2 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0,
+                    IxFunLMAD.LMADDim 32 32 1,
+                    IxFunLMAD.LMADDim 1 32 2
                   ]
 
               lm_w1 =
                 IxFunLMAD.LMAD
                   (1024 * num_blocks * step + 1024 * step)
-                  [ IxFunLMAD.LMADDim 32 32 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 32 1 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim 32 32 0,
+                    IxFunLMAD.LMADDim 1 32 1
                   ]
 
               lm_w2 =
                 IxFunLMAD.LMAD
                   ((1 + step) * 1024 * num_blocks + (1 + step) * 1024)
-                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1024 (num_blocks - step - 1) 1 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1024 1 2 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 32 1 3 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 128 8 4 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 4 8 5 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 32 4 6 IxFunLMAD.Inc,
-                    IxFunLMAD.LMADDim 1 4 7 IxFunLMAD.Inc
+                  [ IxFunLMAD.LMADDim (1024 * num_blocks) (num_blocks - step - 1) 0,
+                    IxFunLMAD.LMADDim 1024 (num_blocks - step - 1) 1,
+                    IxFunLMAD.LMADDim 1024 1 2,
+                    IxFunLMAD.LMADDim 32 1 3,
+                    IxFunLMAD.LMADDim 128 8 4,
+                    IxFunLMAD.LMADDim 4 8 5,
+                    IxFunLMAD.LMADDim 32 4 6,
+                    IxFunLMAD.LMADDim 1 4 7
                   ]
 
               asserts =
diff --git a/unittests/Futhark/IR/Mem/IxFunWrapper.hs b/unittests/Futhark/IR/Mem/IxFunWrapper.hs
--- a/unittests/Futhark/IR/Mem/IxFunWrapper.hs
+++ b/unittests/Futhark/IR/Mem/IxFunWrapper.hs
@@ -12,6 +12,7 @@
   )
 where
 
+import Control.Monad (join)
 import Futhark.IR.Mem.IxFun qualified as I
 import Futhark.IR.Mem.IxFun.Alg qualified as IA
 import Futhark.IR.Syntax (FlatSlice, Slice)
@@ -21,52 +22,52 @@
 
 type Permutation = [Int]
 
-type IxFun num = (I.IxFun num, IA.IxFun num)
+type IxFun num = (Maybe (I.IxFun num), IA.IxFun num)
 
 iota ::
   IntegralExp num =>
   Shape num ->
   IxFun num
-iota x = (I.iota x, IA.iota x)
+iota x = (Just $ I.iota x, IA.iota x)
 
 permute ::
   IntegralExp num =>
   IxFun num ->
   Permutation ->
   IxFun num
-permute (l, a) x = (I.permute l x, IA.permute a x)
+permute (l, a) x = (I.permute <$> l <*> pure x, IA.permute a x)
 
 reshape ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   Shape num ->
   IxFun num
-reshape (l, a) x = (I.reshape l x, IA.reshape a x)
+reshape (l, a) x = (join (I.reshape <$> l <*> pure x), IA.reshape a x)
 
 coerce ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   Shape num ->
   IxFun num
-coerce (l, a) x = (I.coerce l x, IA.coerce a x)
+coerce (l, a) x = (I.coerce <$> l <*> pure x, IA.coerce a x)
 
 slice ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   Slice num ->
   IxFun num
-slice (l, a) x = (I.slice l x, IA.slice a x)
+slice (l, a) x = (I.slice <$> l <*> pure x, IA.slice a x)
 
 flatSlice ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   FlatSlice num ->
   IxFun num
-flatSlice (l, a) x = (I.flatSlice l x, IA.flatSlice a x)
+flatSlice (l, a) x = (join (I.flatSlice <$> l <*> pure x), IA.flatSlice a x)
 
 rebase ::
   (Eq num, IntegralExp num) =>
   IxFun num ->
   IxFun num ->
   IxFun num
-rebase (l, a) (l1, a1) = (I.rebase l l1, IA.rebase a a1)
+rebase (l, a) (l1, a1) = (join (I.rebase <$> l <*> l1), IA.rebase a a1)