diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -2,10 +2,10 @@
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: befc35fd82c38f76939296d95942a59561978305f93fb1599eb91ea20a6af5fe
+-- hash: eb85e7d5c7284742fdf0fda53a17074c98dc9526652f211fdabddb727e2beffa
 
 name:           futhark
-version:        0.10.2
+version:        0.11.1
 synopsis:       An optimising compiler for a functional, array-oriented language.
 description:    Futhark is a small programming language designed to be compiled to
                 efficient parallel code. It is a statically typed, data-parallel,
@@ -59,58 +59,6 @@
   location: https://github.com/diku-dk/futhark
 
 library
-  hs-source-dirs:
-      src
-  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists
-  build-depends:
-      aeson
-    , ansi-terminal >=0.6.3.1
-    , array >=0.4
-    , base >=4 && <5
-    , binary >=0.8.3
-    , blaze-html >=0.9.0.1
-    , bytestring >=0.10.8
-    , containers >=0.5
-    , data-binary-ieee754 >=0.1
-    , directory >=1.3.0.0
-    , directory-tree >=0.12.1
-    , dlist >=0.6.0.1
-    , file-embed >=0.0.9
-    , filepath >=1.4.1.1
-    , free >=4.12.4
-    , gitrev >=1.2.0
-    , haskeline
-    , http-client >=0.5.7.0
-    , http-client-tls >=0.3.5.1
-    , http-conduit >=2.2.4
-    , language-c-quote >=0.12
-    , mainland-pretty >=0.6.1
-    , markdown >=0.1.16
-    , megaparsec >=7.0.1
-    , mtl >=2.2.1
-    , neat-interpolation >=0.3
-    , parallel >=3.2.1.0
-    , parser-combinators >=1.0.0
-    , process >=1.4.3.0
-    , process-extras >=0.7.2
-    , random
-    , regex-tdfa >=1.2
-    , srcloc >=0.4
-    , template-haskell >=2.11.1
-    , temporary
-    , terminal-size >=0.3
-    , text >=1.2.2.2
-    , time >=1.6.0.1
-    , transformers >=0.3
-    , utf8-string >=1
-    , vector >=0.12
-    , vector-binary-instances >=0.2.2.0
-    , versions >=3.3.1
-    , zip-archive >=0.3.1.1
-    , zlib >=0.6.1.2
-  build-tools:
-      alex
-    , happy
   exposed-modules:
       Futhark.Actions
       Futhark.Analysis.AlgSimplify
@@ -179,7 +127,9 @@
       Futhark.CodeGen.ImpGen.Kernels
       Futhark.CodeGen.ImpGen.Kernels.Base
       Futhark.CodeGen.ImpGen.Kernels.SegGenRed
+      Futhark.CodeGen.ImpGen.Kernels.SegMap
       Futhark.CodeGen.ImpGen.Kernels.SegRed
+      Futhark.CodeGen.ImpGen.Kernels.SegScan
       Futhark.CodeGen.ImpGen.Kernels.ToOpenCL
       Futhark.CodeGen.ImpGen.Kernels.Transpose
       Futhark.CodeGen.ImpGen.OpenCL
@@ -233,7 +183,6 @@
       Futhark.Pass.ExtractKernels.Intragroup
       Futhark.Pass.ExtractKernels.ISRWIM
       Futhark.Pass.ExtractKernels.Kernelise
-      Futhark.Pass.ExtractKernels.Segmented
       Futhark.Pass.FirstOrderTransform
       Futhark.Pass.KernelBabysitting
       Futhark.Pass.ResolveAssertions
@@ -310,37 +259,74 @@
       Language.Futhark.Parser.Parser
       Language.Futhark.Parser.Lexer
       Paths_futhark
+  hs-source-dirs:
+      src
+  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists
+  build-depends:
+      aeson >=1.0.0.0
+    , ansi-terminal >=0.6.3.1
+    , array >=0.4
+    , base >=4 && <5
+    , binary >=0.8.3
+    , blaze-html >=0.9.0.1
+    , bytestring >=0.10.8
+    , containers >=0.5
+    , data-binary-ieee754 >=0.1
+    , directory >=1.3.0.0
+    , directory-tree >=0.12.1
+    , dlist >=0.6.0.1
+    , file-embed >=0.0.9
+    , filepath >=1.4.1.1
+    , free >=4.12.4
+    , gitrev >=1.2.0
+    , haskeline
+    , http-client >=0.5.7.0
+    , http-client-tls >=0.3.5.1
+    , http-conduit >=2.2.4
+    , language-c-quote >=0.12
+    , mainland-pretty >=0.6.1
+    , markdown >=0.1.16
+    , megaparsec >=7.0.1
+    , mtl >=2.2.1
+    , neat-interpolation >=0.3
+    , parallel >=3.2.1.0
+    , parser-combinators >=1.0.0
+    , process >=1.4.3.0
+    , process-extras >=0.7.2
+    , random
+    , regex-tdfa >=1.2
+    , srcloc >=0.4
+    , template-haskell >=2.11.1
+    , temporary
+    , terminal-size >=0.3
+    , text >=1.2.2.2
+    , time >=1.6.0.1
+    , transformers >=0.3
+    , utf8-string >=1
+    , vector >=0.12
+    , vector-binary-instances >=0.2.2.0
+    , versions >=3.3.1
+    , zip-archive >=0.3.1.1
+    , zlib >=0.6.1.2
+  build-tools:
+      alex
+    , happy
   default-language: Haskell2010
 
 executable futhark
   main-is: src/futhark.hs
+  other-modules:
+      Paths_futhark
   ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg"
   build-depends:
       base
     , futhark
     , text
-  other-modules:
-      Paths_futhark
   default-language: Haskell2010
 
 test-suite unit
   type: exitcode-stdio-1.0
   main-is: futhark_tests.hs
-  hs-source-dirs:
-      unittests
-  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists
-  build-depends:
-      QuickCheck >=2.8
-    , base
-    , containers
-    , futhark
-    , megaparsec
-    , mtl
-    , parser-combinators
-    , tasty
-    , tasty-hunit
-    , tasty-quickcheck
-    , text
   other-modules:
       Futhark.Analysis.ScalExpTests
       Futhark.Optimise.AlgSimplifyTests
@@ -357,4 +343,19 @@
       Language.Futhark.CoreTests
       Language.Futhark.SyntaxTests
       Paths_futhark
+  hs-source-dirs:
+      unittests
+  ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists
+  build-depends:
+      QuickCheck >=2.8
+    , base
+    , containers
+    , futhark
+    , megaparsec
+    , mtl
+    , parser-combinators
+    , tasty
+    , tasty-hunit
+    , tasty-quickcheck
+    , text
   default-language: Haskell2010
diff --git a/futlib/math.fut b/futlib/math.fut
--- a/futlib/math.fut
+++ b/futlib/math.fut
@@ -122,6 +122,8 @@
   val acos: t -> t
   val atan: t -> t
   val atan2: t -> t -> t
+  val gamma: t -> t
+  val lgamma: t -> t
 
   -- | Natural logarithm.
   val log: t -> t
@@ -828,6 +830,8 @@
   let asin (x: f64) = intrinsics.asin64 x
   let atan (x: f64) = intrinsics.atan64 x
   let atan2 (x: f64) (y: f64) = intrinsics.atan2_64 x y
+  let gamma = intrinsics.gamma64
+  let lgamma = intrinsics.lgamma64
 
   let ceil (x: f64) : f64 =
     let i = to_i64 x
@@ -937,6 +941,8 @@
   let asin (x: f32) = intrinsics.asin32 x
   let atan (x: f32) = intrinsics.atan32 x
   let atan2 (x: f32) (y: f32) = intrinsics.atan2_32 x y
+  let gamma = intrinsics.gamma32
+  let lgamma = intrinsics.lgamma32
 
   let ceil (x: f32) : f32 =
     let i = to_i32 x
diff --git a/futlib/prelude.fut b/futlib/prelude.fut
--- a/futlib/prelude.fut
+++ b/futlib/prelude.fut
@@ -23,12 +23,12 @@
 let opaque 't (x: t): t =
   intrinsics.opaque x
 
--- | Semantically just identity, but in `futharki` the argument value
--- will be printed.
+-- | Semantically just identity, but when run in the interpreter, the
+-- argument value will be printed.
 let trace 't (x: t): t =
   intrinsics.trace x
 
 -- | Semantically just identity, but acts as a break point in
--- `futharki`.
+-- `futhark repl`.
 let break 't (x: t): t =
   intrinsics.break x
diff --git a/futlib/soacs.fut b/futlib/soacs.fut
--- a/futlib/soacs.fut
+++ b/futlib/soacs.fut
@@ -158,19 +158,20 @@
 
 -- | `stream_red op f as` splits `as` into chunks, applies `f` to each
 -- of these in parallel, and uses `op` (which must be associative) to
--- combine the per-chunk results into a final result.  This SOAC is
--- useful when `f` can be given a particularly work-efficient
--- sequential implementation.  Operationally, we can imagine that `as`
--- is divided among as many threads as necessary to saturate the
--- machine, with each thread operating sequentially.
+-- combine the per-chunk results into a final result.  The `i32`
+-- passed to `f` is the size of the chunk.  This SOAC is useful when
+-- `f` can be given a particularly work-efficient sequential
+-- implementation.  Operationally, we can imagine that `as` is divided
+-- among as many threads as necessary to saturate the machine, with
+-- each thread operating sequentially.
 --
--- A chunk may be empty, `f []` must produce the neutral element for
+-- A chunk may be empty, and `f 0 []` must produce the neutral element for
 -- `op`.
 --
 -- **Work:** *O(n)*
 --
 -- **Span:** *O(log(n))*
-let stream_red 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b =
+let stream_red 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b =
   intrinsics.stream_red (op, f, as)
 
 -- | As `stream_red`@term, but the chunks do not necessarily
@@ -180,7 +181,7 @@
 -- **Work:** *O(n)*
 --
 -- **Span:** *O(log(n))*
-let stream_red_per 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b =
+let stream_red_per 'a 'b (op: b -> b -> b) (f: i32 -> []a -> b) (as: []a): b =
   intrinsics.stream_red_per (op, f, as)
 
 -- | Similar to `stream_red`@term, except that each chunk must produce
@@ -190,7 +191,7 @@
 -- **Work:** *O(n)*
 --
 -- **Span:** *O(1)*
-let stream_map 'a 'b (f: []a -> []b) (as: []a): *[]b =
+let stream_map 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b =
   intrinsics.stream_map (f, as)
 
 -- | Similar to `stream_map`@term, but the chunks do not necessarily
@@ -200,7 +201,7 @@
 -- **Work:** *O(n)*
 --
 -- **Span:** *O(1)*
-let stream_map_per 'a 'b (f: []a -> []b) (as: []a): *[]b =
+let stream_map_per 'a 'b (f: i32 -> []a -> []b) (as: []a): *[]b =
   intrinsics.stream_map_per (f, as)
 
 -- | Return `true` if the given function returns `true` for all
diff --git a/rts/c/cuda.h b/rts/c/cuda.h
--- a/rts/c/cuda.h
+++ b/rts/c/cuda.h
@@ -1,4 +1,4 @@
-/* Simple CUDA runtime framework */
+// Start of cuda.h.
 
 #define CUDA_SUCCEED(x) cuda_api_succeed(x, #x, __FILE__, __LINE__)
 #define NVRTC_SUCCEED(x) nvrtc_api_succeed(x, #x, __FILE__, __LINE__)
@@ -98,6 +98,7 @@
   size_t max_grid_size;
   size_t max_tile_size;
   size_t max_threshold;
+  size_t max_shared_memory;
 
   size_t lockstep_width;
 };
@@ -193,7 +194,7 @@
     src_len += strlen(*p);
   }
 
-  char *src = malloc(src_len + 1);
+  char *src = (char*) malloc(src_len + 1);
   size_t n = 0;
   for (p = src_fragments; *p; p++) {
     strcpy(src + n, *p);
@@ -262,7 +263,7 @@
   }
 
   size_t n_opts, i = 0, i_dyn, n_opts_alloc = 20 + num_extra_opts + ctx->cfg.num_sizes;
-  const char **opts = malloc(n_opts_alloc * sizeof(const char *));
+  const char **opts = (const char**) malloc(n_opts_alloc * sizeof(const char *));
   if (!arch_set) {
     opts[i++] = "-arch";
     opts[i++] = cuda_nvrtc_get_arch(ctx->dev);
@@ -302,7 +303,7 @@
   if (res != NVRTC_SUCCESS) {
     size_t log_size;
     if (nvrtcGetProgramLogSize(prog, &log_size) == NVRTC_SUCCESS) {
-      char *log = malloc(log_size);
+      char *log = (char*) malloc(log_size);
       if (nvrtcGetProgramLog(prog, log) == NVRTC_SUCCESS) {
         fprintf(stderr,"Compilation log:\n%s\n", log);
       }
@@ -317,7 +318,7 @@
   char *ptx;
   size_t ptx_size;
   NVRTC_SUCCEED(nvrtcGetPTXSize(prog, &ptx_size));
-  ptx = malloc(ptx_size);
+  ptx = (char*) malloc(ptx_size);
   NVRTC_SUCCEED(nvrtcGetPTX(prog, ptx));
 
   NVRTC_SUCCEED(nvrtcDestroyProgram(&prog));
@@ -405,7 +406,7 @@
   len = ftell(f);
   assert(fseek(f, 0, SEEK_SET) == 0);
 
-  buf = malloc(len + 1);
+  buf = (char*) malloc(len + 1);
   assert(fread(buf, 1, len, f) == len);
   buf[len] = 0;
   *obuf = buf;
@@ -467,6 +468,7 @@
 
   free_list_init(&ctx->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);
@@ -557,3 +559,4 @@
   return CUDA_SUCCESS;
 }
 
+// End of cuda.h.
diff --git a/rts/c/free_list.h b/rts/c/free_list.h
--- a/rts/c/free_list.h
+++ b/rts/c/free_list.h
@@ -1,4 +1,4 @@
-/* Free list management */
+// Start of free_list.h.
 
 /* An entry in the free list.  May be invalid, to avoid having to
    deallocate entries as soon as they are removed.  There is also a
@@ -19,7 +19,7 @@
 void free_list_init(struct free_list *l) {
   l->capacity = 30; // Picked arbitrarily.
   l->used = 0;
-  l->entries = malloc(sizeof(struct free_list_entry) * l->capacity);
+  l->entries = (struct free_list_entry*) malloc(sizeof(struct free_list_entry) * l->capacity);
   for (int i = 0; i < l->capacity; i++) {
     l->entries[i].valid = 0;
   }
@@ -108,3 +108,4 @@
   return 1;
 }
 
+// End of free_list.h.
diff --git a/rts/c/lock.h b/rts/c/lock.h
--- a/rts/c/lock.h
+++ b/rts/c/lock.h
@@ -1,3 +1,5 @@
+// Start of lock.h.
+
 /* A very simple cross-platform implementation of locks.  Uses
    pthreads on Unix and some Windows thing there.  Futhark's
    host-level code is not multithreaded, but user code may be, so we
@@ -51,7 +53,9 @@
 
 static void free_lock(lock_t *lock) {
   /* Nothing to do for pthreads. */
-  lock = lock;
+  (void)lock;
 }
 
 #endif
+
+// End of lock.h.
diff --git a/rts/c/opencl.h b/rts/c/opencl.h
--- a/rts/c/opencl.h
+++ b/rts/c/opencl.h
@@ -1,4 +1,4 @@
-/* The simple OpenCL runtime framework used by Futhark. */
+// Start of opencl.h.
 
 #define OPENCL_SUCCEED_FATAL(e) opencl_succeed_fatal(e, #e, __FILE__, __LINE__)
 #define OPENCL_SUCCEED_NONFATAL(e) opencl_succeed_nonfatal(e, #e, __FILE__, __LINE__)
@@ -97,6 +97,7 @@
   size_t max_num_groups;
   size_t max_tile_size;
   size_t max_threshold;
+  size_t max_local_memory;
 
   size_t lockstep_width;
 };
@@ -117,7 +118,7 @@
 
 static char *strclone(const char *str) {
   size_t size = strlen(str) + 1;
-  char *copy = malloc(size);
+  char *copy = (char*) malloc(size);
   if (copy == NULL) {
     return NULL;
   }
@@ -150,7 +151,7 @@
   return s;
 }
 
-static const char* opencl_error_string(unsigned int err)
+static const char* opencl_error_string(cl_int err)
 {
     switch (err) {
         case CL_SUCCESS:                            return "Success!";
@@ -254,7 +255,7 @@
 
   OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, 0, NULL, &req_bytes));
 
-  info = malloc(req_bytes);
+  info = (char*) malloc(req_bytes);
 
   OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, req_bytes, info, NULL));
 
@@ -268,7 +269,7 @@
 
   OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, 0, NULL, &req_bytes));
 
-  info = malloc(req_bytes);
+  info = (char*) malloc(req_bytes);
 
   OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, req_bytes, info, NULL));
 
@@ -461,7 +462,7 @@
     size_t ret_val_size;
     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &ret_val_size));
 
-    build_log = malloc(ret_val_size+1);
+    build_log = (char*) malloc(ret_val_size+1);
     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL));
 
     // The spec technically does not say whether the build log is zero-terminated, so let's be careful.
@@ -531,6 +532,10 @@
 
   size_t max_tile_size = sqrt(max_group_size);
 
+  cl_ulong max_local_memory;
+  OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_LOCAL_MEM_SIZE,
+                                       sizeof(size_t), &max_local_memory, NULL));
+
   // Make sure this function is defined.
   post_opencl_setup(ctx, &device_option);
 
@@ -553,6 +558,7 @@
   ctx->max_group_size = max_group_size;
   ctx->max_tile_size = max_tile_size; // No limit.
   ctx->max_threshold = ctx->max_num_groups = 0; // No limit.
+  ctx->max_local_memory = max_local_memory;
 
   // Now we go through all the sizes, clamp them to the valid range,
   // or set them to the default.
@@ -608,7 +614,7 @@
       src_size += strlen(*src);
     }
 
-    fut_opencl_src = malloc(src_size + 1);
+    fut_opencl_src = (char*) malloc(src_size + 1);
 
     size_t n, i;
     for (i = 0, n = 0; srcs && srcs[i]; i++) {
@@ -644,7 +650,7 @@
       compile_opts_size += strlen(extra_build_opts[i] + 1);
     }
 
-    char *compile_opts = malloc(compile_opts_size);
+    char *compile_opts = (char*) malloc(compile_opts_size);
 
     int w = snprintf(compile_opts, compile_opts_size,
                      "-DLOCKSTEP_WIDTH=%d ",
@@ -687,7 +693,7 @@
     size_t binary_size;
     OPENCL_SUCCEED_FATAL(clGetProgramInfo(prog, CL_PROGRAM_BINARY_SIZES,
                                           sizeof(size_t), &binary_size, NULL));
-    unsigned char *binary = malloc(binary_size);
+    unsigned char *binary = (unsigned char*) malloc(binary_size);
     unsigned char *binaries[1] = { binary };
     OPENCL_SUCCEED_FATAL(clGetProgramInfo(prog, CL_PROGRAM_BINARIES,
                                           sizeof(unsigned char*), binaries, NULL));
@@ -803,6 +809,9 @@
   int error = opencl_alloc_actual(ctx, min_size, mem_out);
 
   while (error == CL_MEM_OBJECT_ALLOCATION_FAILURE) {
+    if (ctx->cfg.debugging) {
+      fprintf(stderr, "Out of OpenCL memory: releasing entry from the free list...\n");
+    }
     cl_mem mem;
     if (free_list_first(&ctx->free_list, &mem) == 0) {
       error = clReleaseMemObject(mem);
@@ -851,3 +860,5 @@
 
   return CL_SUCCESS;
 }
+
+// End of opencl.h.
diff --git a/rts/c/panic.h b/rts/c/panic.h
--- a/rts/c/panic.h
+++ b/rts/c/panic.h
@@ -1,4 +1,4 @@
-/* Crash and burn. */
+// Start of panic.h.
 
 #include <stdarg.h>
 
@@ -21,8 +21,10 @@
   va_list vl;
   va_start(vl, s);
   size_t needed = 1 + vsnprintf(NULL, 0, s, vl);
-  char *buffer = malloc(needed);
+  char *buffer = (char*) malloc(needed);
   va_start(vl, s); /* Must re-init. */
   vsnprintf(buffer, needed, s, vl);
   return buffer;
 }
+
+// End of panic.h.
diff --git a/rts/c/timing.h b/rts/c/timing.h
--- a/rts/c/timing.h
+++ b/rts/c/timing.h
@@ -1,8 +1,7 @@
-/* Some simple utilities for wall-clock timing.
+// Start of timing.h.
 
-   The function get_wall_time() returns the wall time in microseconds
-   (with an unspecified offset).
-*/
+// The function get_wall_time() returns the wall time in microseconds
+// (with an unspecified offset).
 
 #ifdef _WIN32
 
@@ -28,3 +27,5 @@
 }
 
 #endif
+
+// End of timing.h.
diff --git a/rts/c/tuning.h b/rts/c/tuning.h
--- a/rts/c/tuning.h
+++ b/rts/c/tuning.h
@@ -1,10 +1,10 @@
-// Start of tuning.h
+// Start of tuning.h.
 
 static char* load_tuning_file(const char *fname,
                               void *cfg,
                               int (*set_size)(void*, const char*, size_t)) {
   const int max_line_len = 1024;
-  char* line = malloc(max_line_len);
+  char* line = (char*) malloc(max_line_len);
 
   FILE *f = fopen(fname, "r");
 
@@ -37,4 +37,4 @@
   return NULL;
 }
 
-// End of tuning.h
+// End of tuning.h.
diff --git a/rts/c/values.h b/rts/c/values.h
--- a/rts/c/values.h
+++ b/rts/c/values.h
@@ -1,3 +1,5 @@
+// Start of values.h.
+
 //// Text I/O
 
 typedef int (*writer)(FILE*, void*);
@@ -831,3 +833,5 @@
     return type->write_str(out, src);
   }
 }
+
+// End of values.h.
diff --git a/rts/csharp/opencl.cs b/rts/csharp/opencl.cs
--- a/rts/csharp/opencl.cs
+++ b/rts/csharp/opencl.cs
@@ -401,6 +401,7 @@
    public int MaxNumGroups;
    public int MaxTileSize;
    public int MaxThreshold;
+   public int MaxLocalMemory;
 
    public int LockstepWidth;
 }
@@ -734,7 +735,7 @@
     ComputeErrorCode error;
     CLPlatformHandle platform;
     CLDeviceHandle device;
-    int MaxGroupSize;
+    int MaxGroupSize, MaxLocalMemory;
 
     ctx.OpenCL.LockstepWidth = 0;
 
@@ -777,6 +778,16 @@
 
     int MaxTileSize = (int) Math.Sqrt(MaxGroupSize);
 
+    unsafe
+    {
+        IntPtr throwaway1 = new IntPtr();
+        OPENCL_SUCCEED(CL10.GetDeviceInfo(device,
+                                          ComputeDeviceInfo.LocalMemorySize,
+                                          new IntPtr(sizeof(IntPtr)),
+                                          new IntPtr(&MaxLocalMemory),
+                                          out throwaway1));
+    }
+
     // Make sure this function is defined.
     PostOpenCLSetup(ref ctx, ref device_option);
 
@@ -795,6 +806,7 @@
     ctx.OpenCL.MaxGroupSize = MaxGroupSize;
     ctx.OpenCL.MaxTileSize = MaxTileSize; // No limit.
     ctx.OpenCL.MaxThreshold = ctx.OpenCL.MaxNumGroups; // No limit.
+    ctx.OpenCL.MaxLocalMemory = MaxLocalMemory;
 
     // Now we go through all the sizes, clamp them to the valid range,
     // or set them to the default.
diff --git a/rts/csharp/scalar.cs b/rts/csharp/scalar.cs
--- a/rts/csharp/scalar.cs
+++ b/rts/csharp/scalar.cs
@@ -282,6 +282,8 @@
 private static double futhark_asin64(double x){return Math.Asin(x);}
 private static double futhark_atan64(double x){return Math.Atan(x);}
 private static double futhark_atan2_64(double x, double y){return Math.Atan2(x, y);}
+private static double futhark_gamma64(double x){throw new NotImplementedException();}
+private static double futhark_lgamma64(double x){throw new NotImplementedException();}
 private static bool futhark_isnan64(double x){return double.IsNaN(x);}
 private static bool futhark_isinf64(double x){return double.IsInfinity(x);}
 private static long futhark_to_bits64(double x){return BitConverter.ToInt64(BitConverter.GetBytes(x),0);}
@@ -299,6 +301,8 @@
 private static float futhark_asin32(float x){return (float) Math.Asin(x);}
 private static float futhark_atan32(float x){return (float) Math.Atan(x);}
 private static float futhark_atan2_32(float x, float y){return (float) Math.Atan2(x, y);}
+private static float futhark_gamma32(float x){throw new NotImplementedException();}
+private static float futhark_lgamma32(float x){throw new NotImplementedException();}
 private static bool futhark_isnan32(float x){return float.IsNaN(x);}
 private static bool futhark_isinf32(float x){return float.IsInfinity(x);}
 private static int futhark_to_bits32(float x){return BitConverter.ToInt32(BitConverter.GetBytes(x), 0);}
diff --git a/rts/python/memory.py b/rts/python/memory.py
--- a/rts/python/memory.py
+++ b/rts/python/memory.py
@@ -1,4 +1,4 @@
-# Helper functions dealing with memory blocks.
+# Start of memory.py.
 
 import ctypes as ct
 
@@ -36,3 +36,5 @@
 
   def __repr__(self):
     return "<opaque Futhark value of type {}>".format(self.desc)
+
+# End of memory.py.
diff --git a/rts/python/opencl.py b/rts/python/opencl.py
--- a/rts/python/opencl.py
+++ b/rts/python/opencl.py
@@ -109,6 +109,7 @@
     self.max_tile_size = max_tile_size
     self.max_threshold = 0
     self.max_num_groups = 0
+    self.max_local_memory = int(self.device.local_mem_size)
     self.free_list = {}
 
     if 'default_group_size' in sizes:
diff --git a/rts/python/panic.py b/rts/python/panic.py
--- a/rts/python/panic.py
+++ b/rts/python/panic.py
@@ -1,4 +1,8 @@
+# Start of panic.py.
+
 def panic(exitcode, fmt, *args):
     sys.stderr.write('%s: ' % sys.argv[0])
     sys.stderr.write(fmt % args)
     sys.exit(exitcode)
+
+# End of panic.py.
diff --git a/rts/python/scalar.py b/rts/python/scalar.py
--- a/rts/python/scalar.py
+++ b/rts/python/scalar.py
@@ -1,6 +1,7 @@
-# Scalar functions.
+# Start of scalar.py.
 
 import numpy as np
+import math
 import struct
 
 def signed(x):
@@ -295,6 +296,12 @@
 def futhark_atan2_64(x, y):
   return np.arctan2(x, y)
 
+def futhark_gamma64(x):
+  return np.float64(math.gamma(x))
+
+def futhark_lgamma64(x):
+  return np.float64(math.lgamma(x))
+
 def futhark_round64(x):
   return np.round(x)
 
@@ -348,6 +355,12 @@
 def futhark_atan2_32(x, y):
   return np.arctan2(x, y)
 
+def futhark_gamma32(x):
+  return np.float32(math.gamma(x))
+
+def futhark_lgamma32(x):
+  return np.float32(math.lgamma(x))
+
 def futhark_round32(x):
   return np.round(x)
 
@@ -364,3 +377,5 @@
 def futhark_from_bits32(x):
   s = struct.pack('>l', x)
   return np.float32(struct.unpack('>f', s)[0])
+
+# End of scalar.py.
diff --git a/rts/python/tuning.py b/rts/python/tuning.py
--- a/rts/python/tuning.py
+++ b/rts/python/tuning.py
@@ -1,6 +1,4 @@
-### start of tuning.py
-###
-### Reading the .tuning file.
+# Start of tuning.py
 
 def read_tuning_file(kvs, f):
     for line in f.read().splitlines():
@@ -8,4 +6,4 @@
         kvs[size] = int(value)
     return kvs
 
-### end of tuning.py
+# End of tuning.py.
diff --git a/rts/python/values.py b/rts/python/values.py
--- a/rts/python/values.py
+++ b/rts/python/values.py
@@ -1,3 +1,5 @@
+# Start of values.py.
+
 # Hacky parser/reader/writer for values written in Futhark syntax.
 # Used for reading stdin when compiling standalone programs with the
 # Python code generator.
@@ -76,7 +78,8 @@
             read.append(c)
         return True
     except ValueError:
-        map(f.unget_char, read[::-1])
+        for c in read[::-1]:
+            f.unget_char(c)
         raise
 
 def optional(p, *args):
@@ -112,49 +115,46 @@
     s = b''
     c = f.get_char()
     while c != None:
-        if c in string.hexdigits:
+        if c in b'01234556789ABCDEFabcdef':
             s += c
             c = f.get_char()
-        elif c == '_':
+        elif c == b'_':
             c = f.get_char() # skip _
         else:
             f.unget_char(c)
             break
-    return str(int(s, 16))
-
+    return str(int(s, 16)).encode('utf8') # ugh
 
 def parse_int(f):
     s = b''
     c = f.get_char()
-    if c == b'0' and f.peek_char() in [b'x', b'X']:
+    if c == b'0' and f.peek_char() in b'xX':
         c = f.get_char() # skip X
-        s += parse_hex_int(f)
+        return parse_hex_int(f)
     else:
         while c != None:
             if c.isdigit():
                 s += c
                 c = f.get_char()
-            elif c == '_':
+            elif c == b'_':
                 c = f.get_char() # skip _
             else:
                 f.unget_char(c)
                 break
-    if len(s) == 0:
-        raise ValueError
-    return s
+        if len(s) == 0:
+            raise ValueError
+        return s
 
 def parse_int_signed(f):
     s = b''
     c = f.get_char()
 
     if c == b'-' and f.peek_char().isdigit():
-      s = c + parse_int(f)
+      return c + parse_int(f)
     else:
       if c != b'+':
           f.unget_char(c)
-      s = parse_int(f)
-
-    return s
+      return parse_int(f)
 
 def read_str_comma(f):
     skip_spaces(f)
@@ -568,7 +568,7 @@
             return read_scalar(reader, basetype)
         return (dims, basetype)
 
-def write_value(v, out=sys.stdout):
+def write_value_text(v, out=sys.stdout):
     if type(v) == np.uint8:
         out.write("%uu8" % v)
     elif type(v) == np.uint16:
@@ -625,6 +625,49 @@
     else:
         raise Exception("Cannot print value of type {}: {}".format(type(v), v))
 
-################################################################################
-### end of values.py
-################################################################################
+type_strs = { np.dtype('int8'): b'  i8',
+              np.dtype('int16'): b' i16',
+              np.dtype('int32'): b' i32',
+              np.dtype('int64'): b' i64',
+              np.dtype('uint8'): b'  u8',
+              np.dtype('uint16'): b' u16',
+              np.dtype('uint32'): b' u32',
+              np.dtype('uint64'): b' u64',
+              np.dtype('float32'): b' f32',
+              np.dtype('float64'): b' f64',
+              np.dtype('bool'): b'bool'}
+
+def construct_binary_value(v):
+    t = v.dtype
+    shape = v.shape
+
+    elems = 1
+    for d in shape:
+        elems *= d
+
+    num_bytes = 1 + 1 + 1 + 4 + len(shape) * 8 + elems * t.itemsize
+    bytes = bytearray(num_bytes)
+    bytes[0] = np.int8(ord('b'))
+    bytes[1] = 2
+    bytes[2] = np.int8(len(shape))
+    bytes[3:7] = type_strs[t]
+
+    for i in range(len(shape)):
+        bytes[7+i*8:7+(i+1)*8] = np.int64(shape[i]).tostring()
+
+    bytes[7+len(shape)*8:] = np.ascontiguousarray(v).tostring()
+
+    return bytes
+
+def write_value_binary(v, out=sys.stdout):
+    if sys.version_info >= (3,0):
+        out = out.buffer
+    out.write(construct_binary_value(v))
+
+def write_value(v, out=sys.stdout, binary=False):
+    if binary:
+        return write_value_binary(v, out=out)
+    else:
+        return write_value_text(v, out=out)
+
+# End of values.py.
diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs
--- a/src/Futhark/Analysis/Alias.hs
+++ b/src/Futhark/Analysis/Alias.hs
@@ -50,7 +50,6 @@
 analyseExp = mapExp analyse
   where analyse =
           Mapper { mapOnSubExp = return
-                 , mapOnCertificates = return
                  , mapOnVName = return
                  , mapOnBody = const $ return . analyseBody
                  , mapOnRetType = return
diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs
--- a/src/Futhark/Analysis/DataDependencies.hs
+++ b/src/Futhark/Analysis/DataDependencies.hs
@@ -40,7 +40,7 @@
           in M.unions [branchdeps, deps, tdeps, fdeps]
 
         grow deps (Let pat _ e) =
-          let free = freeIn pat <> freeInExp e
+          let free = freeIn pat <> freeIn e
               freeDeps = S.unions $ map (depsOfVar deps) $ S.toList free
           in M.fromList [ (name, freeDeps) | name <- patternNames pat ] `M.union` deps
 
diff --git a/src/Futhark/Analysis/HORepresentation/MapNest.hs b/src/Futhark/Analysis/HORepresentation/MapNest.hs
--- a/src/Futhark/Analysis/HORepresentation/MapNest.hs
+++ b/src/Futhark/Analysis/HORepresentation/MapNest.hs
@@ -13,7 +13,6 @@
   )
 where
 
-import Control.Monad
 import Data.List
 import Data.Maybe
 import qualified Data.Map.Strict as M
@@ -74,7 +73,7 @@
           -> SOAC lore
           -> m (Maybe (MapNest lore))
 
-fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm (_, []) (_, _, []) lam) inps) = do
+fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm (_, []) [] lam) inps) = do
   maybenest <- case (stmsToList $ bodyStms $ lambdaBody lam,
                      bodyResult $ lambdaBody lam) of
     ([Let pat _ e], res) | res == map Var (patternNames pat) ->
@@ -91,12 +90,11 @@
         unzip <$>
         fixInputs w (zip (map paramName $ lambdaParams lam) inps)
         (zip (params mn) inps')
-      let n' = Nesting {
-            nestingParamNames   = ps
-            , nestingResult     = patternNames pat
-            , nestingReturnType = typeOf mn
-            , nestingWidth      = inner_w
-            }
+      let n' = Nesting { nestingParamNames = ps
+                       , nestingResult     = patternNames pat
+                       , nestingReturnType = typeOf mn
+                       , nestingWidth      = inner_w
+                       }
       return $ Just $ MapNest w body' (n':ns') inps''
     -- No nested MapNest it seems.
     _ -> do
@@ -106,11 +104,11 @@
             | otherwise =
               Nothing
           boundUsedInBody =
-            mapMaybe isBound $ S.toList $ freeInLambda lam
+            mapMaybe isBound $ S.toList $ freeIn lam
       newParams <- mapM (newIdent' (++"_wasfree")) boundUsedInBody
       let subst = M.fromList $
                   zip (map identName boundUsedInBody) (map identName newParams)
-          inps' = map (substituteNames subst) inps ++
+          inps' = inps ++
                   map (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput)
                   boundUsedInBody
           lam' =
@@ -144,35 +142,16 @@
 fixInputs :: MonadFreshNames m =>
              SubExp -> [(VName, SOAC.Input)] -> [(VName, SOAC.Input)]
           -> m [(VName, SOAC.Input)]
-fixInputs w ourInps childInps =
-  reverse . snd <$> foldM inspect (ourInps, []) childInps
+fixInputs w ourInps = mapM inspect
   where
     isParam x (y, _) = x == y
 
-    findParam :: [(VName, SOAC.Input)]
-              -> VName
-              -> Maybe ((VName, SOAC.Input), [(VName, SOAC.Input)])
-    findParam remPs v
-      | ([ourP], remPs') <- partition (isParam v) remPs = Just (ourP, remPs')
-      | otherwise                                       = Nothing
-
-    inspect :: MonadFreshNames m =>
-               ([(VName, SOAC.Input)], [(VName, SOAC.Input)])
-            -> (VName, SOAC.Input)
-            -> m ([(VName, SOAC.Input)], [(VName, SOAC.Input)])
-    inspect (remPs, newInps) (_, SOAC.Input ts v _)
-      | Just ((p,pInp), remPs') <- findParam remPs v =
+    inspect (_, SOAC.Input ts v _)
+      | Just (p,pInp) <- find (isParam v) ourInps = do
           let pInp' = SOAC.transformRows ts pInp
-          in return (remPs',
-                     (p, pInp') : newInps)
-
-      | Just ((p,pInp), _) <- findParam newInps v = do
-          -- The input corresponds to a variable that has already
-          -- been used.
           p' <- newNameFromString $ baseString p
-          return (remPs, (p', pInp) : newInps)
+          return (p', pInp')
 
-    inspect (remPs, newInps) (param, SOAC.Input ts a t) = do
+    inspect (param, SOAC.Input ts a t) = do
       param' <- newNameFromString (baseString param ++ "_rep")
-      return (remPs, (param',
-                      SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t) : newInps)
+      return (param', SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t)
diff --git a/src/Futhark/Analysis/HORepresentation/SOAC.hs b/src/Futhark/Analysis/HORepresentation/SOAC.hs
--- a/src/Futhark/Analysis/HORepresentation/SOAC.hs
+++ b/src/Futhark/Analysis/HORepresentation/SOAC.hs
@@ -597,7 +597,8 @@
       return (Stream w (Sequential nes) strmlam inps,
               map paramIdent inpacc_ids)
 
-      | Just (comm, lamin, nes, _) <- Futhark.isRedomapSOAC form -> do
+      | Just (reds, _) <- Futhark.isRedomapSOAC form,
+        Futhark.Reduce comm lamin nes <- Futhark.singleReduce reds -> do
       -- Redomap(+,lam,nes,a) => is translated in strem's body to:
       -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in
       -- 2. let acc'                   = acc + acc0_ids          in
@@ -613,7 +614,8 @@
       inpacc_ids <- mapM (newParam "inpacc")  accrtps
       acc0_ids   <- mapM (newIdent "acc0"  )  accrtps
       -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in
-      let insoac = Futhark.Screma chvar (Futhark.redomapSOAC comm lamin nes foldlam) $
+      let insoac = Futhark.Screma chvar
+                   (Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam) $
                    map paramName strm_inpids
           insbnd = mkLet [] (acc0_ids++strm_resids) $ Op insoac
       -- 2. let acc'     = acc + acc0_ids    in
diff --git a/src/Futhark/Analysis/PrimExp/Convert.hs b/src/Futhark/Analysis/PrimExp/Convert.hs
--- a/src/Futhark/Analysis/PrimExp/Convert.hs
+++ b/src/Futhark/Analysis/PrimExp/Convert.hs
@@ -7,6 +7,7 @@
   , primExpFromSubExp
   , primExpFromSubExpM
   , replaceInPrimExp
+  , replaceInPrimExpM
   , substituteInPrimExp
 
     -- * Module reexport
@@ -14,6 +15,7 @@
   ) where
 
 import qualified Control.Monad.Fail as Fail
+import           Control.Monad.Identity
 import           Data.Loc
 import qualified Data.Map.Strict as M
 import           Data.Maybe
@@ -84,23 +86,30 @@
 primExpFromSubExp t (Var v)      = LeafExp v t
 primExpFromSubExp _ (Constant v) = ValueExp v
 
--- | Applying a transformation to the leaves in a 'PrimExp'.
-replaceInPrimExp :: (v -> PrimType -> PrimExp v) ->
-                    PrimExp v -> PrimExp v
-replaceInPrimExp f (LeafExp v pt) =
+-- | Applying a monadic transformation to the leaves in a 'PrimExp'.
+replaceInPrimExpM :: Monad m =>
+                     (a -> PrimType -> m (PrimExp b)) ->
+                     PrimExp a -> m (PrimExp b)
+replaceInPrimExpM f (LeafExp v pt) =
   f v pt
-replaceInPrimExp _ (ValueExp v) =
-  ValueExp v
-replaceInPrimExp f (BinOpExp bop pe1 pe2) =
-  constFoldPrimExp $ BinOpExp bop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2)
-replaceInPrimExp f (CmpOpExp cop pe1 pe2) =
-  CmpOpExp cop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2)
-replaceInPrimExp f (UnOpExp uop pe) =
-  UnOpExp uop $ replaceInPrimExp f pe
-replaceInPrimExp f (ConvOpExp cop pe) =
-  ConvOpExp cop $ replaceInPrimExp f pe
-replaceInPrimExp f (FunExp h args t) =
-  FunExp h (map (replaceInPrimExp f) args) t
+replaceInPrimExpM _ (ValueExp v) =
+  return $ ValueExp v
+replaceInPrimExpM f (BinOpExp bop pe1 pe2) =
+  constFoldPrimExp <$>
+  (BinOpExp bop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2)
+replaceInPrimExpM f (CmpOpExp cop pe1 pe2) =
+  CmpOpExp cop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2
+replaceInPrimExpM f (UnOpExp uop pe) =
+  UnOpExp uop <$> replaceInPrimExpM f pe
+replaceInPrimExpM f (ConvOpExp cop pe) =
+  ConvOpExp cop <$> replaceInPrimExpM f pe
+replaceInPrimExpM f (FunExp h args t) =
+  FunExp h <$> mapM (replaceInPrimExpM f) args <*> pure t
+
+replaceInPrimExp :: (a -> PrimType -> PrimExp b) ->
+                    PrimExp a -> PrimExp b
+replaceInPrimExp f e = runIdentity $ replaceInPrimExpM f' e
+  where f' x y = return $ f x y
 
 -- | Substituting names in a PrimExp with other PrimExps
 substituteInPrimExp :: Ord v => M.Map v (PrimExp v)
diff --git a/src/Futhark/Analysis/Range.hs b/src/Futhark/Analysis/Range.hs
--- a/src/Futhark/Analysis/Range.hs
+++ b/src/Futhark/Analysis/Range.hs
@@ -66,15 +66,14 @@
 analyseExp = mapExpM analyse
   where analyse =
           Mapper { mapOnSubExp = return
-                    , mapOnCertificates = return
-                    , mapOnVName = return
-                    , mapOnBody = const analyseBody
-                    , mapOnRetType = return
-                    , mapOnBranchType = return
-                    , mapOnFParam = return
-                    , mapOnLParam = return
-                    , mapOnOp = return . addOpRanges
-                    }
+                 , mapOnVName = return
+                 , mapOnBody = const analyseBody
+                 , mapOnRetType = return
+                 , mapOnBranchType = return
+                 , mapOnFParam = return
+                 , mapOnLParam = return
+                 , mapOnOp = return . addOpRanges
+                 }
 
 analyseLambda :: (Attributes lore, CanBeRanged (Op lore)) =>
                  Lambda lore
diff --git a/src/Futhark/Analysis/ScalExp.hs b/src/Futhark/Analysis/ScalExp.hs
--- a/src/Futhark/Analysis/ScalExp.hs
+++ b/src/Futhark/Analysis/ScalExp.hs
@@ -284,6 +284,8 @@
                    , (Mul t, STimes)
                    , (AST.SDiv t, SDiv)
                    , (AST.Pow t, SPow)
+                   , (AST.SMax t, \x y -> MaxMin False [x,y])
+                   , (AST.SMin t, \x y -> MaxMin True [x,y])
                    ]
 
 instance FreeIn ScalExp where
diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs
--- a/src/Futhark/Analysis/SymbolTable.hs
+++ b/src/Futhark/Analysis/SymbolTable.hs
@@ -566,7 +566,7 @@
                 AST.FParam lore
              -> SymbolTable lore
              -> SymbolTable lore
-insertFParam fparam = insertEntry name entry
+insertFParam fparam = flip (foldr (`isAtLeast` 0)) sizes . insertEntry name entry
   where name = AST.paramName fparam
         entry = FParam FParamEntry { fparamRange = (Nothing, Nothing)
                                    , fparamAttr = AST.paramAttr fparam
@@ -574,12 +574,13 @@
                                    , fparamStmDepth = 0
                                    , fparamConsumed = False
                                    }
+        sizes = subExpVars $ arrayDims $ AST.paramType fparam
 
 insertFParams :: Attributes lore =>
                  [AST.FParam lore]
               -> SymbolTable lore
               -> SymbolTable lore
-insertFParams fparams symtable = foldr insertFParam symtable fparams
+insertFParams fparams symtable = foldl' (flip insertFParam) symtable fparams
 
 insertLParamWithRange :: Attributes lore =>
                          LParam lore -> ScalExpRange -> IndexArray -> SymbolTable lore
diff --git a/src/Futhark/Analysis/Usage.hs b/src/Futhark/Analysis/Usage.hs
--- a/src/Futhark/Analysis/Usage.hs
+++ b/src/Futhark/Analysis/Usage.hs
@@ -1,12 +1,5 @@
 {-# LANGUAGE FlexibleContexts #-}
-module Futhark.Analysis.Usage
-       ( usageInStm
-       , usageInExp
-       , usageInLambda
-
-       , UsageInOp(..)
-       )
-       where
+module Futhark.Analysis.Usage ( usageInStm ) where
 
 import Data.Foldable
 import qualified Data.Set as S
@@ -15,13 +8,12 @@
 import Futhark.Representation.AST.Attributes.Aliases
 import qualified Futhark.Analysis.UsageTable as UT
 
-usageInStm :: (Attributes lore, Aliased lore, UsageInOp (Op lore)) =>
-              Stm lore -> UT.UsageTable
+usageInStm :: (Attributes lore, Aliased lore) => Stm lore -> UT.UsageTable
 usageInStm (Let pat lore e) =
   mconcat [usageInPat,
            usageInExpLore,
            usageInExp e,
-           UT.usages (freeInExp e)]
+           UT.usages (freeIn e)]
   where usageInPat =
           UT.usages (mconcat (map freeIn $ patternElements pat)
                      `S.difference`
@@ -29,7 +21,7 @@
         usageInExpLore =
           UT.usages $ freeIn lore
 
-usageInExp :: (Aliased lore, UsageInOp (Op lore)) => Exp lore -> UT.UsageTable
+usageInExp :: Aliased lore => Exp lore -> UT.UsageTable
 usageInExp (Apply _ args _ _) =
   mconcat [ mconcat $ map UT.consumedUsage $
             S.toList $ subExpAliases arg
@@ -44,22 +36,5 @@
 usageInExp (BasicOp (Update src _ _)) =
   UT.consumedUsage src
 usageInExp (Op op) =
-  mconcat $ usageInOp op : map UT.consumedUsage (S.toList $ consumedInOp op)
+  mconcat $ map UT.consumedUsage (S.toList $ consumedInOp op)
 usageInExp _ = UT.empty
-
-class UsageInOp op where
-  usageInOp :: op -> UT.UsageTable
-
-instance UsageInOp () where
-  usageInOp () = mempty
-
-usageInLambda :: Aliased lore =>
-                 Lambda lore -> [VName] -> UT.UsageTable
-usageInLambda lam arrs =
-  mconcat $
-  map (UT.consumedUsage . snd) $
-  filter ((`S.member` consumed_in_body) . fst) $
-  zip (map paramName arr_params) arrs
-  where arr_params = snd $ splitAt n $ lambdaParams lam
-        consumed_in_body = consumedInBody $ lambdaBody lam
-        n = length arrs
diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs
--- a/src/Futhark/Analysis/UsageTable.hs
+++ b/src/Futhark/Analysis/UsageTable.hs
@@ -19,7 +19,6 @@
   , consumedUsage
   , inResultUsage
   , Usages
-  , leftScope
   )
   where
 
@@ -131,6 +130,3 @@
 -- | x - y, but for Usages.
 withoutU :: Usages -> Usages -> Usages
 withoutU (Usages x) (Usages y) = Usages $ x .&. complement y
-
-leftScope :: UsageTable -> UsageTable
-leftScope (UsageTable table) = UsageTable $ M.map (`withoutU` inResultU) table
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
@@ -30,7 +30,7 @@
              liftIO $ writeFile cpath impl
            ToExecutable -> do
              liftIO $ writeFile cpath $ CCUDA.asExecutable cprog
-             let args = [cpath, "-O3", "-std=c99", "-lm", "-o", outpath]
+             let args = [cpath, "-O", "-std=c99", "-lm", "-o", outpath]
                         ++ extra_options
              ret <- liftIO $ runProgramWithExitCode "gcc" args ""
              case ret of
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
@@ -37,7 +37,7 @@
            ToExecutable -> do
              liftIO $ writeFile cpath $ COpenCL.asExecutable cprog
              ret <- liftIO $ runProgramWithExitCode "gcc"
-                    ([cpath, "-O3", "-std=c99", "-lm", "-o", outpath] ++ extra_options) ""
+                    ([cpath, "-O", "-std=c99", "-lm", "-o", outpath] ++ extra_options) ""
              case ret of
                Left err ->
                  externalErrorS $ "Failed to run gcc: " ++ show err
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
@@ -30,7 +30,7 @@
       let extra = generateBoilerplate cuda_code cuda_prelude
                                       kernel_names sizes
       in Right <$> GC.compileProg operations extra cuda_includes
-                   [Space "device", Space "local", DefaultSpace] cliOptions prog'
+                   [Space "device", DefaultSpace] cliOptions prog'
   where
     operations :: GC.Operations OpenCL ()
     operations = GC.Operations
@@ -106,7 +106,7 @@
   val' <- newVName "write_tmp"
   GC.stm [C.cstm|{$ty:t $id:val' = $exp:val;
                   CUDA_SUCCEED(
-                    cuMemcpyHtoD($exp:mem + $exp:idx,
+                    cuMemcpyHtoD($exp:mem + $exp:idx * sizeof($ty:t),
                                  &$id:val',
                                  sizeof($ty:t)));
                  }|]
@@ -119,7 +119,7 @@
   GC.decl [C.cdecl|$ty:t $id:val;|]
   GC.stm [C.cstm|CUDA_SUCCEED(
                    cuMemcpyDtoH(&$id:val,
-                                $exp:mem + $exp:idx,
+                                $exp:mem + $exp:idx * sizeof($ty:t),
                                 sizeof($ty:t)));
                 |]
   return [C.cexp|$id:val|]
@@ -129,14 +129,12 @@
 allocateCUDABuffer :: GC.Allocate OpenCL ()
 allocateCUDABuffer mem size tag "device" =
   GC.stm [C.cstm|CUDA_SUCCEED(cuda_alloc(&ctx->cuda, $exp:size, $exp:tag, &$exp:mem));|]
-allocateCUDABuffer _ _ _ "local" = return ()
 allocateCUDABuffer _ _ _ space =
   fail $ "Cannot allocate in '" ++ space ++ "' memory space."
 
 deallocateCUDABuffer :: GC.Deallocate OpenCL ()
 deallocateCUDABuffer mem tag "device" =
   GC.stm [C.cstm|CUDA_SUCCEED(cuda_free(&ctx->cuda, $exp:mem, $exp:tag));|]
-deallocateCUDABuffer _ _ "local" = return ()
 deallocateCUDABuffer _ _ space =
   fail $ "Cannot deallocate in '" ++ space ++ "' memory space."
 
@@ -187,7 +185,6 @@
 
 cudaMemoryType :: GC.MemoryType OpenCL ()
 cudaMemoryType "device" = return [C.cty|typename CUdeviceptr|]
-cudaMemoryType "local" = pure [C.cty|unsigned char|] -- dummy type
 cudaMemoryType space =
   fail $ "CUDA backend does not support '" ++ space ++ "' memory space."
 
@@ -206,6 +203,7 @@
     cudaSizeClass SizeGroup = "block_size"
     cudaSizeClass SizeNumGroups = "grid_size"
     cudaSizeClass SizeTile = "tile_size"
+    cudaSizeClass SizeLocalMemory = "shared_memory"
 callKernel (LaunchKernel name args num_blocks block_size) = do
   args_arr <- newVName "kernel_args"
   time_start <- newVName "time_start"
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
@@ -86,13 +86,13 @@
   GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->
     ([C.cedecl|struct $id:cfg* $id:s(void);|],
      [C.cedecl|struct $id:cfg* $id:s(void) {
-                         struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));
+                         struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));
                          if (cfg == NULL) {
                            return NULL;
                          }
 
                          cfg->num_nvrtc_opts = 0;
-                         cfg->nvrtc_opts = malloc(sizeof(const char*));
+                         cfg->nvrtc_opts = (const char**) malloc(sizeof(const char*));
                          cfg->nvrtc_opts[0] = NULL;
                          $stms:size_value_inits
                          cuda_config_init(&cfg->cu_cfg, $int:num_sizes,
@@ -113,7 +113,7 @@
      [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) {
                          cfg->nvrtc_opts[cfg->num_nvrtc_opts] = opt;
                          cfg->num_nvrtc_opts++;
-                         cfg->nvrtc_opts = realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*));
+                         cfg->nvrtc_opts = (const char**) realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*));
                          cfg->nvrtc_opts[cfg->num_nvrtc_opts] = NULL;
                        }|])
 
@@ -227,7 +227,7 @@
   GC.publicDef_ "context_new" GC.InitDecl $ \s ->
     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],
      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {
-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));
+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));
                           if (ctx == NULL) {
                             return NULL;
                           }
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
@@ -29,7 +29,7 @@
     Right (Program opencl_code opencl_prelude kernel_names types sizes prog') ->
       Right <$> GC.compileProg operations
                 (generateBoilerplate opencl_code opencl_prelude kernel_names types sizes)
-                include_opencl_h [Space "device", Space "local", DefaultSpace]
+                include_opencl_h [Space "device", DefaultSpace]
                 cliOptions prog'
   where operations :: GC.Operations OpenCL ()
         operations = GC.Operations
@@ -167,7 +167,7 @@
   GC.stm [C.cstm|{$item:decl
                   OPENCL_SUCCEED_OR_RETURN(
                     clEnqueueWriteBuffer(ctx->opencl.queue, $exp:mem, $exp:blocking,
-                                         $exp:i, sizeof($ty:t),
+                                         $exp:i * sizeof($ty:t), sizeof($ty:t),
                                          &$id:val',
                                          0, NULL, NULL));
                 }|]
@@ -180,7 +180,7 @@
   GC.decl [C.cdecl|$ty:t $id:val;|]
   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(
                    clEnqueueReadBuffer(ctx->opencl.queue, $exp:mem, CL_TRUE,
-                                       $exp:i, sizeof($ty:t),
+                                       $exp:i * sizeof($ty:t), sizeof($ty:t),
                                        &$id:val,
                                        0, NULL, NULL));
               |]
@@ -191,16 +191,12 @@
 allocateOpenCLBuffer :: GC.Allocate OpenCL ()
 allocateOpenCLBuffer mem size tag "device" =
   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_alloc(&ctx->opencl, $exp:size, $exp:tag, &$exp:mem));|]
-allocateOpenCLBuffer _ _ _ "local" =
-  return () -- Hack - these memory blocks do not actually exist.
 allocateOpenCLBuffer _ _ _ space =
-  fail $ "Cannot allocate in '" ++ space ++ "' space"
+  fail $ "Cannot allocate in '" ++ space ++ "' space."
 
 deallocateOpenCLBuffer :: GC.Deallocate OpenCL ()
 deallocateOpenCLBuffer mem tag "device" =
   GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_free(&ctx->opencl, $exp:mem, $exp:tag));|]
-deallocateOpenCLBuffer _ _ "local" =
-  return () -- Hack - these memory blocks do not actually exist.
 deallocateOpenCLBuffer _ _ space =
   fail $ "Cannot deallocate in '" ++ space ++ "' space"
 
@@ -252,7 +248,6 @@
 
 openclMemoryType :: GC.MemoryType OpenCL ()
 openclMemoryType "device" = pure [C.cty|typename cl_mem|]
-openclMemoryType "local" = pure [C.cty|unsigned char|] -- dummy type
 openclMemoryType space =
   fail $ "OpenCL backend does not support '" ++ space ++ "' memory space."
 
diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
--- a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
+++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs
@@ -69,13 +69,13 @@
   GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->
     ([C.cedecl|struct $id:cfg* $id:s(void);|],
      [C.cedecl|struct $id:cfg* $id:s(void) {
-                         struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));
+                         struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));
                          if (cfg == NULL) {
                            return NULL;
                          }
 
                          cfg->num_build_opts = 0;
-                         cfg->build_opts = malloc(sizeof(const char*));
+                         cfg->build_opts = (const char**) malloc(sizeof(const char*));
                          cfg->build_opts[0] = NULL;
                          $stms:size_value_inits
                          opencl_config_init(&cfg->opencl, $int:num_sizes,
@@ -96,7 +96,7 @@
      [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) {
                          cfg->build_opts[cfg->num_build_opts] = opt;
                          cfg->num_build_opts++;
-                         cfg->build_opts = realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*));
+                         cfg->build_opts = (const char**) realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*));
                          cfg->build_opts[cfg->num_build_opts] = NULL;
                        }|])
 
@@ -264,7 +264,7 @@
   GC.publicDef_ "context_new" GC.InitDecl $ \s ->
     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],
      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {
-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));
+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));
                           if (ctx == NULL) {
                             return NULL;
                           }
@@ -281,7 +281,7 @@
   GC.publicDef_ "context_new_with_command_queue" GC.InitDecl $ \s ->
     ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue);|],
      [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue) {
-                          struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));
+                          struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));
                           if (ctx == NULL) {
                             return NULL;
                           }
diff --git a/src/Futhark/CodeGen/Backends/CSOpenCL.hs b/src/Futhark/CodeGen/Backends/CSOpenCL.hs
--- a/src/Futhark/CodeGen/Backends/CSOpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/CSOpenCL.hs
@@ -129,6 +129,7 @@
                      Imp.SizeNumGroups -> "MaxNumGroups"
                      Imp.SizeTile -> "MaxTileSize"
                      Imp.SizeThreshold{} -> "MaxThreshold"
+                     Imp.SizeLocalMemory -> "MaxLocalMemory"
 
 callKernel (Imp.HostCode c) = CS.compileCode c
 
@@ -254,7 +255,8 @@
     , AssignTyped (PointerT VoidT) (Var ptr) (Just $ Addr $ Var scalar)
     , Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer"
         [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True
-        ,CS.toIntPtr i,CS.toIntPtr $ CS.sizeOf bt',CS.toIntPtr $ Var ptr
+        , CS.toIntPtr $ BinOp "*" i (CS.sizeOf bt')
+        , CS.toIntPtr $ CS.sizeOf bt',CS.toIntPtr $ Var ptr
     , Integer 0, Null, Null]
     ]
 
@@ -271,7 +273,8 @@
     [ CS.assignScalarPointer (Var val) (Var ptr)
     , Exp $ CS.simpleCall "CL10.EnqueueReadBuffer"
       [ Var "Ctx.OpenCL.Queue", memblockFromMem mem , Bool True
-      , CS.toIntPtr i, CS.toIntPtr $ CS.sizeOf bt', CS.toIntPtr $ Var ptr
+      , CS.toIntPtr $ BinOp "*" i (CS.sizeOf bt')
+      , CS.toIntPtr $ CS.sizeOf bt', CS.toIntPtr $ Var ptr
       , Integer 0, Null, Null]
     ]
   return $ Var val
@@ -388,32 +391,26 @@
   fail $ "Cannot return array from " ++ sid ++ " space."
 
 unpackArrayInput :: CS.EntryInput Imp.OpenCL ()
-unpackArrayInput mem memsize "device" t _ dims e = do
+unpackArrayInput mem "device" t _ dims e = do
   let size = foldr (BinOp "*") (Integer 1) dims'
   let t' = CS.compilePrimTypeToAST t
   let nbytes = BinOp "*" (CS.sizeOf t') size
   zipWithM_ (CS.unpackDim e) dims [0..]
   ptr <- pretty <$> newVName "ptr"
 
-  CS.stm $ compileMemsize memsize nbytes
-
-  let memsize' = CS.compileDim memsize
-
   CS.stm $ CS.getDefaultDecl (Imp.MemParam mem (Imp.Space "device"))
-  allocateOpenCLBuffer mem memsize' "device"
+  allocateOpenCLBuffer mem nbytes "device"
   CS.stm $ Unsafe [Fixed (Var ptr) (Addr $ Index (Field e "Item1") $ IdxExp $ Integer 0)
-      [ ifNotZeroSize memsize' $
+      [ ifNotZeroSize nbytes $
         Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer"
         [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True
-        , CS.toIntPtr (Integer 0), CS.toIntPtr memsize', CS.toIntPtr (Var ptr)
+        , CS.toIntPtr (Integer 0), CS.toIntPtr nbytes, CS.toIntPtr (Var ptr)
         , Integer 0, Null, Null]
       ]]
 
   where dims' = map CS.compileDim dims
-        compileMemsize (Imp.VarSize v) nbytes = Assign (Var $ CS.compileName v) nbytes
-        compileMemsize _ _                    = Pass
 
-unpackArrayInput _ _ sid _ _ _ _ =
+unpackArrayInput _ sid _ _ _ _ =
   fail $ "Cannot accept array from " ++ sid ++ " space."
 
 futharkSyncContext :: CSStmt
diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs
--- a/src/Futhark/CodeGen/Backends/GenericC.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC.hs
@@ -140,13 +140,13 @@
 -- | The type of a memory block in the given memory space.
 type MemoryType op s = SpaceId -> CompilerM op s C.Type
 
--- | Write a scalar to the given memory block with the given index and
--- in the given memory space.
+-- | Write a scalar to the given memory block with the given element
+-- index and in the given memory space.
 type WriteScalar op s =
   C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s ()
 
--- | Read a scalar from the given memory block with the given index and
--- in the given memory space.
+-- | Read a scalar from the given memory block with the given element
+-- index and in the given memory space.
 type ReadScalar op s =
   C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp
 
@@ -528,14 +528,10 @@
           (long long)size, desc, $string:spacedesc, ctx->$id:usagename);
   }
   int ret = $id:(fatMemUnRef space)(ctx, block, desc);
-  $items:alloc
-  block->references = (int*) malloc(sizeof(int));
-  *(block->references) = 1;
-  block->size = size;
-  block->desc = desc;
+
   ctx->$id:usagename += size;
   if (ctx->detail_memory) {
-    fprintf(stderr, "Allocated %lld bytes for %s in %s (now allocated: %lld bytes)",
+    fprintf(stderr, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)",
             (long long) size,
             desc, $string:spacedesc,
             (long long) ctx->$id:usagename);
@@ -548,6 +544,12 @@
   } else if (ctx->detail_memory) {
     fprintf(stderr, ".\n");
   }
+
+  $items:alloc
+  block->references = (int*) malloc(sizeof(int));
+  *(block->references) = 1;
+  block->size = size;
+  block->desc = desc;
   return ret;
   }|]
 
@@ -673,7 +675,7 @@
 opaqueName s vds = "opaque_" ++ hash (zipWith xor [0..] $ map ord (s ++ concatMap p vds))
   where p (ScalarValue pt signed _) =
           show (pt, signed)
-        p (ArrayValue _ _ space pt signed dims) =
+        p (ArrayValue _ space pt signed dims) =
           show (space, pt, signed, length dims)
 
         -- FIXME: a stupid hash algorithm; may have collisions.
@@ -704,7 +706,7 @@
   shape_array <- publicName $ "shape_" ++ name
 
   let shape_names = [ "dim"++show i | i <- [0..rank-1] ]
-      shape_params = [ [C.cparam|int $id:k|] | k <- shape_names ]
+      shape_params = [ [C.cparam|typename int64_t $id:k|] | k <- shape_names ]
       arr_size = cproduct [ [C.cexp|$id:k|] | k <- shape_names ]
       arr_size_array = cproduct [ [C.cexp|arr->shape[$int:i]|] | i <- [0..rank-1] ]
   copy <- asks envCopy
@@ -759,7 +761,7 @@
   return [C.cunit|
           $ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params) {
             $ty:array_type* bad = NULL;
-            $ty:array_type *arr = malloc(sizeof($ty:array_type));
+            $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type));
             if (arr == NULL) {
               return bad;
             }
@@ -770,7 +772,7 @@
           $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset,
                                             $params:shape_params) {
             $ty:array_type* bad = NULL;
-            $ty:array_type *arr = malloc(sizeof($ty:array_type));
+            $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type));
             if (arr == NULL) {
               return bad;
             }
@@ -808,7 +810,7 @@
 
       freeComponent _ ScalarValue{} =
         return ()
-      freeComponent i (ArrayValue _ _ _ pt signed shape) = do
+      freeComponent i (ArrayValue _ _ pt signed shape) = do
         let rank = length shape
         free_array <- publicName $ "free_" ++ arrayName pt signed rank
         stm [C.cstm|if ((tmp = $id:free_array(ctx, obj->$id:(tupleField i))) != 0) {
@@ -834,7 +836,7 @@
 valueDescToCType :: ValueDesc -> CompilerM op s C.Type
 valueDescToCType (ScalarValue pt signed _) =
   return $ signedPrimTypeToCType signed pt
-valueDescToCType (ArrayValue _ _ space pt signed shape) = do
+valueDescToCType (ArrayValue _ space pt signed shape) = do
   let pt' = signedPrimTypeToCType signed pt
       rank = length shape
   exists <- gets $ lookup (pt',rank) . compArrayStructs
@@ -898,15 +900,11 @@
           stm [C.cstm|$id:name = $exp:src;|]
           return pt'
 
-        prepareValue src vd@(ArrayValue mem mem_size _ _ _ shape) = do
+        prepareValue src vd@(ArrayValue mem _ _ _ shape) = do
           ty <- valueDescToCType vd
 
           stm [C.cstm|$exp:mem = $exp:src->mem;|]
-          case mem_size of
-            VarSize v -> stm [C.cstm|$id:v = $exp:src->mem.size;|]
-            ConstSize _ -> return ()
 
-
           let rank = length shape
               maybeCopyDim (VarSize d) i =
                 Just [C.cstm|$id:d = $exp:src->shape[$int:i];|]
@@ -924,7 +922,7 @@
 
           case vd of
             ArrayValue{} -> do
-              stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|]
+              stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]
               prepareValue [C.cexp|*$id:pname|] vd
               return [C.cparam|$ty:ty **$id:pname|]
             ScalarValue{} -> do
@@ -936,14 +934,14 @@
           ty <- opaqueToCType desc vds
           vd_ts <- mapM valueDescToCType vds
 
-          stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|]
+          stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|]
 
 
           forM_ (zip3 [0..] vd_ts vds) $ \(i,ct,vd) -> do
             let field = [C.cexp|(*$id:pname)->$id:(tupleField i)|]
             case vd of
               ScalarValue{} -> return ()
-              _ -> stm [C.cstm|assert(($exp:field = malloc(sizeof($ty:ct))) != NULL);|]
+              _ -> stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|]
             prepareValue field vd
 
           return [C.cparam|$ty:ty **$id:pname|]
@@ -951,7 +949,7 @@
         prepareValue dest (ScalarValue _ _ name) =
           stm [C.cstm|$exp:dest = $id:name;|]
 
-        prepareValue dest (ArrayValue mem _ _ _ _ shape) = do
+        prepareValue dest (ArrayValue mem _ _ _ shape) = do
           stm [C.cstm|$exp:dest->mem = $id:mem;|]
 
           let rank = length shape
@@ -1039,7 +1037,7 @@
   return [C.cstm|printf("#<opaque %s>", $string:desc);|]
 printStm (TransparentValue (ScalarValue bt ept _)) e =
   return $ printPrimStm [C.cexp|stdout|] e bt ept
-printStm (TransparentValue (ArrayValue _ _ _ bt ept shape)) e = do
+printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = do
   values_array <- publicName $ "values_" ++ name
   shape_array <- publicName $ "shape_" ++ name
   let num_elems = cproduct [ [C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0..rank-1] ]
@@ -1076,7 +1074,7 @@
   item [C.citem|$ty:(primTypeToCType t) $id:dest;|]
   stm $ readPrimStm dest i t ept
   return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|])
-readInput i (TransparentValue vd@(ArrayValue _ _ _ t ept dims)) = do
+readInput i (TransparentValue vd@(ArrayValue _ _ t ept dims)) = do
   dest <- newVName "read_value"
   shape <- newVName "read_shape"
   arr <- newVName "read_arr"
@@ -1124,7 +1122,7 @@
             TransparentValue ScalarValue{} -> do
               item [C.citem|$ty:ty $id:result;|]
               return ([C.cexp|$id:result|], [C.cstm|;|])
-            TransparentValue (ArrayValue _ _ _ t ept dims) -> do
+            TransparentValue (ArrayValue _ _ t ept dims) -> do
               let name = arrayName t ept $ length dims
               free_array <- publicName $ "free_" ++ name
               item [C.citem|$ty:ty *$id:result;|]
@@ -1307,6 +1305,7 @@
       option_parser = generateOptionParser "parse_options" $ benchmarkOptions++options
 
   let headerdefs = [C.cunit|
+$esc:("#pragma once\n")
 $esc:("/*\n * Headers\n*/\n")
 $esc:("#include <stdint.h>")
 $esc:("#include <stddef.h>")
@@ -1463,7 +1462,8 @@
   |]
 
   return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs)
-  where compileProg' = do
+    where
+      compileProg' = do
           (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces
 
           (prototypes, definitions) <- unzip <$> mapM compileFun funs
@@ -1477,29 +1477,30 @@
           ctx_ty <- contextType
           headerDecl MiscDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx);|]
           libDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx) {
-  if (ctx->detail_memory) {
-    $items:memreport
-  }
-  if (ctx->debugging) {
-    $items:debugreport
-  }
-}|]
+                      if (ctx->detail_memory) {
+                        $items:memreport
+                      }
+                      if (ctx->debugging) {
+                        $items:debugreport
+                      }
+                    }|]
 
           return (prototypes, definitions, entry_points)
-        funcToDef func = C.FuncDef func loc
-          where loc = case func of
-                        C.OldFunc _ _ _ _ _ _ l -> l
-                        C.Func _ _ _ _ _ l      -> l
 
-        builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps ++
-                  cFloat32Funs ++ cFloat64Funs
+      funcToDef func = C.FuncDef func loc
+        where loc = case func of
+                       C.OldFunc _ _ _ _ _ _ l -> l
+                       C.Func _ _ _ _ _ l      -> l
 
-        panic_h = $(embedStringFile "rts/c/panic.h")
-        values_h = $(embedStringFile "rts/c/values.h")
-        timing_h = $(embedStringFile "rts/c/timing.h")
-        lock_h = $(embedStringFile "rts/c/lock.h")
-        tuning_h = $(embedStringFile "rts/c/tuning.h")
+      builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps ++
+                cFloat32Funs ++ cFloat64Funs
 
+      panic_h  = $(embedStringFile "rts/c/panic.h")
+      values_h = $(embedStringFile "rts/c/values.h")
+      timing_h = $(embedStringFile "rts/c/timing.h")
+      lock_h   = $(embedStringFile "rts/c/lock.h")
+      tuning_h = $(embedStringFile "rts/c/tuning.h")
+
 compileFun :: (Name, Function op) -> CompilerM op s (C.Definition, C.Func)
 compileFun (fname, Function _ outputs inputs body _ _) = do
   (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs
@@ -1565,7 +1566,7 @@
 
 derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp
 derefPointer ptr i res_t =
-  [C.cexp|*(($ty:res_t)&($exp:ptr[$exp:i]))|]
+  [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|]
 
 writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s
 writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do
@@ -1707,10 +1708,15 @@
               { $items:items }
              |]
 
-compileCode (DebugPrint s _ e) = do
+compileCode (DebugPrint s (Just (_, e))) = do
   e' <- compileExp e
   stm [C.cstm|if (ctx->debugging) {
           fprintf(stderr, "%s: %d\n", $exp:s, (int)$exp:e');
+       }|]
+
+compileCode (DebugPrint s Nothing) =
+  stm [C.cstm|if (ctx->debugging) {
+          fprintf(stderr, "%s\n", $exp:s);
        }|]
 
 compileCode c
diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp.hs b/src/Futhark/CodeGen/Backends/GenericCSharp.hs
--- a/src/Futhark/CodeGen/Backends/GenericCSharp.hs
+++ b/src/Futhark/CodeGen/Backends/GenericCSharp.hs
@@ -123,7 +123,7 @@
                         CompilerM op s CSExp
 
 -- | Unpack the array being passed to an entry point.
-type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId ->
+type EntryInput op s = VName -> Imp.SpaceId ->
                        PrimType -> Imp.Signedness ->
                        [Imp.DimSize] ->
                        CSExp ->
@@ -646,12 +646,12 @@
   return $ cast $ Var $ compileName name
   where cast = compileTypecastExt bt ept
 
-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do
+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt ept dims)) = do
   let src = Var $ compileName mem
   let createTuple = "createTuple_" ++ compilePrimTypeExt bt ept
   return $ simpleCall createTuple [src, CreateArray (Primitive $ CSInt Int64T) $ Right $ map compileDim dims]
 
-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do
+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do
   unRefMem mem (Imp.Space sid)
   pack_output <- asks envEntryOutput
   pack_output mem sid bt ept dims
@@ -666,21 +666,16 @@
       cast = compileTypecast bt
   stm $ Assign vname' (cast e)
 
-entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace bt _ dims), e) = do
+entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt _ dims), e) = do
   zipWithM_ (unpackDim e) dims [0..]
   let arrayData = Field e "Item1"
   let dest = Var $ compileName mem
       unwrap_call = simpleCall "unwrapArray" [arrayData, sizeOf $ compilePrimTypeToAST bt]
-  case memsize of
-    Imp.VarSize sizevar ->
-      stm $ Assign (Var $ compileName sizevar) $ Field e "Item2.Length"
-    Imp.ConstSize _ ->
-      return ()
   stm $ Assign dest unwrap_call
 
-entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do
+entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do
   unpack_input <- asks envEntryInput
-  unpack <- collect $ unpack_input mem memsize sid bt ept dims e
+  unpack <- collect $ unpack_input mem sid bt ept dims e
   stms unpack
 
 extValueDescName :: Imp.ExternalValue -> String
@@ -706,7 +701,7 @@
 
 valueDescVName :: Imp.ValueDesc -> VName
 valueDescVName (Imp.ScalarValue _ _ vname) = vname
-valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname
+valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname
 
 consoleWrite :: String -> [CSExp] -> CSExp
 consoleWrite str exps = simpleCall "Console.Write" $ String str:exps
@@ -780,7 +775,7 @@
 -- TODO: If the type identifier of 'Float32' is changed, currently the error
 -- messages for reading binary input will not use this new name. This is also a
 -- problem for the C runtime system.
-readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ _ bt ept dims)) =
+readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) =
   let rank' = Var $ show $ length dims
       type_enum = String $ readTypeEnum bt ept
       bt' =  compilePrimTypeExt bt ept
@@ -802,11 +797,11 @@
 printStm :: Imp.ValueDesc -> CSExp -> CSExp -> CompilerM op s CSStmt
 printStm Imp.ScalarValue{} _ e =
   return $ printPrimStm e
-printStm (Imp.ArrayValue _ _ _ _ _ []) ind e = do
+printStm (Imp.ArrayValue _ _ _ _ []) ind e = do
   let e' = Index e (IdxExp (PostUnOp "++" ind))
   return $ printPrimStm e'
 
-printStm (Imp.ArrayValue mem memsize space bt ept (outer:shape)) ind e = do
+printStm (Imp.ArrayValue mem space bt ept (outer:shape)) ind e = do
   ptr <- newVName "shapePtr"
   first <- newVName "printFirst"
   let size = callMethod (CreateArray (Primitive $ CSInt Int32T) $ Right $ map compileDim $ outer:shape)
@@ -816,7 +811,7 @@
                              ]
       emptystr = "empty(" ++ ppArrayType bt (length shape) ++ ")"
 
-  printelem <- printStm (Imp.ArrayValue mem memsize space bt ept shape) ind e
+  printelem <- printStm (Imp.ArrayValue mem space bt ept shape) ind e
   return $
     If (BinOp "==" size (Integer 0))
       [puts emptystr]
@@ -926,9 +921,9 @@
         initCopy (varName, Imp.MemParam _ space) = declMem' varName space
         initCopy _ = Pass
 
-        valueDescFun (Imp.ArrayValue mem _ Imp.DefaultSpace _ _ _) =
+        valueDescFun (Imp.ArrayValue mem Imp.DefaultSpace _ _ _) =
             stm $ Assign (Var $ compileName mem ++ "_nbytes") (Var $ compileName mem ++ ".Length")
-        valueDescFun (Imp.ArrayValue mem _ (Imp.Space _) bt _ dims) =
+        valueDescFun (Imp.ArrayValue mem (Imp.Space _) bt _ dims) =
             stm $ Assign (Var $ compileName mem ++ "_nbytes") $ foldr (BinOp "*" . compileDim) (sizeOf $ compilePrimTypeToAST bt) dims
         valueDescFun _ = stm Pass
 
@@ -970,7 +965,7 @@
         getType' :: Imp.ValueDesc -> CSType
         getType' (Imp.ScalarValue primtype signedness _) =
           compilePrimTypeToASText primtype signedness
-        getType' (Imp.ArrayValue _ _ _ primtype signedness _) =
+        getType' (Imp.ArrayValue _ _ primtype signedness _) =
           let t = compilePrimTypeToASText primtype signedness
           in Composite $ SystemTupleT [Composite $ ArrayT t, Composite $ ArrayT $ Primitive $ CSInt Int64T]
 
diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs
--- a/src/Futhark/CodeGen/Backends/GenericPython.hs
+++ b/src/Futhark/CodeGen/Backends/GenericPython.hs
@@ -100,7 +100,7 @@
                         CompilerM op s PyExp
 
 -- | Unpack the array being passed to an entry point.
-type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId ->
+type EntryInput op s = VName -> Imp.SpaceId ->
                        PrimType -> Imp.Signedness ->
                        [Imp.DimSize] ->
                        PyExp ->
@@ -272,11 +272,10 @@
          , optionAction =
            [ Assign (Var "entry_point") $ Var "optarg" ]
          },
-  -- The -b option is just a dummy for now.
   Option { optionLongName = "binary-output"
          , optionShortName = Just 'b'
          , optionArgument = NoArgument
-         , optionAction = [Pass]
+         , optionAction = [Assign (Var "binary_output") $ Bool True]
          },
   Option { optionLongName = "tuning"
          , optionShortName = Nothing
@@ -353,6 +352,7 @@
           Assign (Var "do_warmup_run") (Bool False) :
           Assign (Var "num_runs") (Integer 1) :
           Assign (Var "entry_point") (String "main") :
+          Assign (Var "binary_output") (Bool False) :
           generateOptionParser (standardOptions ++ options)
 
         selectEntryPoint entry_point_names entry_points =
@@ -415,10 +415,10 @@
 entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) =
   return $ simpleCall tf [Var $ compileName name]
   where tf = compilePrimToExtNp bt ept
-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do
+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace bt ept dims)) = do
   let cast = Cast (Var $ compileName mem) (compilePrimTypeExt bt ept)
   return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims]
-entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do
+entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do
   pack_output <- asks envEntryOutput
   pack_output mem sid bt ept dims
 
@@ -455,7 +455,7 @@
     [Catch (Tuple [Var "TypeError", Var "AssertionError"])
      [badInput i e $ prettySigned (s==Imp.TypeUnsigned) bt]]
 
-entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace t s dims), e) = do
+entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem Imp.DefaultSpace t s dims), e) = do
   let type_is_wrong =
         UnOp "not" $
         BinOp "and"
@@ -470,18 +470,11 @@
   let dest = Var $ compileName mem
       unwrap_call = simpleCall "unwrapArray" [e]
 
-  case memsize of
-    Imp.VarSize sizevar ->
-      stm $ Assign (Var $ compileName sizevar) $
-      simpleCall "np.int32" [Field e "nbytes"]
-    Imp.ConstSize _ ->
-      return ()
-
   stm $ Assign dest unwrap_call
 
-entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do
+entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do
   unpack_input <- asks envEntryInput
-  unpack <- collect $ unpack_input mem memsize sid bt ept dims e
+  unpack <- collect $ unpack_input mem sid bt ept dims e
   stm $ Try unpack
     [Catch (Tuple [Var "TypeError", Var "AssertionError"])
      [badInput i e $ concat (replicate (length dims) "[]") ++
@@ -501,7 +494,7 @@
 
 valueDescVName :: Imp.ValueDesc -> VName
 valueDescVName (Imp.ScalarValue _ _ vname) = vname
-valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname
+valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname
 
 -- Key into the FUTHARK_PRIMTYPES dict.
 readTypeEnum :: PrimType -> Imp.Signedness -> String
@@ -527,7 +520,7 @@
   let type_name = readTypeEnum bt ept
   in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" [String type_name]
 
-readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ _ bt ept dims)) =
+readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) =
   let type_name = readTypeEnum bt ept
   in Assign (Var $ extValueDescName decl) $ simpleCall "read_value"
      [String $ concat (replicate (length dims) "[]") ++ type_name]
@@ -542,11 +535,13 @@
   where printValue' (Imp.OpaqueValue desc _) _ =
           return [Exp $ simpleCall "sys.stdout.write"
                   [String $ "#<opaque " ++ desc ++ ">"]]
-        printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize (Space _) bt ept shape)) e =
-          printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize DefaultSpace bt ept shape)) $
+        printValue' (Imp.TransparentValue (Imp.ArrayValue mem (Space _) bt ept shape)) e =
+          printValue' (Imp.TransparentValue (Imp.ArrayValue mem DefaultSpace bt ept shape)) $
           simpleCall (pretty e ++ ".get") []
         printValue' (Imp.TransparentValue _) e =
-          return [Exp $ simpleCall "write_value" [e],
+          return [Exp $ Call (Var "write_value")
+                   [Arg e,
+                    ArgKeyword "binary" (Var "binary_output")],
                   Exp $ simpleCall "sys.stdout.write" [String "\n"]]
 
 prepareEntry :: (Name, Imp.Function op) -> CompilerM op s
@@ -616,7 +611,7 @@
                    map desc $ Imp.functionResult func)
   where desc (Imp.OpaqueValue d _) = d
         desc (Imp.TransparentValue (Imp.ScalarValue pt s _)) = readTypeEnum pt s
-        desc (Imp.TransparentValue (Imp.ArrayValue _ _ _ pt s dims)) =
+        desc (Imp.TransparentValue (Imp.ArrayValue _ _ pt s dims)) =
           concat (replicate (length dims) "[]") ++ readTypeEnum pt s
 
 callEntryFun :: [PyStmt] -> (Name, Imp.Function op)
@@ -751,7 +746,7 @@
     (IntType Int64, _) -> "np.int64"
     (FloatType Float32, _) -> "np.float32"
     (FloatType Float64, _) -> "np.float64"
-    (Imp.Bool, _) -> "np.bool"
+    (Imp.Bool, _) -> "np.bool_"
     (Cert, _) -> "np.byte"
 
 compilePrimValue :: Imp.PrimValue -> PyExp
diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
--- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs
+++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs
@@ -179,7 +179,7 @@
               [Arg val, ArgKeyword "dtype" $ Var $ Py.compilePrimType bt]
   Py.stm $ Exp $ Call (Var "cl.enqueue_copy")
     [Arg $ Var "self.queue", Arg mem', Arg nparr,
-     ArgKeyword "device_offset" $ asLong i,
+     ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),
      ArgKeyword "is_blocking" $ Var "synchronous"]
 
 writeOpenCLScalar _ _ _ space _ =
@@ -196,7 +196,7 @@
   Py.stm $ Assign val' nparr
   Py.stm $ Exp $ Call (Var "cl.enqueue_copy")
     [Arg $ Var "self.queue", Arg val', Arg mem',
-     ArgKeyword "device_offset" $ asLong i,
+     ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt),
      ArgKeyword "is_blocking" $ Bool True]
   return $ Index val' $ IdxExp $ Integer 0
 
@@ -303,7 +303,7 @@
   fail $ "Cannot return array from " ++ sid ++ " space."
 
 unpackArrayInput :: Py.EntryInput Imp.OpenCL ()
-unpackArrayInput mem memsize "device" t s dims e = do
+unpackArrayInput mem "device" t s dims e = do
   let type_is_ok =
         BinOp "and"
         (BinOp "in" (Py.simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"]))
@@ -312,14 +312,7 @@
 
   zipWithM_ (Py.unpackDim e) dims [0..]
 
-  case memsize of
-    Imp.VarSize sizevar ->
-      Py.stm $ Assign (Var $ Py.compileName sizevar) $
-      Py.simpleCall "np.int64" [Field e "nbytes"]
-    Imp.ConstSize _ ->
-      return ()
-
-  let memsize' = Py.compileDim memsize
+  let memsize' = Py.simpleCall "np.int64" [Field e "nbytes"]
       pyOpenCLArrayCase =
         [Assign mem_dest $ Field e "data"]
   numpyArrayCase <- Py.collect $ do
@@ -335,7 +328,7 @@
     pyOpenCLArrayCase
     numpyArrayCase
   where mem_dest = Var $ Py.compileName mem
-unpackArrayInput _ _ sid _ _ _ _ =
+unpackArrayInput _ sid _ _ _ _ =
   fail $ "Cannot accept array from " ++ sid ++ " space."
 
 ifNotZeroSize :: PyExp -> PyStmt -> PyStmt
diff --git a/src/Futhark/CodeGen/Backends/SequentialC.hs b/src/Futhark/CodeGen/Backends/SequentialC.hs
--- a/src/Futhark/CodeGen/Backends/SequentialC.hs
+++ b/src/Futhark/CodeGen/Backends/SequentialC.hs
@@ -38,7 +38,7 @@
           GC.publicDef_ "context_config_new" GC.InitDecl $ \s ->
             ([C.cedecl|struct $id:cfg* $id:s();|],
              [C.cedecl|struct $id:cfg* $id:s() {
-                                 struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg));
+                                 struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg));
                                  if (cfg == NULL) {
                                    return NULL;
                                  }
@@ -80,7 +80,7 @@
           GC.publicDef_ "context_new" GC.InitDecl $ \s ->
             ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|],
              [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) {
-                                  struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx));
+                                  struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx));
                                   if (ctx == NULL) {
                                     return NULL;
                                   }
diff --git a/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs b/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs
--- a/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs
+++ b/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs
@@ -24,7 +24,7 @@
 import qualified Language.C.Quote.C as C
 
 import Futhark.CodeGen.ImpCode
-import Futhark.Util.Pretty (pretty)
+import Futhark.Util.Pretty (pretty, prettyOneLine)
 import Futhark.Util (zEncodeString)
 
 -- | The C type corresponding to a signed integer type.
@@ -158,6 +158,12 @@
         mkSlt = intCmpOp "slt" [C.cexp|x < y|]
         mkSle = intCmpOp "sle" [C.cexp|x <= y|]
 
+        -- We define some operations as macros rather than functions,
+        -- because this allows us to use them as constant expressions
+        -- in things like array sizes and static initialisers.
+        macro name rhs =
+          [C.cedecl|$esc:("#define " ++ name ++ "(x) (" ++ prettyOneLine rhs ++ ")")|]
+
         mkPow t =
           let ct = intTypeToCType t
           in [C.cedecl|static inline $ty:ct $id:(taggedI "pow" t)($ty:ct x, $ty:ct y) {
@@ -172,16 +178,12 @@
                          return res;
               }|]
 
-        mkSExt from_t to_t =
-          [C.cedecl|static inline $ty:to_ct
-                    $id:name($ty:from_ct x) { return x;} |]
+        mkSExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|]
           where name = "sext_"++pretty from_t++"_"++pretty to_t
                 from_ct = intTypeToCType from_t
                 to_ct = intTypeToCType to_t
 
-        mkZExt from_t to_t =
-          [C.cedecl|static inline $ty:to_ct
-                    $id:name($ty:from_ct x) { return x;} |]
+        mkZExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|]
           where name = "zext_"++pretty from_t++"_"++pretty to_t
                 from_ct = uintTypeToCType from_t
                 to_ct = uintTypeToCType to_t
@@ -315,6 +317,14 @@
       return atan2(x,y);
     }
 
+    static inline float $id:(funName' "gamma32")(float x) {
+      return tgamma(x);
+    }
+
+    static inline float $id:(funName' "lgamma32")(float x) {
+      return lgamma(x);
+    }
+
     static inline float $id:(funName' "round32")(float x) {
       return rint(x);
     }
@@ -394,6 +404,14 @@
 
     static inline double $id:(funName' "atan2_64")(double x, double y) {
       return atan2(x,y);
+    }
+
+    static inline double $id:(funName' "gamma64")(double x) {
+      return tgamma(x);
+    }
+
+    static inline double $id:(funName' "lgamma64")(double x) {
+      return lgamma(x);
     }
 
     static inline double $id:(funName' "round64")(double x) {
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
@@ -31,6 +31,7 @@
   , Volatility (..)
   , Arg (..)
   , var
+  , vi32
   , index
   , ErrorMsg(..)
   , ErrorMsgPart(..)
@@ -105,7 +106,7 @@
                 deriving (Eq, Show)
 
 -- | A description of an externally meaningful value.
-data ValueDesc = ArrayValue VName MemSize Space PrimType Signedness [DimSize]
+data ValueDesc = ArrayValue VName Space PrimType Signedness [DimSize]
                -- ^ An array with memory block, memory block size,
                -- memory space, element type, signedness of element
                -- type (if applicable), and shape.
@@ -174,7 +175,7 @@
               -- ^ Destination, offset in destination, destination
               -- space, source, offset in source, offset space, number
               -- of bytes.
-            | Write VName (Count Bytes) PrimType Space Volatility Exp
+            | Write VName (Count Elements) PrimType Space Volatility Exp
             | SetScalar VName Exp
             | SetMem VName VName Space
               -- ^ Must be in same space.
@@ -185,12 +186,13 @@
               -- ^ Has the same semantics as the contained code, but
               -- the comment should show up in generated code for ease
               -- of inspection.
-            | DebugPrint String PrimType Exp
+            | DebugPrint String (Maybe (PrimType, Exp))
               -- ^ Print the given value (of the given type) to the
               -- screen, somehow annotated with the given string as a
-              -- description.  This has no semantic meaning, but is
-              -- used entirely for debugging.  Code generators are
-              -- free to ignore this statement.
+              -- description.  If no type/value pair, just print the
+              -- string.  This has no semantic meaning, but is used
+              -- entirely for debugging.  Code generators are free to
+              -- ignore this statement.
             | Op a
             deriving (Show)
 
@@ -208,7 +210,7 @@
 
 data ExpLeaf = ScalarVar VName
              | SizeOf PrimType
-             | Index VName (Count Bytes) PrimType Space Volatility
+             | Index VName (Count Elements) PrimType Space Volatility
            deriving (Eq, Show)
 
 type Exp = PrimExp ExpLeaf
@@ -253,7 +255,11 @@
 var :: VName -> PrimType -> Exp
 var = LeafExp . ScalarVar
 
-index :: VName -> Count Bytes -> PrimType -> Space -> Volatility -> Exp
+-- | Turn a 'VName' into a 'int32' 'Imp.ScalarVar'.
+vi32 :: VName -> Exp
+vi32 = flip var $ IntType Int32
+
+index :: VName -> Count Elements -> PrimType -> Space -> Volatility -> Exp
 index arr i t s vol = LeafExp (Index arr i t s vol) t
 
 -- Prettyprinting definitions.
@@ -286,8 +292,8 @@
     ppr t <+> ppr name <> ept'
     where ept' = case ept of TypeUnsigned -> text " (unsigned)"
                              TypeDirect   -> mempty
-  ppr (ArrayValue mem memsize space et ept shape) =
-    foldr f (ppr et) shape <+> text "at" <+> ppr mem <> parens (ppr memsize) <> space' <+> ept'
+  ppr (ArrayValue mem space et ept shape) =
+    foldr f (ppr et) shape <+> text "at" <+> ppr mem <> space' <+> ept'
     where f e s = brackets $ s <> comma <> ppr e
           ept' = case ept of TypeUnsigned -> text " (unsigned)"
                              TypeDirect   -> mempty
@@ -361,8 +367,10 @@
     ppr fname <> parens (commasep $ map ppr args)
   ppr (Comment s code) =
     text "--" <+> text s </> ppr code
-  ppr (DebugPrint desc pt e) =
+  ppr (DebugPrint desc (Just (pt, e))) =
     text "debug" <+> parens (commasep [text (show desc), ppr pt, ppr e])
+  ppr (DebugPrint desc Nothing) =
+    text "debug" <+> parens (text (show desc))
 
 instance Pretty Arg where
   ppr (MemArg m) = ppr m
@@ -445,8 +453,8 @@
     pure $ Call dests fname args
   traverse f (Comment s code) =
     Comment s <$> traverse f code
-  traverse _ (DebugPrint s t e) =
-    pure $ DebugPrint s t e
+  traverse _ (DebugPrint s v) =
+    pure $ DebugPrint s v
 
 declaredIn :: Code a -> Names
 declaredIn (DeclareMem name _) = S.singleton name
@@ -496,8 +504,8 @@
     freeIn op
   freeIn (Comment _ code) =
     freeIn code
-  freeIn (DebugPrint _ _ e) =
-    freeIn e
+  freeIn (DebugPrint _ v) =
+    maybe mempty (freeIn . snd) v
 
 instance FreeIn ExpLeaf where
   freeIn (Index v e _ _ _) = freeIn v <> freeIn e
diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs
--- a/src/Futhark/CodeGen/ImpCode/Kernels.hs
+++ b/src/Futhark/CodeGen/ImpCode/Kernels.hs
@@ -56,8 +56,6 @@
 -- | A generic kernel containing arbitrary kernel code.
 data Kernel = Kernel
               { kernelBody :: Imp.Code KernelOp
-              , kernelLocalMemory :: [LocalMemoryUse]
-              -- ^ The local memory used by this kernel.
 
               , kernelUses :: [KernelUse]
                 -- ^ The host variables referenced by the kernel.
@@ -71,7 +69,7 @@
             deriving (Show)
 
 -- ^ In-kernel name and per-workgroup size in bytes.
-type LocalMemoryUse = (VName, Either MemSize KernelConstExp)
+type LocalMemoryUse = (VName, Either (Count Bytes) KernelConstExp)
 
 data KernelUse = ScalarUse VName PrimType
                | MemoryUse VName
@@ -87,7 +85,7 @@
         sameKernel _ _ = False
 
 -- | Get an atomic operator corresponding to a binary operator.
-atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Bytes -> Exp -> AtomicOp)
+atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Elements -> Exp -> AtomicOp)
 atomicBinOp = flip lookup [ (Add Int32, AtomicAdd)
                           , (SMax Int32, AtomicSMax)
                           , (SMin Int32, AtomicSMin)
@@ -140,15 +138,8 @@
     text "kernel" <+> brace
     (text "groups" <+> brace (ppr $ kernelNumGroups kernel) </>
      text "group_size" <+> brace (ppr $ kernelGroupSize kernel) </>
-     text "local_memory" <+> brace (commasep $
-                                    map ppLocalMemory $
-                                    kernelLocalMemory kernel) </>
      text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) </>
      text "body" <+> brace (ppr $ kernelBody kernel))
-    where ppLocalMemory (name, Left size) =
-            ppr name <+> parens (ppr size <+> text "bytes")
-          ppLocalMemory (name, Right size) =
-            ppr name <+> parens (ppr size <+> text "bytes (const)")
 
 data KernelOp = GetGroupId VName Int
               | GetLocalId VName Int
@@ -156,24 +147,27 @@
               | GetGlobalSize VName Int
               | GetGlobalId VName Int
               | GetLockstepWidth VName
-              | Atomic AtomicOp
+              | Atomic Space AtomicOp
               | LocalBarrier
               | GlobalBarrier
-              | MemFence
+              | MemFenceLocal
+              | MemFenceGlobal
+              | PrivateAlloc VName (Count Bytes)
+              | LocalAlloc VName (Either (Count Bytes) KernelConstExp)
               deriving (Show)
 
 -- Atomic operations return the value stored before the update.
 -- This value is stored in the first VName.
-data AtomicOp = AtomicAdd VName VName (Count Bytes) Exp
-              | AtomicSMax VName VName (Count Bytes) Exp
-              | AtomicSMin VName VName (Count Bytes) Exp
-              | AtomicUMax VName VName (Count Bytes) Exp
-              | AtomicUMin VName VName (Count Bytes) Exp
-              | AtomicAnd VName VName (Count Bytes) Exp
-              | AtomicOr VName VName (Count Bytes) Exp
-              | AtomicXor VName VName (Count Bytes) Exp
-              | AtomicCmpXchg VName VName (Count Bytes) Exp Exp
-              | AtomicXchg VName VName (Count Bytes) Exp
+data AtomicOp = AtomicAdd VName VName (Count Elements) Exp
+              | AtomicSMax VName VName (Count Elements) Exp
+              | AtomicSMin VName VName (Count Elements) Exp
+              | AtomicUMax VName VName (Count Elements) Exp
+              | AtomicUMin VName VName (Count Elements) Exp
+              | AtomicAnd VName VName (Count Elements) Exp
+              | AtomicOr VName VName (Count Elements) Exp
+              | AtomicXor VName VName (Count Elements) Exp
+              | AtomicCmpXchg VName VName (Count Elements) Exp Exp
+              | AtomicXchg VName VName (Count Elements) Exp
               deriving (Show)
 
 instance FreeIn AtomicOp where
@@ -211,41 +205,49 @@
     text "local_barrier()"
   ppr GlobalBarrier =
     text "global_barrier()"
-  ppr MemFence =
-    text "mem_fence()"
-  ppr (Atomic (AtomicAdd old arr ind x)) =
+  ppr MemFenceLocal =
+    text "mem_fence_local()"
+  ppr MemFenceGlobal =
+    text "mem_fence_global()"
+  ppr (PrivateAlloc name size) =
+    ppr name <+> equals <+> text "private_alloc" <> parens (ppr size)
+  ppr (LocalAlloc name size) =
+    ppr name <+> equals <+> text "local_alloc" <>
+    parens (either ppr constCase size)
+    where constCase e = text "(constant)" <+> ppr e
+  ppr (Atomic _ (AtomicAdd old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_add" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicSMax old arr ind x)) =
+  ppr (Atomic _ (AtomicSMax old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_smax" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicSMin old arr ind x)) =
+  ppr (Atomic _ (AtomicSMin old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_smin" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicUMax old arr ind x)) =
+  ppr (Atomic _ (AtomicUMax old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_umax" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicUMin old arr ind x)) =
+  ppr (Atomic _ (AtomicUMin old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_umin" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicAnd old arr ind x)) =
+  ppr (Atomic _ (AtomicAnd old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_and" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicOr old arr ind x)) =
+  ppr (Atomic _ (AtomicOr old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_or" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicXor old arr ind x)) =
+  ppr (Atomic _ (AtomicXor old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_xor" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic (AtomicCmpXchg old arr ind x y)) =
+  ppr (Atomic _ (AtomicCmpXchg old arr ind x y)) =
     ppr old <+> text "<-" <+> text "atomic_cmp_xchg" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y])
-  ppr (Atomic (AtomicXchg old arr ind x)) =
+  ppr (Atomic _ (AtomicXchg old arr ind x)) =
     ppr old <+> text "<-" <+> text "atomic_xchg" <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
 
 instance FreeIn KernelOp where
-  freeIn (Atomic op) = freeIn op
+  freeIn (Atomic _ op) = freeIn op
   freeIn _ = mempty
 
 brace :: Doc -> Doc
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
@@ -10,6 +10,7 @@
   , ExpCompiler
   , CopyCompiler
   , StmsCompiler
+  , AllocCompiler
   , Operations (..)
   , defaultOperations
   , ValueDestination
@@ -30,6 +31,7 @@
   , emitFunction
   , hasFunction
   , collect
+  , collect'
   , comment
   , VarEntry (..)
   , ArrayEntry (..)
@@ -40,10 +42,7 @@
   , lookupMemory
 
     -- * Building Blocks
-  , compileSubExp
-  , compileSubExpOfType
-  , compileSubExpTo
-  , compilePrimExp
+  , ToExp(..)
   , compileAlloc
   , subExpToDimSize
   , everythingVolatile
@@ -58,13 +57,13 @@
   , strideArray
   , fullyIndexArray
   , fullyIndexArray'
-  , varIndex
   , Imp.dimSizeToExp
   , dimSizeToSubExp
   , copy
   , copyDWIM
   , copyDWIMDest
   , copyElementWise
+  , typeSize
 
   -- * Constructing code.
   , dLParams
@@ -72,15 +71,16 @@
   , dScope
   , dScopes
   , dArray
-  , dPrim, dPrim_, dPrimV
+  , dPrim, dPrim_, dPrimV_, dPrimV
 
   , sFor, sWhile
   , sComment
   , sIf, sWhen, sUnless
   , sOp
   , sDeclareMem, sAlloc, sAlloc_
-  , sArray, sAllocArray, sStaticArray
+  , sArray, sAllocArray, sAllocArrayPerm, sStaticArray
   , sWrite, sUpdate
+  , sLoopNest
   , (<--)
   )
   where
@@ -101,7 +101,7 @@
 import Futhark.CodeGen.ImpCode
   (Count (..),
    Bytes, Elements,
-   bytes, withElemType)
+   bytes, elements, withElemType)
 import Futhark.Representation.ExplicitMemory
 import Futhark.Representation.SOACS (SOACS)
 import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
@@ -114,7 +114,7 @@
 type OpCompiler lore op = Pattern lore -> Op lore -> ImpM lore op ()
 
 -- | How to compile some 'Stms'.
-type StmsCompiler lore op = Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()
+type StmsCompiler lore op = Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()
 
 -- | How to compile an 'Exp'.
 type ExpCompiler lore op = Pattern lore -> Exp lore -> ImpM lore op ()
@@ -125,10 +125,14 @@
                            -> Count Elements -- ^ Number of row elements of the source.
                            -> ImpM lore op ()
 
+-- | An alternate way of compiling an allocation.
+type AllocCompiler lore op = VName -> Count Bytes -> ImpM lore op ()
+
 data Operations lore op = Operations { opsExpCompiler :: ExpCompiler lore op
                                      , opsOpCompiler :: OpCompiler lore op
                                      , opsStmsCompiler :: StmsCompiler lore op
                                      , opsCopyCompiler :: CopyCompiler lore op
+                                     , opsAllocCompilers :: M.Map Space (AllocCompiler lore op)
                                      }
 
 -- | An operations set for which the expression compiler always
@@ -139,6 +143,7 @@
                                    , opsOpCompiler = opc
                                    , opsStmsCompiler = defCompileStms
                                    , opsCopyCompiler = defaultCopy
+                                   , opsAllocCompilers = mempty
                                    }
 
 -- | When an array is dared, this is where it is stored.
@@ -157,10 +162,7 @@
 entryArrayShape :: ArrayEntry -> [Imp.DimSize]
 entryArrayShape = memLocationShape . entryArrayLocation
 
-data MemEntry = MemEntry {
-      entryMemSize  :: Imp.MemSize
-    , entryMemSpace :: Imp.Space
-  }
+newtype MemEntry = MemEntry { entryMemSpace :: Imp.Space }
   deriving (Show)
 
 newtype ScalarEntry = ScalarEntry {
@@ -200,23 +202,22 @@
   , envStmsCompiler :: StmsCompiler lore op
   , envOpCompiler :: OpCompiler lore op
   , envCopyCompiler :: CopyCompiler lore op
+  , envAllocCompilers :: M.Map Space (AllocCompiler lore op)
   , envDefaultSpace :: Imp.Space
   , envVolatility :: Imp.Volatility
-  , envFakeMemory :: [Space]
-    -- ^ Do not actually generate allocations for these memory spaces.
   , envFunction :: Name
     -- ^ Name of the function we are compiling.
   }
 
-newEnv :: Operations lore op -> Imp.Space -> [Imp.Space] -> Name -> Env lore op
-newEnv ops ds fake fname =
+newEnv :: Operations lore op -> Imp.Space -> Name -> Env lore op
+newEnv ops ds fname =
   Env { envExpCompiler = opsExpCompiler ops
       , envStmsCompiler = opsStmsCompiler ops
       , envOpCompiler = opsOpCompiler ops
       , envCopyCompiler = opsCopyCompiler ops
+      , envAllocCompilers = mempty
       , envDefaultSpace = ds
       , envVolatility = Imp.Nonvolatile
-      , envFakeMemory = fake
       , envFunction = fname
       }
 
@@ -245,10 +246,12 @@
   getNameSource = gets stateNameSource
   putNameSource src = modify $ \s -> s { stateNameSource = src }
 
+-- Cannot be an ExplicitMemory scope because the index functions have
+-- the wrong leaves (VName instead of Imp.Exp).
 instance HasScope SOACS (ImpM lore op) where
   askScope = M.map (LetInfo . entryType) <$> gets stateVTable
     where entryType (MemVar _ memEntry) =
-            Mem (dimSizeToSubExp $ entryMemSize memEntry) (entryMemSpace memEntry)
+            Mem (entryMemSpace memEntry)
           entryType (ArrayVar _ arrayEntry) =
             Array
             (entryArrayElemType arrayEntry)
@@ -258,10 +261,9 @@
             Prim $ entryScalarType scalarEntry
 
 runImpM :: ImpM lore op a
-        -> Operations lore op -> Imp.Space -> [Imp.Space] -> Name -> State lore op
+        -> Operations lore op -> Imp.Space -> Name -> State lore op
         -> Either InternalError (a, State lore op, Imp.Code op)
-runImpM (ImpM m) comp space fake fname =
-  runRWST m (newEnv comp space fake fname)
+runImpM (ImpM m) ops space fname = runRWST m $ newEnv ops space fname
 
 subImpM_ :: Operations lore' op' -> ImpM lore' op' a
          -> ImpM lore op (Imp.Code op')
@@ -276,6 +278,7 @@
                      , envStmsCompiler = opsStmsCompiler ops
                      , envCopyCompiler = opsCopyCompiler ops
                      , envOpCompiler = opsOpCompiler ops
+                     , envAllocCompilers = opsAllocCompilers ops
                      }
                  s { stateVTable = M.map scrubExps $ stateVTable s
                    , stateFunctions = mempty } of
@@ -321,16 +324,16 @@
                                  in isJust $ lookup fname fs
 
 compileProg :: (ExplicitMemorish lore, MonadFreshNames m) =>
-               Operations lore op -> Imp.Space -> [Imp.Space]
+               Operations lore op -> Imp.Space
             -> Prog lore -> m (Either InternalError (Imp.Functions op))
-compileProg ops space fake prog =
+compileProg ops space prog =
   modifyNameSource $ \src ->
   case foldM compileFunDef' (newState src) (progFunctions prog) of
     Left err -> (Left err, src)
     Right s -> (Right $ stateFunctions s, stateNameSource s)
   where compileFunDef' s fdef = do
           ((), s', _) <-
-            runImpM (compileFunDef fdef) ops space fake (funDefName fdef) s
+            runImpM (compileFunDef fdef) ops space (funDefName fdef) s
           return s'
 
 compileInParam :: ExplicitMemorish lore =>
@@ -338,20 +341,18 @@
 compileInParam fparam = case paramAttr fparam of
   MemPrim bt ->
     return $ Left $ Imp.ScalarParam name bt
-  MemMem _ space ->
+  MemMem space ->
     return $ Left $ Imp.MemParam name space
   MemArray bt shape _ (ArrayIn mem ixfun) -> do
     shape' <- mapM subExpToDimSize $ shapeDims shape
     return $ Right $ ArrayDecl name bt $
-      MemLocation mem shape' $ fmap compilePrimExp ixfun
+      MemLocation mem shape' $ fmap (toExp' int32) ixfun
   where name = paramName fparam
 
 data ArrayDecl = ArrayDecl VName PrimType MemLocation
 
 fparamSizes :: Typed attr => Param attr -> S.Set VName
-fparamSizes fparam
-  | Mem (Var size) _ <- paramType fparam = S.singleton size
-  | otherwise = S.fromList $ subExpVars $ arrayDims $ paramType fparam
+fparamSizes = S.fromList . subExpVars . arrayDims . paramType
 
 compileInParams :: ExplicitMemorish lore =>
                    [FParam lore] -> [EntryPointType]
@@ -365,21 +366,19 @@
 
       summaries = M.fromList $ mapMaybe memSummary params
         where memSummary param
-                | MemMem (Constant (IntValue (Int64Value size))) space <- paramAttr param =
-                    Just (paramName param, (Imp.ConstSize size, space))
-                | MemMem (Var size) space <- paramAttr param =
-                    Just (paramName param, (Imp.VarSize size, space))
+                | MemMem space <- paramAttr param =
+                    Just (paramName param, space)
                 | otherwise =
                     Nothing
 
-      findMemInfo :: VName -> Maybe (Imp.MemSize, Space)
+      findMemInfo :: VName -> Maybe Space
       findMemInfo = flip M.lookup summaries
 
       mkValueDesc fparam signedness =
         case (findArray $ paramName fparam, paramType fparam) of
           (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do
-            (memsize, memspace) <- findMemInfo mem
-            Just $ Imp.ArrayValue mem memsize memspace bt signedness shape
+            memspace <- findMemInfo mem
+            Just $ Imp.ArrayValue mem memspace bt signedness shape
           (_, Prim bt)
             | paramName fparam `S.member` sizes ->
               Nothing
@@ -440,18 +439,16 @@
           return (Imp.ScalarValue t ept out, ScalarDestination out)
         mkParam (MemArray t shape _ attr) ept = do
           space <- asks envDefaultSpace
-          (memout, memsize) <- case attr of
-            ReturnsNewBlock _ x x_size _ixfun -> do
+          memout <- case attr of
+            ReturnsNewBlock _ x _ixfun -> do
               memout <- imp $ newVName "out_mem"
-              sizeout <- ensureMemSizeOut x_size
               tell ([Imp.MemParam memout space],
                     M.singleton x $ MemoryDestination memout)
-              return (memout, sizeout)
-            ReturnsInBlock memout _ -> do
-              memsize <- imp $ entryMemSize <$> lookupMemory memout
-              return (memout, memsize)
+              return memout
+            ReturnsInBlock memout _ ->
+              return memout
           resultshape <- mapM inspectExtSize $ shapeDims shape
-          return (Imp.ArrayValue memout memsize space t ept resultshape,
+          return (Imp.ArrayValue memout space t ept resultshape,
                   ArrayDestination Nothing)
 
         inspectExtSize (Ext x) = do
@@ -468,19 +465,6 @@
         inspectExtSize (Free se) =
           imp $ subExpToDimSize se
 
-        -- | Return the name of the out-parameter for the memory size
-        -- 'x', creating it if it does not already exist.
-        ensureMemSizeOut (Ext x) = do
-          (memseen, arrseen) <- get
-          case M.lookup x memseen of
-            Nothing -> do sizeout <- imp $ newVName "out_memsize"
-                          tell ([Imp.ScalarParam sizeout int64],
-                                M.singleton x $ ScalarDestination sizeout)
-                          put (M.insert x sizeout memseen, arrseen)
-                          return $ Imp.VarSize sizeout
-            Just sizeout -> return $ Imp.VarSize sizeout
-        ensureMemSizeOut (Free v) = imp $ subExpToDimSize v
-
 compileFunDef :: ExplicitMemorish lore =>
                  FunDef lore
               -> ImpM lore op ()
@@ -496,7 +480,7 @@
           addArrays arrayds
 
           let Body _ stms ses = body
-          compileStms (freeIn ses) (stmsToList stms) $
+          compileStms (freeIn ses) stms $
             forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []
 
           return (outparams, inparams, results, args)
@@ -504,50 +488,49 @@
 compileBody :: (ExplicitMemorish lore) => Pattern lore -> Body lore -> ImpM lore op ()
 compileBody pat (Body _ bnds ses) = do
   Destination _ dests <- destinationFromPattern pat
-  compileStms (freeIn ses) (stmsToList bnds) $
+  compileStms (freeIn ses) bnds $
     forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []
 
 compileBody' :: (ExplicitMemorish lore, attr ~ LetAttr lore)
              => [Param attr] -> Body lore -> ImpM lore op ()
 compileBody' = compileBody . patternFromParams
 
-compileLoopBody :: [VName] -> Body lore -> ImpM lore op ()
-compileLoopBody mergenames (Body _ bnds ses) = do
+compileLoopBody :: Typed attr => [Param attr] -> Body lore -> ImpM lore op ()
+compileLoopBody mergeparams (Body _ bnds ses) = do
   -- We cannot write the results to the merge parameters immediately,
   -- as some of the results may actually *be* merge parameters, and
   -- would thus be clobbered.  Therefore, we first copy to new
   -- variables mirroring the merge parameters, and then copy this
   -- buffer to the merge parameters.  This is efficient, because the
   -- operations are all scalar operations.
-  tmpnames <- mapM (newVName . (++"_tmp") . baseString) mergenames
-  compileStms (freeIn ses) (stmsToList bnds) $ do
-    copy_to_merge_params <- forM (zip3 mergenames tmpnames ses) $ \(d,tmp,se) ->
-      subExpType se >>= \case
-        Prim bt  -> do
-          se' <- compileSubExp se
-          emit $ Imp.DeclareScalar tmp bt
-          emit $ Imp.SetScalar tmp se'
-          return $ emit $ Imp.SetScalar d $ Imp.var tmp bt
-        Mem _ space | Var v <- se -> do
+  tmpnames <- mapM (newVName . (++"_tmp") . baseString . paramName) mergeparams
+  compileStms (freeIn ses) bnds $ do
+    copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p,tmp,se) ->
+      case typeOf p of
+        Prim pt  -> do
+          emit $ Imp.DeclareScalar tmp pt
+          emit $ Imp.SetScalar tmp $ toExp' pt se
+          return $ emit $ Imp.SetScalar (paramName p) $ Imp.var tmp pt
+        Mem space | Var v <- se -> do
           emit $ Imp.DeclareMem tmp space
           emit $ Imp.SetMem tmp v space
-          return $ emit $ Imp.SetMem d tmp space
+          return $ emit $ Imp.SetMem (paramName p) tmp space
         _ -> return $ return ()
     sequence_ copy_to_merge_params
 
-compileStms :: Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()
+compileStms :: Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()
 compileStms alive_after_stms all_stms m = do
   cb <- asks envStmsCompiler
   cb alive_after_stms all_stms m
 
 defCompileStms :: (ExplicitMemorish lore, FreeIn op) =>
-                  Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op ()
+                  Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()
 defCompileStms alive_after_stms all_stms m =
   -- We keep track of any memory blocks produced by the statements,
   -- and after the last time that memory block is used, we insert a
   -- Free.  This is very conservative, but can cut down on lifetimes
   -- in some cases.
-  void $ compileStms' mempty all_stms
+  void $ compileStms' mempty $ stmsToList all_stms
   where compileStms' allocs (Let pat _ e:bs) = do
           dVars (Just e) (patternElements pat)
 
@@ -569,8 +552,8 @@
 
         patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements
         isMemPatElem pe = case patElemType pe of
-                            Mem _ space -> Just (patElemName pe, space)
-                            _           -> Nothing
+                            Mem space -> Just (patElemName pe, space)
+                            _         -> Nothing
 
 compileExp :: Pattern lore -> Exp lore -> ImpM lore op ()
 compileExp pat e = do
@@ -581,10 +564,9 @@
                  Pattern lore -> Exp lore -> ImpM lore op ()
 
 defCompileExp pat (If cond tbranch fbranch _) = do
-  cond' <- compileSubExp cond
   tcode <- collect $ compileBody pat tbranch
   fcode <- collect $ compileBody pat fbranch
-  emit $ Imp.If cond' tcode fcode
+  emit $ Imp.If (toExp' Bool cond) tcode fcode
 
 defCompileExp pat (Apply fname args _ _) = do
   dest <- destinationFromPattern pat
@@ -594,7 +576,7 @@
   where compileArg (se, _) = do
           t <- subExpType se
           case (se, t) of
-            (_, Prim pt)    -> return $ Just $ Imp.ExpArg $ compileSubExpOfType pt se
+            (_, Prim pt)   -> return $ Just $ Imp.ExpArg $ toExp' pt se
             (Var v, Mem{}) -> return $ Just $ Imp.MemArg v
             _              -> return Nothing
 
@@ -602,25 +584,23 @@
 
 defCompileExp pat (DoLoop ctx val form body) = do
   dFParams mergepat
-  forM_ merge $ \(p, se) -> do
-    na <- subExpNotArray se
-    when na $
-      copyDWIM (paramName p) [] se []
+  forM_ merge $ \(p, se) ->
+    when ((==0) $ arrayRank $ paramType p) $
+    copyDWIM (paramName p) [] se []
 
-  let doBody = compileLoopBody mergenames body
+  let doBody = compileLoopBody mergepat body
 
   case form of
     ForLoop i it bound loopvars -> do
-      bound' <- compileSubExp bound
-
       let setLoopParam (p,a)
             | Prim _ <- paramType p =
-                copyDWIM (paramName p) [] (Var a) [varIndex i]
+                copyDWIM (paramName p) [] (Var a) [Imp.vi32 i]
             | otherwise =
                 return ()
 
       dLParams $ map fst loopvars
-      sFor i it bound' $ mapM_ setLoopParam loopvars >> doBody
+      sFor i it (toExp' (IntType it) bound) $
+        mapM_ setLoopParam loopvars >> doBody
     WhileLoop cond ->
       sWhile (Imp.var cond Bool) doBody
 
@@ -630,7 +610,6 @@
 
   where merge = ctx ++ val
         mergepat = map fst merge
-        mergenames = map paramName mergepat
 
 defCompileExp pat (Op op) = do
   opc <- asks envOpCompiler
@@ -646,42 +625,42 @@
   copyDWIM (patElemName pe) [] se []
 
 defCompileBasicOp (Pattern _ [pe]) (UnOp op e) = do
-  e' <- compileSubExp e
+  e' <- toExp e
   patElemName pe <-- Imp.UnOpExp op e'
 
 defCompileBasicOp (Pattern _ [pe]) (ConvOp conv e) = do
-  e' <- compileSubExp e
+  e' <- toExp e
   patElemName pe <-- Imp.ConvOpExp conv e'
 
 defCompileBasicOp (Pattern _ [pe]) (BinOp bop x y) = do
-  x' <- compileSubExp x
-  y' <- compileSubExp y
+  x' <- toExp x
+  y' <- toExp y
   patElemName pe <-- Imp.BinOpExp bop x' y'
 
 defCompileBasicOp (Pattern _ [pe]) (CmpOp bop x y) = do
-  x' <- compileSubExp x
-  y' <- compileSubExp y
+  x' <- toExp x
+  y' <- toExp y
   patElemName pe <-- Imp.CmpOpExp bop x' y'
 
 defCompileBasicOp _ (Assert e msg loc) = do
-  e' <- compileSubExp e
-  msg' <- traverse compileSubExp msg
+  e' <- toExp e
+  msg' <- traverse toExp msg
   emit $ Imp.Assert e' msg' loc
 
 defCompileBasicOp (Pattern _ [pe]) (Index src slice)
   | Just idxs <- sliceIndices slice =
-      copyDWIM (patElemName pe) [] (Var src) $ map (compileSubExpOfType int32) idxs
+      copyDWIM (patElemName pe) [] (Var src) $ map (toExp' int32) idxs
 
 defCompileBasicOp _ Index{} =
   return ()
 
 defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) =
-  sUpdate (patElemName pe) (map (fmap (compileSubExpOfType int32)) slice) se
+  sUpdate (patElemName pe) (map (fmap (toExp' int32)) slice) se
 
 defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do
-  ds' <- mapM compileSubExp ds
+  ds' <- mapM toExp ds
   is <- replicateM (length ds) (newVName "i")
-  copy_elem <- collect $ copyDWIM (patElemName pe) (map varIndex is) se []
+  copy_elem <- collect $ copyDWIM (patElemName pe) (map Imp.vi32 is) se []
   emit $ foldl (.) id (zipWith (`Imp.For` Int32) is ds') copy_elem
 
 defCompileBasicOp _ Scratch{} =
@@ -690,14 +669,14 @@
 defCompileBasicOp (Pattern [] [pe]) (Iota n e s et) = do
   i <- newVName "i"
   x <- newVName "x"
-  n' <- compileSubExp n
-  e' <- compileSubExp e
-  s' <- compileSubExp s
+  n' <- toExp n
+  e' <- toExp e
+  s' <- toExp s
   let i' = ConvOpExp (SExt Int32 et) $ Imp.var i $ IntType Int32
   dPrim_ x $ IntType et
   sFor i Int32 n' $ do
     x <-- e' + i' * s'
-    copyDWIM (patElemName pe) [varIndex i] (Var x) []
+    copyDWIM (patElemName pe) [Imp.vi32 i] (Var x) []
 
 defCompileBasicOp (Pattern _ [pe]) (Copy src) =
   copyDWIM (patElemName pe) [] (Var src) []
@@ -708,14 +687,13 @@
 defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do
     MemLocation destmem destshape destixfun <-
       entryArrayLocation <$> lookupArray (patElemName pe)
-    xtype <- lookupType x
     offs_glb <- dPrim "tmp_offs" int32
     emit $ Imp.SetScalar offs_glb 0
     let perm = [i] ++ [0..i-1] ++ [i+1..length destshape-1]
         invperm = rearrangeInverse perm
         destloc = MemLocation destmem destshape
                   (IxFun.permute (IxFun.offsetIndex (IxFun.permute destixfun perm) $
-                                  varIndex offs_glb)
+                                  Imp.vi32 offs_glb)
                    invperm)
 
     forM_ (x:ys) $ \y -> do
@@ -724,7 +702,7 @@
           rows = case drop i $ entryArrayShape yentry of
                   []  -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y
                   r:_ -> innerExp $ Imp.dimSizeToExp r
-      copy (elemType xtype) destloc srcloc $ arrayOuterSize yentry
+      copy (elemType $ patElemType pe) destloc srcloc $ arrayOuterSize yentry
       emit $ Imp.SetScalar offs_glb $ Imp.var offs_glb int32 + rows
 
 defCompileBasicOp (Pattern [] [pe]) (ArrayLit es _)
@@ -736,13 +714,12 @@
       emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs
       let static_src = MemLocation static_array [Imp.ConstSize $ fromIntegral $ length es] $
                        IxFun.iota [fromIntegral $ length es]
-          num_bytes = Imp.ConstSize $ fromIntegral (length es) * primByteSize t
-          entry = MemVar Nothing $ MemEntry num_bytes dest_space
+          entry = MemVar Nothing $ MemEntry dest_space
       addVar static_array entry
       copy t dest_mem static_src $ fromIntegral $ length es
   | otherwise =
     forM_ (zip [0..] es) $ \(i,e) ->
-      copyDWIM (patElemName pe) [constIndex i] e []
+      copyDWIM (patElemName pe) [fromInteger i] e []
 
   where isLiteral (Constant v) = Just v
         isLiteral _ = Nothing
@@ -806,6 +783,10 @@
                   dPrim_ name' t
                   return name'
 
+dPrimV_ :: VName -> Imp.Exp -> ImpM lore op ()
+dPrimV_ name e = do dPrim_ name $ primExpType e
+                    name <-- e
+
 dPrimV :: String -> Imp.Exp -> ImpM lore op VName
 dPrimV name e = do name' <- dPrim name $ primExpType e
                    name' <-- e
@@ -815,14 +796,11 @@
                    -> ImpM lore op (VarEntry lore)
 memBoundToVarEntry e (MemPrim bt) =
   return $ ScalarVar e ScalarEntry { entryScalarType = bt }
-memBoundToVarEntry e (MemMem size space) = do
-  size' <- subExpToDimSize size
-  return $ MemVar e MemEntry { entryMemSize = size'
-                             , entryMemSpace = space
-                             }
+memBoundToVarEntry e (MemMem space) =
+  return $ MemVar e $ MemEntry space
 memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) = do
   shape' <- mapM subExpToDimSize $ shapeDims shape
-  let location = MemLocation mem shape' $ fmap compilePrimExp ixfun
+  let location = MemLocation mem shape' $ fmap (toExp' int32) ixfun
   return $ ArrayVar e ArrayEntry { entryArrayLocation = location
                                  , entryArrayElemType = bt
                                  }
@@ -879,30 +857,29 @@
 subExpToDimSize Constant{} =
   compilerBugS "Size subexp is not an int32 or int64 constant."
 
-compileSubExpTo :: VName -> SubExp -> ImpM lore op ()
-compileSubExpTo d se = copyDWIM d [] se []
-
-compileSubExp :: SubExp -> ImpM lore op Imp.Exp
-compileSubExp (Constant v) =
-  return $ Imp.ValueExp v
-compileSubExp (Var v) = do
-  t <- lookupType v
-  case t of
-    Prim pt -> return $ Imp.var v pt
-    _       -> compilerBugS $ "compileSubExp: SubExp is not a primitive type: " ++ pretty v
-
-compileSubExpOfType :: PrimType -> SubExp -> Imp.Exp
-compileSubExpOfType _ (Constant v) = Imp.ValueExp v
-compileSubExpOfType t (Var v) = Imp.var v t
+-- | Compile things to 'Imp.Exp'.
+class ToExp a where
+  -- | Compile to an 'Imp.Exp', where the type (must must still be a
+  -- primitive) is deduced monadically.
+  toExp :: a -> ImpM lore op Imp.Exp
+  -- | Compile where we know the type in advance.
+  toExp' :: PrimType -> a -> Imp.Exp
 
-compilePrimExp :: PrimExp VName -> Imp.Exp
-compilePrimExp = fmap Imp.ScalarVar
+instance ToExp SubExp where
+  toExp (Constant v) =
+    return $ Imp.ValueExp v
+  toExp (Var v) =
+    lookupVar v >>= \case
+    ScalarVar _ (ScalarEntry pt) ->
+      return $ Imp.var v pt
+    _       -> compilerBugS $ "toExp SubExp: SubExp is not a primitive type: " ++ pretty v
 
-varIndex :: VName -> Imp.Exp
-varIndex name = LeafExp (Imp.ScalarVar name) int32
+  toExp' _ (Constant v) = Imp.ValueExp v
+  toExp' t (Var v) = Imp.var v t
 
-constIndex :: Int -> Imp.Exp
-constIndex = fromIntegral
+instance ToExp (PrimExp VName) where
+  toExp = pure . fmap Imp.ScalarVar
+  toExp' _ = fmap Imp.ScalarVar
 
 addVar :: VName -> VarEntry lore -> ImpM lore op ()
 addVar name entry =
@@ -966,17 +943,17 @@
               return $ ScalarDestination name
 
 fullyIndexArray :: VName -> [Imp.Exp]
-                -> ImpM lore op (VName, Imp.Space, Count Bytes)
+                -> ImpM lore op (VName, Imp.Space, Count Elements)
 fullyIndexArray name indices = do
   arr <- lookupArray name
-  fullyIndexArray' (entryArrayLocation arr) indices $ entryArrayElemType arr
+  fullyIndexArray' (entryArrayLocation arr) indices
 
-fullyIndexArray' :: MemLocation -> [Imp.Exp] -> PrimType
-                 -> ImpM lore op (VName, Imp.Space, Count Bytes)
-fullyIndexArray' (MemLocation mem _ ixfun) indices bt = do
+fullyIndexArray' :: MemLocation -> [Imp.Exp]
+                 -> ImpM lore op (VName, Imp.Space, Count Elements)
+fullyIndexArray' (MemLocation mem _ ixfun) indices = do
   space <- entryMemSpace <$> lookupMemory mem
   return (mem, space,
-          bytes $ IxFun.index ixfun indices $ primByteSize bt)
+          elements $ IxFun.index ixfun indices)
 
 sliceArray :: MemLocation
            -> Slice Imp.Exp
@@ -999,11 +976,6 @@
 strideArray (MemLocation mem shape ixfun) stride =
   MemLocation mem shape $ IxFun.strideIndex ixfun stride
 
-subExpNotArray :: SubExp -> ImpM lore op Bool
-subExpNotArray se = subExpType se >>= \case
-  Array {} -> return False
-  _        -> return True
-
 arrayOuterSize :: ArrayEntry -> Count Elements
 arrayOuterSize = arrayDimSize 0
 
@@ -1045,17 +1017,16 @@
 copyElementWise :: CopyCompiler lore op
 copyElementWise bt (MemLocation destmem _ destIxFun) (MemLocation srcmem srcshape srcIxFun) n = do
     is <- replicateM (IxFun.rank destIxFun) (newVName "i")
-    let ivars = map varIndex is
-        destidx = IxFun.index destIxFun ivars bt_size
-        srcidx = IxFun.index srcIxFun ivars bt_size
+    let ivars = map Imp.vi32 is
+        destidx = IxFun.index destIxFun ivars
+        srcidx = IxFun.index srcIxFun ivars
         bounds = map innerExp $ n : drop 1 (map Imp.dimSizeToExp srcshape)
     srcspace <- entryMemSpace <$> lookupMemory srcmem
     destspace <- entryMemSpace <$> lookupMemory destmem
     vol <- asks envVolatility
     emit $ foldl (.) id (zipWith (`Imp.For` Int32) is bounds) $
-      Imp.Write destmem (bytes destidx) bt destspace vol $
-      Imp.index srcmem (bytes srcidx) bt srcspace vol
-  where bt_size = primByteSize bt
+      Imp.Write destmem (elements destidx) bt destspace vol $
+      Imp.index srcmem (elements srcidx) bt srcspace vol
 
 -- | Copy from here to there; both destination and source may be
 -- indexeded.
@@ -1069,9 +1040,9 @@
 
   | length srcis == length srcshape, length destis == length destshape = do
   (targetmem, destspace, targetoffset) <-
-    fullyIndexArray' destlocation destis bt
+    fullyIndexArray' destlocation destis
   (srcmem, srcspace, srcoffset) <-
-    fullyIndexArray' srclocation srcis bt
+    fullyIndexArray' srclocation srcis
   vol <- asks envVolatility
   return $ Imp.Write targetmem targetoffset bt destspace vol $
     Imp.index srcmem srcoffset bt srcspace vol
@@ -1083,7 +1054,15 @@
           srclocation'  =
             sliceArray srclocation $
             fullSliceNum (IxFun.shape src_ixfun) $ map DimFix srcis
-      if destlocation' == srclocation'
+          destrank = length (memLocationShape destlocation')
+          srcrank = length (memLocationShape srclocation')
+      if destrank /= srcrank
+        then fail $ "copyArrayDWIM: cannot copy to " ++
+             pretty (memLocationName destlocation') ++
+             " from " ++ pretty (memLocationName srclocation') ++
+             " because ranks do not match (" ++ pretty destrank ++
+             " vs " ++ pretty srcrank ++ ")"
+      else if destlocation' == srclocation'
         then return mempty -- Copy would be no-op.
         else collect $ copy bt destlocation' srclocation' $
              product $ map Imp.dimSizeToExp $
@@ -1106,7 +1085,7 @@
     unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."]
   ArrayDestination (Just dest_loc) -> do
     (dest_mem, dest_space, dest_i) <-
-      fullyIndexArray' dest_loc dest_is bt
+      fullyIndexArray' dest_loc dest_is
     vol <- asks envVolatility
     emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v
   ArrayDestination Nothing ->
@@ -1116,7 +1095,7 @@
 copyDWIMDest dest dest_is (Var src) src_is = do
   src_entry <- lookupVar src
   case (dest, src_entry) of
-    (MemoryDestination mem, MemVar _ (MemEntry _ space)) ->
+    (MemoryDestination mem, MemVar _ (MemEntry space)) ->
       emit $ Imp.SetMem mem src space
 
     (MemoryDestination{}, _) ->
@@ -1142,7 +1121,7 @@
     (ScalarDestination name, ArrayVar _ arr) -> do
       let bt = entryArrayElemType arr
       (mem, space, i) <-
-        fullyIndexArray' (entryArrayLocation arr) src_is bt
+        fullyIndexArray' (entryArrayLocation arr) src_is
       vol <- asks envVolatility
       emit $ Imp.SetScalar name $ Imp.index mem i bt space vol
 
@@ -1152,8 +1131,7 @@
       emit =<< copyArrayDWIM bt dest_loc dest_is src_loc src_is
 
     (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt)) -> do
-      (dest_mem, dest_space, dest_i) <-
-        fullyIndexArray' dest_loc dest_is bt
+      (dest_mem, dest_space, dest_i) <- fullyIndexArray' dest_loc dest_is
       vol <- asks envVolatility
       emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt)
 
@@ -1188,9 +1166,11 @@
                 Pattern lore -> SubExp -> Space
              -> ImpM lore op ()
 compileAlloc (Pattern [] [mem]) e space = do
-  e' <- compileSubExp e
-  fake <- asks $ elem space . envFakeMemory
-  unless fake $ emit $ Imp.Allocate (patElemName mem) (Imp.bytes e') space
+  e' <- Imp.bytes <$> toExp e
+  allocator <- asks $ M.lookup space . envAllocCompilers
+  case allocator of
+    Nothing -> emit $ Imp.Allocate (patElemName mem) e' space
+    Just allocator' -> allocator' (patElemName mem) e'
 compileAlloc pat _ _ =
   compilerBugS $ "compileAlloc: Invalid pattern: " ++ pretty pat
 
@@ -1199,8 +1179,14 @@
 dimSizeToSubExp (Imp.VarSize v) = Var v
 
 dimSizeToExp :: Imp.Size -> Imp.Exp
-dimSizeToExp = compilePrimExp . primExpFromSubExp int32 . dimSizeToSubExp
+dimSizeToExp = toExp' int32 . primExpFromSubExp int32 . dimSizeToSubExp
 
+-- | The number of bytes needed to represent the array in a
+-- straightforward contiguous format.
+typeSize :: Type -> Count Bytes
+typeSize t = Imp.bytes $ Imp.LeafExp (Imp.SizeOf $ elemType t) int32 *
+             product (map (toExp' int32) (arrayDims t))
+
 --- Building blocks for constructing code.
 
 sFor :: VName -> IntType -> Imp.Exp -> ImpM lore op () -> ImpM lore op ()
@@ -1234,33 +1220,24 @@
 sOp :: op -> ImpM lore op ()
 sOp = emit . Imp.Op
 
-dSize :: Imp.Count u -> ImpM lore op Imp.Size
-dSize size =
-  case Imp.innerExp size of
-    Imp.LeafExp (Imp.ScalarVar size') _ -> return $ Imp.VarSize size'
-    Imp.ValueExp (IntValue (Int64Value v)) -> return $ Imp.ConstSize v
-    _ -> do size_var <- dPrim "local_buf_size" int32
-            size_var <-- Imp.innerExp size
-            return $ Imp.VarSize size_var
-
-sDeclareMem :: String -> Count Bytes -> Space -> ImpM lore op (VName, Imp.MemSize)
-sDeclareMem name size space = do
-
+sDeclareMem :: String -> Space -> ImpM lore op VName
+sDeclareMem name space = do
   name' <- newVName name
-  size' <- dSize size
   emit $ Imp.DeclareMem name' space
-  addVar name' $ MemVar Nothing $ MemEntry size' space
-  return (name', size')
+  addVar name' $ MemVar Nothing $ MemEntry space
+  return name'
 
-sAlloc_ :: VName -> Imp.MemSize -> Space -> ImpM lore op ()
+sAlloc_ :: VName -> Count Bytes -> Space -> ImpM lore op ()
 sAlloc_ name' size' space = do
-  fake <- asks $ elem space . envFakeMemory
-  unless fake $ emit $ Imp.Allocate name' (Imp.memSizeToExp size') space
+  allocator <- asks $ M.lookup space . envAllocCompilers
+  case allocator of
+    Nothing -> emit $ Imp.Allocate name' size' space
+    Just allocator' -> allocator' name' size'
 
 sAlloc :: String -> Count Bytes -> Space -> ImpM lore op VName
 sAlloc name size space = do
-  (name', size') <- sDeclareMem name size space
-  sAlloc_ name' size' space
+  name' <- sDeclareMem name space
+  sAlloc_ name' size space
   return name'
 
 sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore op VName
@@ -1269,25 +1246,29 @@
   dArray name' bt shape membind
   return name'
 
+-- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified.
+sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore op VName
+sAllocArrayPerm name pt shape space perm = do
+  let permuted_dims = rearrangeShape perm $ shapeDims shape
+  mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space
+  let iota_ixfun = IxFun.iota $ map (primExpFromSubExp int32) permuted_dims
+  sArray name pt shape $
+    ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm
+
 -- | Uses linear/iota index function.
 sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore op VName
-sAllocArray name pt shape space = do
-  let arr_bytes = Imp.bytes $ Imp.LeafExp (Imp.SizeOf pt) int32 *
-                  product (map (compileSubExpOfType int32) (shapeDims shape))
-  mem <- sAlloc (name ++ "_mem") arr_bytes space
-  sArray name pt shape $
-    ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape
+sAllocArray name pt shape space =
+  sAllocArrayPerm name pt shape space [0..shapeRank shape-1]
 
 -- | Uses linear/iota index function.
 sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore op VName
 sStaticArray name space pt vs = do
-  let num_elems = case vs of Imp.ArrayValues vs' -> genericLength vs'
+  let num_elems = case vs of Imp.ArrayValues vs' -> length vs'
                              Imp.ArrayZeros n -> fromIntegral n
-      shape = Shape [constant num_elems]
-      mem_size = Imp.ConstSize $ num_elems * primByteSize pt
+      shape = Shape [intConst Int32 $ toInteger num_elems]
   mem <- newVName $ name ++ "_mem"
   emit $ Imp.DeclareArray mem space pt vs
-  addVar mem $ MemVar Nothing $ MemEntry mem_size space
+  addVar mem $ MemVar Nothing $ MemEntry space
   sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems]
 
 sWrite :: VName -> [Imp.Exp] -> PrimExp Imp.ExpLeaf -> ImpM lore op ()
@@ -1301,6 +1282,16 @@
   MemLocation mem shape ixfun <- entryArrayLocation <$> lookupArray arr
   let memdest = sliceArray (MemLocation mem shape ixfun) slice
   copyDWIMDest (ArrayDestination $ Just memdest) [] v []
+
+sLoopNest :: Shape
+          -> ([Imp.Exp] -> ImpM lore op ())
+          -> ImpM lore op ()
+sLoopNest = sLoopNest' [] . shapeDims
+  where sLoopNest' is [] f = f $ reverse is
+        sLoopNest' is (d:ds) f = do
+          i <- newVName "nest_i"
+          d' <- toExp d
+          sFor i Int32 d' $ sLoopNest' (Imp.var i int32:is) ds f
 
 -- | ASsignment.
 (<--) :: VName -> Imp.Exp -> ImpM lore op ()
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels.hs
@@ -3,14 +3,13 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE ConstraintKinds #-}
 module Futhark.CodeGen.ImpGen.Kernels
-  ( compileProg
+  ( Futhark.CodeGen.ImpGen.Kernels.compileProg
   )
   where
 
 import Control.Monad.Except
 import Control.Monad.Reader
 import Data.Maybe
-import qualified Data.Map.Strict as M
 import Data.List
 
 import Prelude hiding (quot)
@@ -20,49 +19,50 @@
 import Futhark.Representation.ExplicitMemory
 import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
 import Futhark.CodeGen.ImpCode.Kernels (bytes)
-import qualified Futhark.CodeGen.ImpGen as ImpGen
+import Futhark.CodeGen.ImpGen
 import Futhark.CodeGen.ImpGen.Kernels.Base
+import Futhark.CodeGen.ImpGen.Kernels.SegMap
 import Futhark.CodeGen.ImpGen.Kernels.SegRed
+import Futhark.CodeGen.ImpGen.Kernels.SegScan
 import Futhark.CodeGen.ImpGen.Kernels.SegGenRed
-import Futhark.CodeGen.ImpGen (sFor, sWhen,
-                               sOp)
 import Futhark.CodeGen.ImpGen.Kernels.Transpose
 import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
 import Futhark.CodeGen.SetDefaultSpace
-import Futhark.Util.IntegralExp (quotRoundingUp, quot, IntegralExp)
+import Futhark.Util.IntegralExp (quot, IntegralExp)
 
-callKernelOperations :: ImpGen.Operations ExplicitMemory Imp.HostOp
+callKernelOperations :: Operations ExplicitMemory Imp.HostOp
 callKernelOperations =
-  ImpGen.Operations { ImpGen.opsExpCompiler = expCompiler
-                    , ImpGen.opsCopyCompiler = callKernelCopy
-                    , ImpGen.opsOpCompiler = opCompiler
-                    , ImpGen.opsStmsCompiler = ImpGen.defCompileStms
+  Operations { opsExpCompiler = expCompiler
+                    , opsCopyCompiler = callKernelCopy
+                    , opsOpCompiler = opCompiler
+                    , opsStmsCompiler = defCompileStms
+                    , opsAllocCompilers = mempty
                     }
 
 compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program)
 compileProg prog =
   fmap (setDefaultSpace (Imp.Space "device")) <$>
-  ImpGen.compileProg callKernelOperations (Imp.Space "device") [Imp.Space "local"] prog
+  Futhark.CodeGen.ImpGen.compileProg callKernelOperations (Imp.Space "device") prog
 
 opCompiler :: Pattern ExplicitMemory -> Op ExplicitMemory
            -> CallKernelGen ()
 opCompiler dest (Alloc e space) =
-  ImpGen.compileAlloc dest e space
+  compileAlloc dest e space
 opCompiler (Pattern _ [pe]) (Inner (GetSize key size_class)) = do
-  fname <- asks ImpGen.envFunction
+  fname <- asks envFunction
   sOp $ Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $
     sizeClassWithEntryPoint fname size_class
 opCompiler (Pattern _ [pe]) (Inner (CmpSizeLe key size_class x)) = do
-  fname <- asks ImpGen.envFunction
+  fname <- asks envFunction
   let size_class' = sizeClassWithEntryPoint fname size_class
   sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class'
-    =<< ImpGen.compileSubExp x
+    =<< toExp x
 opCompiler (Pattern _ [pe]) (Inner (GetSizeMax size_class)) =
   sOp $ Imp.GetSizeMax (patElemName pe) size_class
 opCompiler dest (Inner (HostOp kernel)) =
   kernelCompiler dest kernel
 opCompiler pat e =
-  compilerBugS $ "ImpGen.opCompiler: Invalid pattern\n  " ++
+  compilerBugS $ "opCompiler: Invalid pattern\n  " ++
   pretty pat ++ "\nfor expression\n  " ++ pretty e
 
 sizeClassWithEntryPoint :: Name -> Imp.SizeClass -> Imp.SizeClass
@@ -75,18 +75,7 @@
                -> CallKernelGen ()
 
 kernelCompiler pat (Kernel desc space _ kernel_body) = do
-  (constants, init_constants) <- kernelInitialisation space
-
-  kernel_body' <-
-    makeAllMemoryGlobal $ ImpGen.subImpM_ (inKernelOperations constants) $ do
-    init_constants
-    compileKernelBody pat constants kernel_body
-
-  let bound_in_kernel =
-        M.keys $
-        scopeOfKernelSpace space <>
-        scopeOf (kernelBodyStms kernel_body)
-  (uses, local_memory) <- computeKernelUses kernel_body' bound_in_kernel
+  (constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
 
   forM_ (kernelHints desc) $ \(s,v) -> do
     ty <- case v of
@@ -96,31 +85,40 @@
                                          , " in kernel '", kernelName desc, "'"
                                          , " did not have primType value." ]
 
-    ImpGen.compileSubExp v >>= ImpGen.emit . Imp.DebugPrint s (elemType ty)
+    emit $ Imp.DebugPrint s $ Just (elemType ty, toExp' (elemType ty) v)
 
-  sOp $ Imp.CallKernel Imp.Kernel
-            { Imp.kernelBody = kernel_body'
-            , Imp.kernelLocalMemory = local_memory
-            , Imp.kernelUses = uses
-            , Imp.kernelNumGroups = [ImpGen.compileSubExpOfType int32 $ spaceNumGroups space]
-            , Imp.kernelGroupSize = [ImpGen.compileSubExpOfType int32 $ spaceGroupSize space]
-            , Imp.kernelName = nameFromString $ kernelName desc ++ "_" ++
-                               show (baseTag $ kernelGlobalThreadIdVar constants)
-            }
+  let virt_groups = toExp' int32 (spaceNumVirtGroups space)
+  sKernel constants (kernelName desc) $ do
+    init_constants
+    virtualiseGroups constants virt_groups $ \group_id -> do
+      let flat_id =
+            if kernelGroupIdVar constants /= group_id
+            then Imp.vi32 group_id * kernelGroupSize constants + kernelLocalThreadId constants
+            else kernelGlobalThreadId constants
+      setSpaceIndices flat_id space
+      compileKernelStms constants (kernelBodyStms kernel_body) $
+        zipWithM_ (compileKernelResult constants) (patternElements pat) $
+        kernelBodyResult kernel_body
 
-kernelCompiler pat (SegRed space comm red_op nes _ body) =
-  compileSegRed pat space comm red_op nes body
+kernelCompiler pat (SegMap space _ body) =
+  compileSegMap pat space body
 
+kernelCompiler pat (SegRed space reds _ body) =
+  compileSegRed pat space reds body
+
+kernelCompiler pat (SegScan space red_op nes _ kbody) =
+  compileSegScan pat space red_op nes kbody
+
 kernelCompiler pat (SegGenRed space ops _ body) =
   compileSegGenRed pat space ops body
 
-expCompiler :: ImpGen.ExpCompiler ExplicitMemory Imp.HostOp
+expCompiler :: ExpCompiler ExplicitMemory Imp.HostOp
 
 -- We generate a simple kernel for itoa and replicate.
 expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do
-  n' <- ImpGen.compileSubExp n
-  x' <- ImpGen.compileSubExp x
-  s' <- ImpGen.compileSubExp s
+  n' <- toExp n
+  x' <- toExp x
+  s' <- toExp s
 
   sIota (patElemName pe) n' x' s' et
 
@@ -132,19 +130,19 @@
   return ()
 
 expCompiler dest e =
-  ImpGen.defCompileExp dest e
+  defCompileExp dest e
 
-callKernelCopy :: ImpGen.CopyCompiler ExplicitMemory Imp.HostOp
+callKernelCopy :: CopyCompiler ExplicitMemory Imp.HostOp
 callKernelCopy bt
-  destloc@(ImpGen.MemLocation destmem destshape destIxFun)
-  srcloc@(ImpGen.MemLocation srcmem srcshape srcIxFun)
+  destloc@(MemLocation destmem destshape destIxFun)
+  srcloc@(MemLocation srcmem srcshape srcIxFun)
   n
   | Just (destoffset, srcoffset,
           num_arrays, size_x, size_y,
           src_elems, dest_elems) <- isMapTransposeKernel bt destloc srcloc = do
 
       fname <- mapTransposeForType bt
-      ImpGen.emit $ Imp.Call [] fname
+      emit $ Imp.Call [] fname
         [Imp.MemArg destmem, Imp.ExpArg destoffset,
          Imp.MemArg srcmem, Imp.ExpArg srcoffset,
          Imp.ExpArg num_arrays, Imp.ExpArg size_x, Imp.ExpArg size_y,
@@ -159,25 +157,25 @@
       IxFun.linearWithOffset destIxFun bt_size,
     Just srcoffset  <-
       IxFun.linearWithOffset srcIxFun bt_size = do
-        let row_size = product $ map ImpGen.dimSizeToExp $ drop 1 srcshape
-        srcspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory srcmem
-        destspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory destmem
-        ImpGen.emit $ Imp.Copy
+        let row_size = product $ map dimSizeToExp $ drop 1 srcshape
+        srcspace <- entryMemSpace <$> lookupMemory srcmem
+        destspace <- entryMemSpace <$> lookupMemory destmem
+        emit $ Imp.Copy
           destmem (bytes destoffset) destspace
           srcmem (bytes srcoffset) srcspace $
           (n * row_size) `Imp.withElemType` bt
 
   | otherwise = sCopy bt destloc srcloc n
 
-mapTransposeForType :: PrimType -> ImpGen.ImpM ExplicitMemory Imp.HostOp Name
+mapTransposeForType :: PrimType -> ImpM ExplicitMemory Imp.HostOp Name
 mapTransposeForType bt = do
   -- XXX: The leading underscore is to avoid clashes with a
   -- programmer-defined function of the same name (this is a bad
   -- solution...).
   let fname = nameFromString $ "_" <> mapTransposeName bt
 
-  exists <- ImpGen.hasFunction fname
-  unless exists $ ImpGen.emitFunction fname $ mapTransposeFunction bt
+  exists <- hasFunction fname
+  unless exists $ emitFunction fname $ mapTransposeFunction bt
 
   return fname
 
@@ -288,13 +286,13 @@
             v32 mulx, v32 muly, v32 num_arrays,
             block) bt
 
-isMapTransposeKernel :: PrimType -> ImpGen.MemLocation -> ImpGen.MemLocation
+isMapTransposeKernel :: PrimType -> MemLocation -> MemLocation
                      -> Maybe (Imp.Exp, Imp.Exp,
                                Imp.Exp, Imp.Exp, Imp.Exp,
                                Imp.Exp, Imp.Exp)
 isMapTransposeKernel bt
-  (ImpGen.MemLocation _ _ destIxFun)
-  (ImpGen.MemLocation _ _ srcIxFun)
+  (MemLocation _ _ destIxFun)
+  (MemLocation _ _ srcIxFun)
   | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun bt_size,
     (perm, destshape) <- unzip perm_and_destshape,
     srcshape' <- IxFun.shape srcIxFun,
@@ -322,100 +320,3 @@
           let (mapped, notmapped) = splitAt r1 shape
               (pretrans, posttrans) = f $ splitAt r2 notmapped
           in (product mapped, product pretrans, product posttrans)
-
-compileKernelBody :: Pattern InKernel
-                  -> KernelConstants
-                  -> KernelBody InKernel
-                  -> InKernelGen ()
-compileKernelBody pat constants kbody =
-  compileKernelStms constants (stmsToList $ kernelBodyStms kbody) $
-  zipWithM_ (compileKernelResult constants) (patternElements pat) $
-  kernelBodyResult kbody
-
-compileKernelResult :: KernelConstants -> PatElem InKernel -> KernelResult
-                    -> InKernelGen ()
-
-compileKernelResult constants pe (ThreadsReturn OneResultPerGroup what) = do
-  i <- newVName "i"
-
-  in_local_memory <- arrayInLocalMemory what
-  let me = kernelLocalThreadId constants
-
-  if not in_local_memory then do
-    who' <- ImpGen.compileSubExp $ intConst Int32 0
-    sWhen (me .==. who') $
-      ImpGen.copyDWIM (patElemName pe) [kernelGroupId constants] what []
-    else do
-      -- If the result of the group is an array in local memory, we
-      -- store it by collective copying among all the threads of the
-      -- group.  TODO: also do this if the array is in global memory
-      -- (but this is a bit more tricky, synchronisation-wise).
-      --
-      -- We do the reads/writes multidimensionally, but the loop is
-      -- single-dimensional.
-      ws <- mapM ImpGen.compileSubExp . arrayDims =<< subExpType what
-      -- Compute how many elements this thread is responsible for.
-      -- Formula: (w - ltid) / group_size (rounded up).
-      let w = product ws
-          ltid = kernelLocalThreadId constants
-          group_size = kernelGroupSize constants
-          to_write = (w - ltid) `quotRoundingUp` group_size
-          is = unflattenIndex ws $ ImpGen.varIndex i * group_size + ltid
-
-      sFor i Int32 to_write $
-        ImpGen.copyDWIM (patElemName pe) (kernelGroupId constants : is) what is
-
-compileKernelResult constants pe (ThreadsReturn AllThreads what) =
-  ImpGen.copyDWIM (patElemName pe) [kernelGlobalThreadId constants] what []
-
-compileKernelResult constants pe (ThreadsReturn (ThreadsPerGroup limit) what) =
-  sWhen (isActive limit) $
-  ImpGen.copyDWIM (patElemName pe) [kernelGroupId constants] what []
-
-compileKernelResult constants pe (ThreadsReturn ThreadsInSpace what) = do
-  let is = map (ImpGen.varIndex . fst) $ kernelDimensions constants
-  sWhen (kernelThreadActive constants) $ ImpGen.copyDWIM (patElemName pe) is what []
-
-compileKernelResult constants pe (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do
-  dest_loc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray (patElemName pe)
-  let dest_loc_offset = ImpGen.offsetArray dest_loc offset
-      dest' = ImpGen.arrayDestination dest_loc_offset
-  ImpGen.copyDWIMDest dest' [] (Var what) []
-  where offset = case moffset of
-                   Nothing -> ImpGen.compileSubExpOfType int32 per_thread_elems *
-                              kernelGlobalThreadId constants
-                   Just se -> ImpGen.compileSubExpOfType int32 se
-
-compileKernelResult constants pe (ConcatReturns (SplitStrided stride) _ _ moffset what) = do
-  dest_loc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray (patElemName pe)
-  let dest_loc' = ImpGen.strideArray
-                  (ImpGen.offsetArray dest_loc offset) $
-                  ImpGen.compileSubExpOfType int32 stride
-      dest' = ImpGen.arrayDestination dest_loc'
-  ImpGen.copyDWIMDest dest' [] (Var what) []
-  where offset = case moffset of
-                   Nothing -> kernelGlobalThreadId constants
-                   Just se -> ImpGen.compileSubExpOfType int32 se
-
-compileKernelResult constants pe (WriteReturn rws _arr dests) = do
-  rws' <- mapM ImpGen.compileSubExp rws
-  forM_ dests $ \(is, e) -> do
-    is' <- mapM ImpGen.compileSubExp is
-    let condInBounds i rw = 0 .<=. i .&&. i .<. rw
-        write = foldl (.&&.) (kernelThreadActive constants) $
-                zipWith condInBounds is' rws'
-    sWhen write $ ImpGen.copyDWIM (patElemName pe) (map (ImpGen.compileSubExpOfType int32) is) e []
-
-compileKernelResult _ _ KernelInPlaceReturn{} =
-  -- Already in its place... said it was a hack.
-  return ()
-
-arrayInLocalMemory :: SubExp -> InKernelGen Bool
-arrayInLocalMemory (Var name) = do
-  res <- ImpGen.lookupVar name
-  case res of
-    ImpGen.ArrayVar _ entry ->
-      (Space "local"==) . ImpGen.entryMemSpace <$>
-      ImpGen.lookupMemory (ImpGen.memLocationName (ImpGen.entryArrayLocation entry))
-    _ -> return False
-arrayInLocalMemory Constant{} = return False
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs
@@ -4,12 +4,13 @@
 module Futhark.CodeGen.ImpGen.Kernels.Base
   ( KernelConstants (..)
   , inKernelOperations
-  , computeKernelUses
   , keyWithEntryPoint
   , CallKernelGen
   , InKernelGen
   , computeThreadChunkSize
+  , simpleKernelConstants
   , kernelInitialisation
+  , kernelInitialisationSimple
   , kernelInitialisationSetSpace
   , setSpaceIndices
   , makeAllMemoryGlobal
@@ -22,6 +23,8 @@
   , sReplicate
   , sIota
   , sCopy
+  , compileKernelResult
+  , virtualiseGroups
 
   , atomicUpdate
   , atomicUpdateLocking
@@ -45,21 +48,17 @@
 import Futhark.Transform.Rename
 import Futhark.Representation.ExplicitMemory
 import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
-import Futhark.CodeGen.ImpCode.Kernels (bytes)
-import qualified Futhark.CodeGen.ImpGen as ImpGen
-import Futhark.CodeGen.ImpGen ((<--),
-                               sFor, sWhile, sComment, sIf, sWhen, sUnless,
-                               sOp,
-                               dPrim, dPrim_, dPrimV)
-import Futhark.Tools (partitionChunkedKernelLambdaParameters)
+import Futhark.CodeGen.ImpCode.Kernels (elements)
+import Futhark.CodeGen.ImpGen
 import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem, IntegralExp)
 import Futhark.Util (splitAt3, maybeNth, takeLast)
 
-type CallKernelGen = ImpGen.ImpM ExplicitMemory Imp.HostOp
-type InKernelGen = ImpGen.ImpM InKernel Imp.KernelOp
+type CallKernelGen = ImpM ExplicitMemory Imp.HostOp
+type InKernelGen = ImpM InKernel Imp.KernelOp
 
 data KernelConstants = KernelConstants
-                       { kernelGlobalThreadId :: Imp.Exp
+                       { kernelOuterVTable :: VTable ExplicitMemory -- XXX
+                       , kernelGlobalThreadId :: Imp.Exp
                        , kernelLocalThreadId :: Imp.Exp
                        , kernelGroupId :: Imp.Exp
                        , kernelGlobalThreadIdVar :: VName
@@ -76,12 +75,22 @@
                        -- for unrolling.
                        }
 
-inKernelOperations :: KernelConstants -> ImpGen.Operations InKernel Imp.KernelOp
-inKernelOperations constants = (ImpGen.defaultOperations $ compileInKernelOp constants)
-                               { ImpGen.opsCopyCompiler = inKernelCopy
-                               , ImpGen.opsExpCompiler = inKernelExpCompiler
-                               , ImpGen.opsStmsCompiler = \_ -> compileKernelStms constants
-                               }
+inKernelOperations :: KernelConstants -> Operations InKernel Imp.KernelOp
+inKernelOperations constants =
+  (defaultOperations $ compileInKernelOp constants)
+  { opsCopyCompiler = inKernelCopy
+  , opsExpCompiler = inKernelExpCompiler
+  , opsStmsCompiler = \_ -> compileKernelStms constants
+  , opsAllocCompilers =
+      M.fromList [ (Space "local", allocLocal)
+                 , (Space "private", allocPrivate) ]
+  }
+  where allocLocal :: AllocCompiler InKernel Imp.KernelOp
+        allocLocal mem size = do
+          size' <- localMemSize (kernelOuterVTable constants) size
+          sOp $ Imp.LocalAlloc mem size'
+        allocPrivate mem size =
+          sOp $ Imp.PrivateAlloc mem size
 
 keyWithEntryPoint :: Name -> Name -> Name
 keyWithEntryPoint fname key =
@@ -89,11 +98,18 @@
 
 -- | We have no bulk copy operation (e.g. memmove) inside kernels, so
 -- turn any copy into a loop.
-inKernelCopy :: ImpGen.CopyCompiler InKernel Imp.KernelOp
-inKernelCopy = ImpGen.copyElementWise
+inKernelCopy :: CopyCompiler InKernel Imp.KernelOp
+inKernelCopy = copyElementWise
 
 compileInKernelOp :: KernelConstants -> Pattern InKernel -> Op InKernel
                   -> InKernelGen ()
+compileInKernelOp _ (Pattern _ [mem]) (Alloc size (Space "private")) = do
+  size' <- toExp size
+  sOp $ Imp.PrivateAlloc (patElemName mem) $ Imp.bytes size'
+compileInKernelOp constants (Pattern _ [mem]) (Alloc size (Space "local")) = do
+  size' <- localMemSize (kernelOuterVTable constants) . Imp.bytes =<<
+           toExp size
+  sOp $ Imp.LocalAlloc (patElemName mem) size'
 compileInKernelOp _ (Pattern _ [mem]) Alloc{} =
   compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel."
 compileInKernelOp _ dest Alloc{} =
@@ -101,7 +117,7 @@
 compileInKernelOp constants pat (Inner op) =
   compileKernelExp constants pat op
 
-inKernelExpCompiler :: ImpGen.ExpCompiler InKernel Imp.KernelOp
+inKernelExpCompiler :: ExpCompiler InKernel Imp.KernelOp
 inKernelExpCompiler _ (BasicOp (Assert _ _ (loc, locs))) =
   compilerLimitationS $
   unlines [ "Cannot compile assertion at " ++
@@ -111,22 +127,22 @@
 -- The static arrays stuff does not work inside kernels.
 inKernelExpCompiler (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =
   forM_ (zip [0..] es) $ \(i,e) ->
-  ImpGen.copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e []
+  copyDWIM (patElemName dest) [fromIntegral (i::Int32)] e []
 inKernelExpCompiler dest e =
-  ImpGen.defCompileExp dest e
+  defCompileExp dest e
 
 compileKernelExp :: KernelConstants -> Pattern InKernel -> KernelExp InKernel
                  -> InKernelGen ()
 
 compileKernelExp _ pat (Barrier ses) = do
   forM_ (zip (patternNames pat) ses) $ \(d, se) ->
-    ImpGen.copyDWIM d [] se []
+    copyDWIM d [] se []
   sOp Imp.LocalBarrier
 
 compileKernelExp _ (Pattern [] [size]) (SplitSpace o w i elems_per_thread) = do
-  num_elements <- Imp.elements <$> ImpGen.compileSubExp w
-  i' <- ImpGen.compileSubExp i
-  elems_per_thread' <- Imp.elements <$> ImpGen.compileSubExp elems_per_thread
+  num_elements <- Imp.elements <$> toExp w
+  i' <- toExp i
+  elems_per_thread' <- Imp.elements <$> toExp elems_per_thread
   computeThreadChunkSize o i' elems_per_thread' num_elements (patElemName size)
 
 compileKernelExp constants pat (Combine (CombineSpace scatter cspace) _ aspace body) = do
@@ -166,56 +182,51 @@
 
     -- Execute the body if we are within bounds.
     sWhen (isActive cspace .&&. isActive aspace) $ allThreads constants $
-      ImpGen.compileStms (freeIn $ bodyResult body) (stmsToList $ bodyStms body) $ do
+      compileStms (freeIn $ bodyResult body) (bodyStms body) $ do
 
       forM_ (zip4 scatter_ws_repl res_is res_vs scatter_pes) $
         \(w, res_i, res_v, scatter_pe) -> do
-          let res_i' = ImpGen.compileSubExpOfType int32 res_i
-              w'     = ImpGen.compileSubExpOfType int32 w
+          let res_i' = toExp' int32 res_i
+              w'     = toExp' int32 w
               -- We have to check that 'res_i' is in-bounds wrt. an array of size 'w'.
               in_bounds = 0 .<=. res_i' .&&. res_i' .<. w'
-          sWhen in_bounds $ ImpGen.copyDWIM (patElemName scatter_pe) [res_i'] res_v []
+          sWhen in_bounds $ copyDWIM (patElemName scatter_pe) [res_i'] res_v []
 
       forM_ (zip normal_pes res_normal) $ \(pe, res) ->
-        ImpGen.copyDWIM (patElemName pe) local_index res []
+        copyDWIM (patElemName pe) local_index res []
 
   sOp Imp.LocalBarrier
 
   where streamBounded (Var v)
           | Just x <- lookup v $ kernelStreamed constants =
               Imp.sizeToExp x
-        streamBounded se = ImpGen.compileSubExpOfType int32 se
+        streamBounded se = toExp' int32 se
 
-        local_index = map (ImpGen.compileSubExpOfType int32 . Var . fst) cspace
+        local_index = map (toExp' int32 . Var . fst) cspace
 
 compileKernelExp constants (Pattern _ dests) (GroupReduce w lam input) = do
-  let [my_index_param, offset_param] = take 2 $ lambdaParams lam
-      lam' = lam { lambdaParams = drop 2 $ lambdaParams lam }
-
-  dPrim_ (paramName my_index_param) int32
-  dPrim_ (paramName offset_param) int32
-  paramName my_index_param <-- kernelGlobalThreadId constants
-  w' <- ImpGen.compileSubExp w
-  groupReduceWithOffset constants (paramName offset_param) w' lam' $ map snd input
+  w' <- toExp w
+  groupReduce constants w' lam $ map snd input
 
   sOp Imp.LocalBarrier
 
   -- The final result will be stored in element 0 of the local memory array.
   forM_ (zip dests input) $ \(dest, (_, arr)) ->
-    ImpGen.copyDWIM (patElemName dest) [] (Var arr) [0]
+    copyDWIM (patElemName dest) [] (Var arr) [0]
 
 compileKernelExp constants _ (GroupScan w lam input) = do
-  w' <- ImpGen.compileSubExp w
+  w' <- toExp w
   groupScan constants Nothing w' lam $ map snd input
 
 compileKernelExp constants (Pattern _ final) (GroupStream w maxchunk lam accs _arrs) = do
   let GroupStreamLambda block_size block_offset acc_params arr_params body = lam
       block_offset' = Imp.var block_offset int32
-  w' <- ImpGen.compileSubExp w
-  max_block_size <- ImpGen.compileSubExp maxchunk
+  w' <- toExp w
+  max_block_size <- toExp maxchunk
 
-  ImpGen.dLParams (acc_params++arr_params)
-  zipWithM_ ImpGen.compileSubExpTo (map paramName acc_params) accs
+  dLParams (acc_params++arr_params)
+  forM_ (zip acc_params accs) $ \(p, acc) ->
+    copyDWIM (paramName p) [] acc []
   dPrim_ block_size int32
 
   -- If the GroupStream is morally just a do-loop, generate simpler code.
@@ -223,7 +234,7 @@
     Just stms' | ValueExp x <- max_block_size, oneIsh x -> do
       let body' = body { bodyStms = stmsFromList stms' }
           body'' = allThreads constants $
-                   ImpGen.compileLoopBody (map paramName acc_params) body'
+                   compileLoopBody acc_params body'
       block_size <-- 1
 
       -- Check if loop is candidate for unrolling.
@@ -244,7 +255,7 @@
     _ -> do
       dPrim_ block_offset int32
       let body' = streaming constants block_size maxchunk $
-                  ImpGen.compileBody' acc_params body
+                  compileBody' acc_params body
 
       block_offset <-- 0
 
@@ -270,23 +281,23 @@
              set_block_size >> body' >> increase_offset
 
   forM_ (zip final acc_params) $ \(pe, p) ->
-    ImpGen.copyDWIM (patElemName pe) [] (Var $ paramName p) []
+    copyDWIM (patElemName pe) [] (Var $ paramName p) []
 
   where isSimpleThreadInSpace (Let _ _ Op{}) = Nothing
         isSimpleThreadInSpace bnd = Just bnd
 
 compileKernelExp _ _ (GroupGenReduce w arrs op bucket values locks) = do
   -- Check if bucket is in-bounds
-  bucket' <- mapM ImpGen.compileSubExp bucket
-  w' <- mapM ImpGen.compileSubExp w
-  num_locks <- ImpGen.compileSubExpOfType int32 . arraySize 0 <$> lookupType locks
+  bucket' <- mapM toExp bucket
+  w' <- mapM toExp w
+  num_locks <- toExp' int32 . arraySize 0 <$> lookupType locks
   let locking = Locking locks 0 1 0 $ (`rem` num_locks) . sum
       values_params = takeLast (length values) $ lambdaParams op
 
   sWhen (indexInBounds bucket' w') $ do
     forM_ (zip values_params values) $ \(p, v) ->
-      ImpGen.copyDWIM (paramName p) [] v []
-    atomicUpdate arrs bucket' op locking
+      copyDWIM (paramName p) [] v []
+    atomicUpdate DefaultSpace arrs bucket' op locking
   where indexInBounds inds bounds =
           foldl1 (.&&.) $ zipWith checkBound inds bounds
           where checkBound ind bound = 0 .<=. ind .&&. ind .<. bound
@@ -297,10 +308,10 @@
 
 streaming :: KernelConstants -> VName -> SubExp -> InKernelGen () -> InKernelGen ()
 streaming constants chunksize bound m = do
-  bound' <- ImpGen.subExpToDimSize bound
+  bound' <- subExpToDimSize bound
   let constants' =
         constants { kernelStreamed = (chunksize, bound') : kernelStreamed constants }
-  ImpGen.emit =<< ImpGen.subImpM_ (inKernelOperations constants') m
+  emit =<< subImpM_ (inKernelOperations constants') m
 
 -- | Locking strategy used for an atomic update.
 data Locking =
@@ -321,15 +332,15 @@
 -- | A function for generating code for an atomic update.  Assumes
 -- that the bucket is in-bounds.
 type AtomicUpdate lore =
-  [VName] -> [Imp.Exp] -> ImpGen.ImpM lore Imp.KernelOp ()
+  Space -> [VName] -> [Imp.Exp] -> ImpM lore Imp.KernelOp ()
 
 atomicUpdate :: ExplicitMemorish lore =>
-                [VName] -> [Imp.Exp] -> Lambda lore -> Locking
-             -> ImpGen.ImpM lore Imp.KernelOp ()
-atomicUpdate arrs bucket lam locking =
+                Space -> [VName] -> [Imp.Exp] -> Lambda lore -> Locking
+             -> ImpM lore Imp.KernelOp ()
+atomicUpdate space arrs bucket lam locking =
   case atomicUpdateLocking lam of
-    Left f -> f arrs bucket
-    Right f -> f locking arrs bucket
+    Left f -> f space arrs bucket
+    Right f -> f locking space arrs bucket
 
 -- | 'atomicUpdate', but where it is explicitly visible whether a
 -- locking strategy is necessary.
@@ -339,7 +350,7 @@
 
 atomicUpdateLocking lam
   | Just ops_and_ts <- splitOp lam,
-    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = Left $ \arrs bucket ->
+    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = Left $ \space arrs bucket ->
   -- If the operator is a vectorised binary operator on 32-bit values,
   -- we can use a particularly efficient implementation. If the
   -- operator has an atomic implementation we use that, otherwise it
@@ -350,16 +361,15 @@
   -- Common variables.
   old <- dPrim "old" t
 
-  (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray a bucket
+  (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket
 
-  case opHasAtomicSupport old arr' bucket_offset op of
+  case opHasAtomicSupport space old arr' bucket_offset op of
     Just f -> sOp $ f $ Imp.var y t
-
-    Nothing -> atomicUpdateCAS t a old bucket x $
+    Nothing -> atomicUpdateCAS space t a old bucket x $
       x <-- Imp.BinOpExp op (Imp.var x t) (Imp.var y t)
 
-  where opHasAtomicSupport old arr' bucket' bop = do
-          let atomic f = Imp.Atomic . f old arr' bucket'
+  where opHasAtomicSupport space old arr' bucket' bop = do
+          let atomic f = Imp.Atomic space . f old arr' bucket'
           atomic <$> Imp.atomicBinOp bop
 
 -- If the operator functions purely on single 32-bit values, we can
@@ -368,28 +378,28 @@
 atomicUpdateLocking op
   | [Prim t] <- lambdaReturnType op,
     [xp, _] <- lambdaParams op,
-    primBitSize t == 32 = Left $ \[arr] bucket -> do
+    primBitSize t == 32 = Left $ \space [arr] bucket -> do
       old <- dPrim "old" t
-      atomicUpdateCAS t arr old bucket (paramName xp) $
-        ImpGen.compileBody' [xp] $ lambdaBody op
+      atomicUpdateCAS space t arr old bucket (paramName xp) $
+        compileBody' [xp] $ lambdaBody op
 
-atomicUpdateLocking op = Right $ \locking arrs bucket -> do
+atomicUpdateLocking op = Right $ \locking space arrs bucket -> do
   old <- dPrim "old" int32
   continue <- dPrimV "continue" true
 
   -- Correctly index into locks.
   (locks', _locks_space, locks_offset) <-
-    ImpGen.fullyIndexArray (lockingArray locking) [lockingMapping locking bucket]
+    fullyIndexArray (lockingArray locking) [lockingMapping locking bucket]
 
   -- Critical section
   let try_acquire_lock =
-        sOp $ Imp.Atomic $
+        sOp $ Imp.Atomic space $
         Imp.AtomicCmpXchg old locks' locks_offset (lockingIsUnlocked locking) (lockingToLock locking)
       lock_acquired = Imp.var old int32 .==. lockingIsUnlocked locking
       -- Even the releasing is done with an atomic rather than a
       -- simple write, for memory coherency reasons.
       release_lock =
-        sOp $ Imp.Atomic $
+        sOp $ Imp.Atomic space $
         Imp.AtomicCmpXchg old locks' locks_offset (lockingToLock locking) (lockingToUnlock locking)
       break_loop = continue <-- false
 
@@ -405,39 +415,43 @@
   -- synchronisation without atomics (naughty!).
   let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op
       bind_acc_params =
-        ImpGen.everythingVolatile $
-        ImpGen.sComment "bind lhs" $
+        everythingVolatile $
+        sComment "bind lhs" $
         forM_ (zip acc_params arrs) $ \(acc_p, arr) ->
-        ImpGen.copyDWIM (paramName acc_p) [] (Var arr) bucket
+        copyDWIM (paramName acc_p) [] (Var arr) bucket
 
-  let op_body = ImpGen.sComment "execute operation" $
-                ImpGen.compileBody' acc_params $ lambdaBody op
+  let op_body = sComment "execute operation" $
+                compileBody' acc_params $ lambdaBody op
 
       do_gen_reduce =
-        ImpGen.everythingVolatile $
-        ImpGen.sComment "update global result" $
+        everythingVolatile $
+        sComment "update global result" $
         zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params
 
+      fence = case space of Space "local" -> sOp Imp.MemFenceLocal
+                            _             -> sOp Imp.MemFenceGlobal
+
+
   -- While-loop: Try to insert your value
   sWhile (Imp.var continue Bool) $ do
     try_acquire_lock
     sWhen lock_acquired $ do
-      ImpGen.dLParams acc_params
+      dLParams acc_params
       bind_acc_params
       op_body
       do_gen_reduce
-      sOp Imp.MemFence
+      fence
       release_lock
       break_loop
-    sOp Imp.MemFence
-  where writeArray bucket arr val = ImpGen.copyDWIM arr bucket val []
+    fence
+  where writeArray bucket arr val = copyDWIM arr bucket val []
 
-atomicUpdateCAS :: PrimType
+atomicUpdateCAS :: Space -> PrimType
                 -> VName -> VName
                 -> [Imp.Exp] -> VName
-                -> ImpGen.ImpM lore Imp.KernelOp ()
-                -> ImpGen.ImpM lore Imp.KernelOp ()
-atomicUpdateCAS t arr old bucket x do_op = do
+                -> ImpM lore Imp.KernelOp ()
+                -> ImpM lore Imp.KernelOp ()
+atomicUpdateCAS space t arr old bucket x do_op = do
   -- Code generation target:
   --
   -- old = d_his[idx];
@@ -448,9 +462,9 @@
   -- } while(assumed != old);
   assumed <- dPrim "assumed" t
   run_loop <- dPrimV "run_loop" 1
-  ImpGen.copyDWIM old [] (Var arr) bucket
+  copyDWIM old [] (Var arr) bucket
 
-  (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray arr bucket
+  (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket
 
   -- While-loop: Try to insert your value
   let (toBits, fromBits) =
@@ -462,7 +476,7 @@
     x <-- Imp.var assumed t
     do_op
     old_bits <- dPrim "old_bits" int32
-    sOp $ Imp.Atomic $
+    sOp $ Imp.Atomic space $
       Imp.AtomicCmpXchg old_bits arr' bucket_offset
       (toBits (Imp.var assumed t)) (toBits (Imp.var x t))
     old <-- fromBits (Imp.var old_bits int32)
@@ -488,62 +502,50 @@
 
 computeKernelUses :: FreeIn a =>
                      a -> [VName]
-                  -> CallKernelGen ([Imp.KernelUse], [Imp.LocalMemoryUse])
+                  -> CallKernelGen [Imp.KernelUse]
 computeKernelUses kernel_body bound_in_kernel = do
   let actually_free = freeIn kernel_body `S.difference` S.fromList bound_in_kernel
-
   -- Compute the variables that we need to pass to the kernel.
-  reads_from <- readsFromSet actually_free
-
-  -- Are we using any local memory?
-  local_memory <- computeLocalMemoryUse actually_free
-  return (nub reads_from, nub local_memory)
+  nub <$> readsFromSet actually_free
 
 readsFromSet :: Names -> CallKernelGen [Imp.KernelUse]
 readsFromSet free =
   fmap catMaybes $
   forM (S.toList free) $ \var -> do
     t <- lookupType var
+    vtable <- getVTable
     case t of
       Array {} -> return Nothing
-      Mem _ (Space "local") -> return Nothing
+      Mem (Space "local") -> return Nothing
       Mem {} -> return $ Just $ Imp.MemoryUse var
       Prim bt ->
-        isConstExp var >>= \case
+        isConstExp vtable (Imp.var var bt) >>= \case
           Just ce -> return $ Just $ Imp.ConstUse var ce
           Nothing | bt == Cert -> return Nothing
                   | otherwise  -> return $ Just $ Imp.ScalarUse var bt
 
-computeLocalMemoryUse :: Names -> CallKernelGen [Imp.LocalMemoryUse]
-computeLocalMemoryUse free =
-  fmap catMaybes $
-  forM (S.toList free) $ \var -> do
-    t <- lookupType var
-    case t of
-      Mem memsize (Space "local") -> do
-        memsize' <- localMemSize =<< ImpGen.subExpToDimSize memsize
-        return $ Just (var, memsize')
-      _ -> return Nothing
-
-localMemSize :: Imp.MemSize -> CallKernelGen (Either Imp.MemSize Imp.KernelConstExp)
-localMemSize (Imp.ConstSize x) =
-  return $ Right $ ValueExp $ IntValue $ Int64Value x
-localMemSize (Imp.VarSize v) = isConstExp v >>= \case
-  Just e | isStaticExp e -> return $ Right e
-  _ -> return $ Left $ Imp.VarSize v
+localMemSize :: VTable ExplicitMemory -> Imp.Count Imp.Bytes
+             -> ImpM lore op (Either (Imp.Count Imp.Bytes) Imp.KernelConstExp)
+localMemSize vtable e = isConstExp vtable (Imp.innerExp e) >>= \case
+  Just e' | isStaticExp e' -> return $ Right e'
+  _ -> return $ Left e
 
-isConstExp :: VName -> CallKernelGen (Maybe Imp.KernelConstExp)
-isConstExp v = do
-  vtable <- ImpGen.getVTable
-  fname <- asks ImpGen.envFunction
-  let lookupConstExp name = constExp =<< hasExp =<< M.lookup name vtable
+isConstExp :: VTable ExplicitMemory -> Imp.Exp
+           -> ImpM lore op (Maybe Imp.KernelConstExp)
+isConstExp vtable size = do
+  fname <- asks envFunction
+  let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name
+      onLeaf (Imp.SizeOf pt) _ = Just $ primByteSize pt
+      onLeaf Imp.Index{} _ = Nothing
+      lookupConstExp name =
+        constExp =<< hasExp =<< M.lookup name vtable
       constExp (Op (Inner (GetSize key _))) =
         Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32
       constExp e = primExpFromExp lookupConstExp e
-  return $ lookupConstExp v
-  where hasExp (ImpGen.ArrayVar e _) = e
-        hasExp (ImpGen.ScalarVar e _) = e
-        hasExp (ImpGen.MemVar e _) = e
+  return $ replaceInPrimExpM onLeaf size
+  where hasExp (ArrayVar e _) = e
+        hasExp (ScalarVar e _) = e
+        hasExp (MemVar e _) = e
 
 -- | Only some constant expressions qualify as *static* expressions,
 -- which we can use for static memory allocation.  This is a bit of a
@@ -552,6 +554,8 @@
 isStaticExp :: Imp.KernelConstExp -> Bool
 isStaticExp LeafExp{} = True
 isStaticExp ValueExp{} = True
+isStaticExp (ConvOpExp ZExt{} x) = isStaticExp x
+isStaticExp (ConvOpExp SExt{} x) = isStaticExp x
 isStaticExp (BinOpExp Add{} x y) = isStaticExp x && isStaticExp y
 isStaticExp (BinOpExp Sub{} x y) = isStaticExp x && isStaticExp y
 isStaticExp (BinOpExp Mul{} x y) = isStaticExp x && isStaticExp y
@@ -562,9 +566,9 @@
                        -> Imp.Count Imp.Elements
                        -> Imp.Count Imp.Elements
                        -> VName
-                       -> ImpGen.ImpM lore op ()
+                       -> ImpM lore op ()
 computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do
-  stride' <- ImpGen.compileSubExp stride
+  stride' <- toExp stride
   chunk_var <--
     Imp.BinOpExp (SMin Int32)
     (Imp.innerExp elements_per_thread)
@@ -590,36 +594,37 @@
           Imp.innerExp num_elements .<.
           (thread_index + 1) * Imp.innerExp elements_per_thread
 
-kernelInitialisationSetSpace :: KernelSpace -> InKernelGen ()
-                             -> ImpGen.ImpM lore op (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())
-kernelInitialisationSetSpace space set_space = do
-  group_size' <- ImpGen.compileSubExp $ spaceGroupSize space
-  num_threads' <- ImpGen.compileSubExp $ spaceNumThreads space
-  num_groups <- ImpGen.compileSubExp $ spaceNumGroups space
-
-  let global_tid = spaceGlobalId space
-      local_tid = spaceLocalId space
-      group_id = spaceGroupId space
+kernelInitialisationSimple :: Imp.Exp -> Imp.Exp
+                           -> Maybe (VName, VName, VName)
+                           -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())
+kernelInitialisationSimple num_groups group_size names = do
+  (global_tid, local_tid, group_id) <-
+    case names of Nothing ->
+                    (,,)
+                    <$> newVName "global_tid"
+                    <*> newVName "local_tid"
+                    <*> newVName "group_id"
+                  Just (global_tid, local_tid, group_id) ->
+                    return (global_tid, local_tid, group_id)
   wave_size <- newVName "wave_size"
   inner_group_size <- newVName "group_size"
-
-  let (space_is, space_dims) = unzip $ spaceDimensions space
-  space_dims' <- mapM ImpGen.compileSubExp space_dims
+  vtable <- getVTable
   let constants =
-        KernelConstants
+        KernelConstants vtable
         (Imp.var global_tid int32)
         (Imp.var local_tid int32)
         (Imp.var group_id int32)
         global_tid local_tid group_id
-        group_size' num_groups num_threads'
-        (Imp.var wave_size int32) (zip space_is space_dims')
-        (if null (spaceDimensions space)
-         then true else isActive (spaceDimensions space)) mempty
+        group_size num_groups (group_size*num_groups)
+        (Imp.var wave_size int32) []
+        true mempty
 
   let set_constants = do
-        dPrim_ wave_size int32
+        dPrim_ global_tid int32
+        dPrim_ local_tid int32
         dPrim_ inner_group_size int32
-        ImpGen.dScope Nothing (scopeOfKernelSpace space)
+        dPrim_ wave_size int32
+        dPrim_ group_id int32
 
         sOp (Imp.GetGlobalId global_tid 0)
         sOp (Imp.GetLocalId local_tid 0)
@@ -627,12 +632,42 @@
         sOp (Imp.GetLockstepWidth wave_size)
         sOp (Imp.GetGroupId group_id 0)
 
+  return (constants, set_constants)
+
+kernelInitialisationSetSpace :: KernelSpace -> InKernelGen ()
+                             -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())
+kernelInitialisationSetSpace space set_space = do
+  group_size <- toExp $ spaceGroupSize space
+  num_groups <- toExp $ spaceNumGroups space
+
+  (constants, set_constants) <-
+    kernelInitialisationSimple num_groups group_size $
+    Just (spaceGlobalId space, spaceLocalId space, spaceGroupId space)
+
+  let set_constants' = do
+        set_constants
+        case spaceStructure space of
+          FlatThreadSpace is_and_dims ->
+            mapM_ ((`dPrim_` int32) . fst) is_and_dims
+          NestedThreadSpace is_and_dims -> do
+            let (gtids, _, ltids, _) = unzip4 is_and_dims
+            mapM_ (`dPrim_` int32) $ gtids ++ ltids
         set_space
 
-  return (constants, set_constants)
+  let (space_is, space_dims) = unzip $ spaceDimensions space
+  space_dims' <- mapM toExp space_dims
 
+  return (constants { kernelThreadActive =
+                        if null $ spaceDimensions space
+                        then true
+                        else isActive $ spaceDimensions space
+                    , kernelDimensions =
+                        zip space_is space_dims'
+                    },
+          set_constants')
+
 kernelInitialisation :: KernelSpace
-                     -> ImpGen.ImpM lore op (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())
+                     -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())
 kernelInitialisation space =
   kernelInitialisationSetSpace space $
   setSpaceIndices (Imp.var (spaceGlobalId space) int32) space
@@ -644,23 +679,22 @@
       flatSpaceWith gtid is_and_dims
     NestedThreadSpace is_and_dims -> do
       let (gtids, gdims, ltids, ldims) = unzip4 is_and_dims
-      gdims' <- mapM ImpGen.compileSubExp gdims
-      ldims' <- mapM ImpGen.compileSubExp ldims
+      gdims' <- mapM toExp gdims
+      ldims' <- mapM toExp ldims
       let (gtid_es, ltid_es) = unzip $ unflattenNestedIndex gdims' ldims' gtid
       zipWithM_ (<--) gtids gtid_es
       zipWithM_ (<--) ltids ltid_es
   where flatSpaceWith base is_and_dims = do
           let (is, dims) = unzip is_and_dims
-          dims' <- mapM ImpGen.compileSubExp dims
-          let index_expressions = unflattenIndex dims' base
-          zipWithM_ (<--) is index_expressions
+          dims' <- mapM toExp dims
+          zipWithM_ (<--) is $ unflattenIndex dims' base
 
 isActive :: [(VName, SubExp)] -> Imp.Exp
 isActive limit = case actives of
                     [] -> Imp.ValueExp $ BoolValue True
                     x:xs -> foldl (.&&.) x xs
   where (is, ws) = unzip limit
-        actives = zipWith active is $ map (ImpGen.compileSubExpOfType Bool) ws
+        actives = zipWith active is $ map (toExp' Bool) ws
         active i = (Imp.var i int32 .<.)
 
 unflattenNestedIndex :: IntegralExp num => [num] -> [num] -> num -> [(num,num)]
@@ -683,16 +717,16 @@
 -- kernel).
 makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a
 makeAllMemoryGlobal =
-  local (\env -> env { ImpGen.envDefaultSpace = Imp.Space "global" }) .
-  ImpGen.localVTable (M.map globalMemory)
-  where globalMemory (ImpGen.MemVar _ entry)
-          | ImpGen.entryMemSpace entry /= Space "local" =
-              ImpGen.MemVar Nothing entry { ImpGen.entryMemSpace = Imp.Space "global" }
+  local (\env -> env { envDefaultSpace = Imp.Space "global" }) .
+  localVTable (M.map globalMemory)
+  where globalMemory (MemVar _ entry)
+          | entryMemSpace entry /= Space "local" =
+              MemVar Nothing entry { entryMemSpace = Imp.Space "global" }
         globalMemory entry =
           entry
 
 allThreads :: KernelConstants -> InKernelGen () -> InKernelGen ()
-allThreads constants = ImpGen.emit <=< ImpGen.subImpM_ (inKernelOperations constants')
+allThreads constants = emit <=< subImpM_ (inKernelOperations constants')
   where constants' =
           constants { kernelThreadActive = Imp.ValueExp (BoolValue True) }
 
@@ -700,34 +734,32 @@
 
 writeParamToLocalMemory :: Typed (MemBound u) =>
                            Imp.Exp -> (VName, t) -> Param (MemBound u)
-                        -> ImpGen.ImpM lore op ()
+                        -> ImpM lore op ()
 writeParamToLocalMemory i (mem, _) param
   | Prim t <- paramType param =
-      ImpGen.emit $
-      Imp.Write mem (bytes i') bt (Space "local") Imp.Volatile $
+      emit $
+      Imp.Write mem (elements i) bt (Space "local") Imp.Volatile $
       Imp.var (paramName param) t
   | otherwise =
       return ()
-  where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32
-        bt = elemType $ paramType param
+  where bt = elemType $ paramType param
 
 readParamFromLocalMemory :: Typed (MemBound u) =>
-                            VName -> Imp.Exp -> Param (MemBound u) -> (VName, t)
-                         -> ImpGen.ImpM lore op ()
-readParamFromLocalMemory index i param (l_mem, _)
+                            Imp.Exp -> Param (MemBound u) -> (VName, t)
+                         -> ImpM lore op ()
+readParamFromLocalMemory i param (l_mem, _)
   | Prim _ <- paramType param =
       paramName param <--
-      Imp.index l_mem (bytes i') bt (Space "local") Imp.Volatile
-  | otherwise = index <-- i
-  where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32
-        bt = elemType $ paramType param
+      Imp.index l_mem (elements i) bt (Space "local") Imp.Volatile
+  | otherwise = return ()
+  where bt = elemType $ paramType param
 
 groupReduce :: ExplicitMemorish lore =>
                KernelConstants
             -> Imp.Exp
             -> Lambda lore
             -> [VName]
-            -> ImpGen.ImpM lore Imp.KernelOp ()
+            -> ImpM lore Imp.KernelOp ()
 groupReduce constants w lam arrs = do
   offset <- dPrim "offset" int32
   groupReduceWithOffset constants offset w lam arrs
@@ -738,26 +770,26 @@
                       -> Imp.Exp
                       -> Lambda lore
                       -> [VName]
-                      -> ImpGen.ImpM lore Imp.KernelOp ()
+                      -> ImpM lore Imp.KernelOp ()
 groupReduceWithOffset constants offset w lam arrs = do
   let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam
 
   skip_waves <- dPrim "skip_waves" int32
-  ImpGen.dLParams $ lambdaParams lam
+  dLParams $ lambdaParams lam
 
   offset <-- 0
 
-  ImpGen.comment "participating threads read initial accumulator" $
+  comment "participating threads read initial accumulator" $
     sWhen (local_tid .<. w) $
     zipWithM_ readReduceArgument reduce_acc_params arrs
 
-  let do_reduce = do ImpGen.comment "read array element" $
+  let do_reduce = do comment "read array element" $
                        zipWithM_ readReduceArgument reduce_arr_params arrs
-                     ImpGen.comment "apply reduction operation" $
-                       ImpGen.compileBody' reduce_acc_params $ lambdaBody lam
-                     ImpGen.comment "write result of operation" $
+                     comment "apply reduction operation" $
+                       compileBody' reduce_acc_params $ lambdaBody lam
+                     comment "write result of operation" $
                        zipWithM_ writeReduceOpResult reduce_acc_params arrs
-      in_wave_reduce = ImpGen.everythingVolatile do_reduce
+      in_wave_reduce = everythingVolatile do_reduce
 
       wave_size = kernelWaveSize constants
       group_size = kernelGroupSize constants
@@ -805,15 +837,15 @@
 
         readReduceArgument param arr
           | Prim _ <- paramType param = do
-              let i = local_tid + ImpGen.varIndex offset
-              ImpGen.copyDWIM (paramName param) [] (Var arr) [i]
+              let i = local_tid + Imp.vi32 offset
+              copyDWIM (paramName param) [] (Var arr) [i]
           | otherwise = do
-              let i = global_tid + ImpGen.varIndex offset
-              ImpGen.copyDWIM (paramName param) [] (Var arr) [i]
+              let i = global_tid + Imp.vi32 offset
+              copyDWIM (paramName param) [] (Var arr) [i]
 
         writeReduceOpResult param arr
           | Prim _ <- paramType param =
-              ImpGen.copyDWIM arr [local_tid] (Var $ paramName param) []
+              copyDWIM arr [local_tid] (Var $ paramName param) []
           | otherwise =
               return ()
 
@@ -822,7 +854,7 @@
           -> Imp.Exp
           -> Lambda InKernel
           -> [VName]
-          -> ImpGen.ImpM InKernel Imp.KernelOp ()
+          -> ImpM InKernel Imp.KernelOp ()
 groupScan constants seg_flag w lam arrs = do
   when (any (not . primType . paramType) $ lambdaParams lam) $
     compilerLimitationS "Cannot compile parallel scans with array element type."
@@ -830,16 +862,13 @@
   renamed_lam <- renameLambda lam
 
   acc_local_mem <- flip zip (repeat ()) <$>
-                   mapM (fmap (ImpGen.memLocationName . ImpGen.entryArrayLocation) .
-                         ImpGen.lookupArray) arrs
+                   mapM (fmap (memLocationName . entryArrayLocation) .
+                         lookupArray) arrs
 
   let ltid = kernelLocalThreadId constants
-      (lam_i, other_index_param, actual_params) =
-        partitionChunkedKernelLambdaParameters $ lambdaParams lam
-      (x_params, y_params) = splitAt (length arrs) actual_params
+      (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam
 
-  ImpGen.dLParams (lambdaParams lam++lambdaParams renamed_lam)
-  lam_i <-- ltid
+  dLParams (lambdaParams lam++lambdaParams renamed_lam)
 
   -- The scan works by splitting the group into blocks, which are
   -- scanned separately.  Typically, these blocks are smaller than
@@ -873,23 +902,22 @@
         flag_true <- seg_flag
         Just $ \from to ->
           flag_true (from*block_size+block_size-1) (to*block_size+block_size-1)
-  ImpGen.comment
+  comment
     "scan the first block, after which offset 'i' contains carry-in for warp 'i+1'" $
     doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam
 
   sOp Imp.LocalBarrier
 
   let read_carry_in =
-        zipWithM_ (readParamFromLocalMemory
-                   (paramName other_index_param) (block_id - 1))
+        zipWithM_ (readParamFromLocalMemory (block_id - 1))
         x_params acc_local_mem
 
   let op_to_y
         | Nothing <- seg_flag =
-            ImpGen.compileBody' y_params $ lambdaBody lam
+            compileBody' y_params $ lambdaBody lam
         | Just flag_true <- seg_flag =
             sUnless (flag_true (block_id*block_size-1) ltid) $
-              ImpGen.compileBody' y_params $ lambdaBody lam
+              compileBody' y_params $ lambdaBody lam
       write_final_result =
         zipWithM_ (writeParamToLocalMemory ltid) acc_local_mem y_params
 
@@ -914,30 +942,27 @@
             -> [(VName, t)]
             -> Lambda InKernel
             -> InKernelGen ()
-inBlockScan seg_flag lockstep_width block_size active ltid acc_local_mem scan_lam = ImpGen.everythingVolatile $ do
+inBlockScan seg_flag lockstep_width block_size active ltid acc_local_mem scan_lam = everythingVolatile $ do
   skip_threads <- dPrim "skip_threads" int32
   let in_block_thread_active =
         Imp.var skip_threads int32 .<=. in_block_id
-      (scan_lam_i, other_index_param, actual_params) =
-        partitionChunkedKernelLambdaParameters $ lambdaParams scan_lam
+      actual_params = lambdaParams scan_lam
       (x_params, y_params) =
         splitAt (length actual_params `div` 2) actual_params
       read_operands =
-        zipWithM_ (readParamFromLocalMemory (paramName other_index_param) $
-                   ltid - Imp.var skip_threads int32)
+        zipWithM_ (readParamFromLocalMemory $ ltid - Imp.var skip_threads int32)
         x_params acc_local_mem
 
   -- Set initial y values
   sWhen active $
-    zipWithM_ (readParamFromLocalMemory scan_lam_i ltid)
-    y_params acc_local_mem
+    zipWithM_ (readParamFromLocalMemory ltid) y_params acc_local_mem
 
   let op_to_y
         | Nothing <- seg_flag =
-            ImpGen.compileBody' y_params $ lambdaBody scan_lam
+            compileBody' y_params $ lambdaBody scan_lam
         | Just flag_true <- seg_flag =
             sUnless (flag_true (ltid-Imp.var skip_threads int32) ltid) $
-              ImpGen.compileBody' y_params $ lambdaBody scan_lam
+              compileBody' y_params $ lambdaBody scan_lam
       write_operation_result =
         zipWithM_ (writeParamToLocalMemory ltid) acc_local_mem y_params
       maybeLocalBarrier = sWhen (lockstep_width .<=. Imp.var skip_threads int32) $
@@ -962,14 +987,14 @@
   where block_id = ltid `quot` block_size
         in_block_id = ltid - block_id * block_size
 
-compileKernelStms :: KernelConstants -> [Stm InKernel]
+compileKernelStms :: KernelConstants -> Stms InKernel
                   -> InKernelGen a
                   -> InKernelGen a
 compileKernelStms constants ungrouped_bnds m =
-  compileGroupedKernelStms' $ groupStmsByGuard constants ungrouped_bnds
+  compileGroupedKernelStms' $ groupStmsByGuard constants $ stmsToList ungrouped_bnds
   where compileGroupedKernelStms' [] = m
         compileGroupedKernelStms' ((g, bnds):rest_bnds) = do
-          ImpGen.dScopes (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds)
+          dScopes (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds)
           protect g $ mapM_ compileKernelStm bnds
           compileGroupedKernelStms' rest_bnds
 
@@ -980,7 +1005,7 @@
         protect (Just g) body_m =
           sWhen g $ allThreads constants body_m
 
-        compileKernelStm (Let pat _ e) = ImpGen.compileExp pat e
+        compileKernelStm (Let pat _ e) = compileExp pat e
 
 groupStmsByGuard :: KernelConstants
                      -> [Stm InKernel]
@@ -1000,7 +1025,7 @@
 computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp)
 computeMapKernelGroups kernel_size = do
   group_size <- dPrim "group_size" int32
-  fname <- asks ImpGen.envFunction
+  fname <- asks envFunction
   let group_size_var = Imp.var group_size int32
       group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size
   sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup
@@ -1008,7 +1033,7 @@
   return (Imp.var group_size int32, Imp.var num_groups int32)
 
 simpleKernelConstants :: Imp.Exp -> String
-                      -> CallKernelGen (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ())
+                      -> CallKernelGen (KernelConstants, ImpM InKernel Imp.KernelOp ())
 simpleKernelConstants kernel_size desc = do
   thread_gtid <- newVName $ desc ++ "_gtid"
   thread_ltid <- newVName $ desc ++ "_ltid"
@@ -1022,7 +1047,9 @@
         sOp (Imp.GetLocalId thread_ltid 0)
         sOp (Imp.GetGroupId group_id 0)
 
-  return (KernelConstants
+
+  vtable <- getVTable
+  return (KernelConstants vtable
           (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32)
           thread_gtid thread_ltid group_id
           group_size num_groups (group_size*num_groups) 0
@@ -1030,20 +1057,43 @@
 
           set_constants)
 
-sKernel :: KernelConstants -> String -> ImpGen.ImpM InKernel Imp.KernelOp a -> CallKernelGen ()
+-- | For many kernels, we may not have enough physical groups to cover
+-- the logical iteration space.  Some groups thus have to perform
+-- double duty; we put an outer loop to accomplish this.  The
+-- advantage over just launching a bazillion threads is that the cost
+-- of memory expansion should be proportional to the number of
+-- *physical* threads (hardware parallelism), not the amount of
+-- application parallelism.
+virtualiseGroups :: KernelConstants
+                 -> Imp.Exp
+                 -> (VName -> InKernelGen ())
+                 -> InKernelGen ()
+virtualiseGroups constants required_groups m
+  | kernelNumGroups constants == required_groups =
+      m $ kernelGroupIdVar constants
+  | otherwise = do
+  phys_group_id <- dPrim "phys_group_id" int32
+  sOp $ Imp.GetGroupId phys_group_id 0
+  let iterations = (required_groups - Imp.vi32 phys_group_id) `quotRoundingUp`
+                   kernelNumGroups constants
+  i <- newVName "i"
+  sFor i Int32 iterations $
+    m =<< dPrimV "virt_group_id" (Imp.vi32 phys_group_id + Imp.vi32 i * kernelNumGroups constants)
+
+sKernel :: KernelConstants -> String -> ImpM InKernel Imp.KernelOp a -> CallKernelGen ()
 sKernel constants name m = do
-  body <- makeAllMemoryGlobal $
-          ImpGen.subImpM_ (inKernelOperations constants) m
-  (uses, local_memory) <- computeKernelUses body mempty
-  ImpGen.emit $ Imp.Op $ Imp.CallKernel Imp.Kernel
+  body <- makeAllMemoryGlobal $ subImpM_ (inKernelOperations constants) m
+  uses <- computeKernelUses body mempty
+
+  emit $ Imp.Op $ Imp.CallKernel Imp.Kernel
     { Imp.kernelBody = body
-    , Imp.kernelLocalMemory = local_memory
     , Imp.kernelUses = uses
     , Imp.kernelNumGroups = [kernelNumGroups constants]
     , Imp.kernelGroupSize = [kernelGroupSize constants]
     , Imp.kernelName =
-        nameFromString $ name ++ "_" ++ show (baseTag $ kernelGlobalThreadIdVar constants)
+        nameFromString $ name ++ "_" ++ show tag
     }
+  where tag = baseTag $ kernelGlobalThreadIdVar constants
 
 -- | Perform a Replicate with a kernel.
 sReplicate :: VName -> Shape -> SubExp
@@ -1051,7 +1101,7 @@
 sReplicate arr (Shape ds) se = do
   t <- subExpType se
 
-  dims <- mapM ImpGen.compileSubExp $ ds ++ arrayDims t
+  dims <- mapM toExp $ ds ++ arrayDims t
   (constants, set_constants) <-
     simpleKernelConstants (product dims) "replicate"
 
@@ -1060,39 +1110,38 @@
   sKernel constants "replicate" $ do
     set_constants
     sWhen (kernelThreadActive constants) $
-      ImpGen.copyDWIM arr is' se $ drop (length ds) is'
+      copyDWIM arr is' se $ drop (length ds) is'
 
 -- | Perform an Iota with a kernel.
 sIota :: VName -> Imp.Exp -> Imp.Exp -> Imp.Exp -> IntType
       -> CallKernelGen ()
 sIota arr n x s et = do
-  destloc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray arr
+  destloc <- entryArrayLocation <$> lookupArray arr
   (constants, set_constants) <- simpleKernelConstants n "iota"
 
   sKernel constants "iota" $ do
     set_constants
     let gtid = kernelGlobalThreadId constants
     sWhen (kernelThreadActive constants) $ do
-      (destmem, destspace, destidx) <-
-        ImpGen.fullyIndexArray' destloc [gtid] (IntType et)
+      (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid]
 
-      ImpGen.emit $
+      emit $
         Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $
         Imp.ConvOpExp (SExt Int32 et) gtid * s + x
 
 sCopy :: PrimType
-      -> ImpGen.MemLocation
-      -> ImpGen.MemLocation
+      -> MemLocation
+      -> MemLocation
       -> Imp.Count Imp.Elements
       -> CallKernelGen ()
 sCopy bt
-  destloc@(ImpGen.MemLocation destmem _ _)
-  srcloc@(ImpGen.MemLocation srcmem srcshape _)
+  destloc@(MemLocation destmem _ _)
+  srcloc@(MemLocation srcmem srcshape _)
   n = do
   -- Note that the shape of the destination and the source are
   -- necessarily the same.
   let shape = map Imp.sizeToExp srcshape
-      shape_se = map (Imp.innerExp . ImpGen.dimSizeToExp) srcshape
+      shape_se = map (Imp.innerExp . dimSizeToExp) srcshape
       kernel_size = Imp.innerExp n * product (drop 1 shape)
 
   (constants, set_constants) <- simpleKernelConstants kernel_size "copy"
@@ -1104,9 +1153,87 @@
         dest_is = unflattenIndex shape_se gtid
         src_is = dest_is
 
-    (_, destspace, destidx) <- ImpGen.fullyIndexArray' destloc dest_is bt
-    (_, srcspace, srcidx) <- ImpGen.fullyIndexArray' srcloc src_is bt
+    (_, destspace, destidx) <- fullyIndexArray' destloc dest_is
+    (_, srcspace, srcidx) <- fullyIndexArray' srcloc src_is
 
-    sWhen (gtid .<. kernel_size) $ ImpGen.emit $
+    sWhen (gtid .<. kernel_size) $ emit $
       Imp.Write destmem destidx bt destspace Imp.Nonvolatile $
       Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile
+
+
+compileKernelResult :: KernelConstants -> PatElem InKernel -> KernelResult
+                    -> InKernelGen ()
+
+compileKernelResult constants pe (GroupsReturn what) = do
+  i <- newVName "i"
+
+  in_local_memory <- arrayInLocalMemory what
+  let me = kernelLocalThreadId constants
+
+  if not in_local_memory then do
+    who' <- toExp $ intConst Int32 0
+    sWhen (me .==. who') $
+      copyDWIM (patElemName pe) [kernelGroupId constants] what []
+    else do
+      -- If the result of the group is an array in local memory, we
+      -- store it by collective copying among all the threads of the
+      -- group.  TODO: also do this if the array is in global memory
+      -- (but this is a bit more tricky, synchronisation-wise).
+      --
+      -- We do the reads/writes multidimensionally, but the loop is
+      -- single-dimensional.
+      ws <- mapM toExp . arrayDims =<< subExpType what
+      -- Compute how many elements this thread is responsible for.
+      -- Formula: (w - ltid) / group_size (rounded up).
+      let w = product ws
+          ltid = kernelLocalThreadId constants
+          group_size = kernelGroupSize constants
+          to_write = (w - ltid) `quotRoundingUp` group_size
+          is = unflattenIndex ws $ Imp.vi32 i * group_size + ltid
+
+      sFor i Int32 to_write $
+        copyDWIM (patElemName pe) (kernelGroupId constants : is) what is
+
+compileKernelResult constants pe (ThreadsReturn what) = do
+  let is = map (Imp.vi32 . fst) $ kernelDimensions constants
+  sWhen (kernelThreadActive constants) $ copyDWIM (patElemName pe) is what []
+
+compileKernelResult constants pe (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do
+  dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe)
+  let dest_loc_offset = offsetArray dest_loc offset
+      dest' = arrayDestination dest_loc_offset
+  copyDWIMDest dest' [] (Var what) []
+  where offset = case moffset of
+                   Nothing -> toExp' int32 per_thread_elems *
+                              kernelGlobalThreadId constants
+                   Just se -> toExp' int32 se
+
+compileKernelResult constants pe (ConcatReturns (SplitStrided stride) _ _ moffset what) = do
+  dest_loc <- entryArrayLocation <$> lookupArray (patElemName pe)
+  let dest_loc' = strideArray
+                  (offsetArray dest_loc offset) $
+                  toExp' int32 stride
+      dest' = arrayDestination dest_loc'
+  copyDWIMDest dest' [] (Var what) []
+  where offset = case moffset of
+                   Nothing -> kernelGlobalThreadId constants
+                   Just se -> toExp' int32 se
+
+compileKernelResult constants pe (WriteReturn rws _arr dests) = do
+  rws' <- mapM toExp rws
+  forM_ dests $ \(is, e) -> do
+    is' <- mapM toExp is
+    let condInBounds i rw = 0 .<=. i .&&. i .<. rw
+        write = foldl (.&&.) (kernelThreadActive constants) $
+                zipWith condInBounds is' rws'
+    sWhen write $ copyDWIM (patElemName pe) (map (toExp' int32) is) e []
+
+arrayInLocalMemory :: SubExp -> InKernelGen Bool
+arrayInLocalMemory (Var name) = do
+  res <- lookupVar name
+  case res of
+    ArrayVar _ entry ->
+      (Space "local"==) . entryMemSpace <$>
+      lookupMemory (memLocationName (entryArrayLocation entry))
+    _ -> return False
+arrayInLocalMemory Constant{} = return False
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE OverloadedStrings #-}
 -- | Our compilation strategy for 'SegGenRed' is based around avoiding
 -- bin conflicts.  We do this by splitting the input into chunks, and
 -- for each chunk computing a single subhistogram.  Then we combine
@@ -8,11 +9,40 @@
 -- There are some branches around to efficiently handle the case where
 -- we use only a single subhistogram (because it's large), so that we
 -- respect the asymptotics, and do not copy the destination array.
+--
+-- We also use a heuristic strategy for computing subhistograms in
+-- local memory when possible.  Given:
+--
+-- H: total size of histograms in bytes, including any lock arrays.
+--
+-- G: group size
+--
+-- T: number of bytes of local memory each thread can be given without
+-- impacting occupancy (determined experimentally, e.g. 32).
+--
+-- LMAX: maximum amount of local memory per workgroup (hard limit).
+--
+-- We wish to compute:
+--
+-- COOP: cooperation level (number of threads per subhistogram)
+--
+-- LH: number of local memory subhistograms
+--
+-- We do this as:
+--
+-- COOP = ceil(H / T)
+-- LH = ceil((G*T)/H)
+-- if COOP <= G && H <= LMAX then
+--   use local memory
+-- else
+--   use global memory
+
 module Futhark.CodeGen.ImpGen.Kernels.SegGenRed
   ( compileSegGenRed )
   where
 
 import Control.Monad.Except
+import Data.Either
 import Data.Maybe
 import Data.List
 
@@ -23,25 +53,102 @@
 import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
 import Futhark.Pass.ExplicitAllocations()
 import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
-import qualified Futhark.CodeGen.ImpGen as ImpGen
-import Futhark.CodeGen.ImpGen ((<--),
-                               sFor, sComment, sIf, sWhen, sArray,
-                               dPrim_, dPrimV)
+import Futhark.CodeGen.ImpGen
 import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed')
 import Futhark.CodeGen.ImpGen.Kernels.Base
 import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)
 import Futhark.Util (chunks, mapAccumLM, splitFromEnd, takeLast)
 import Futhark.Construct (fullSliceNum)
 
-prepareAtomicUpdate :: Maybe Locking -> [VName] -> Lambda InKernel
-                    -> CallKernelGen (Maybe Locking,
-                                      [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ())
-prepareAtomicUpdate l dests lam =
-  -- We need a separate lock array if the opterators are not all of a
+i32Toi64 :: PrimExp v -> PrimExp v
+i32Toi64 = ConvOpExp (SExt Int32 Int64)
+
+data SubhistosInfo = SubhistosInfo { subhistosArray :: VName
+                                   , subhistosAlloc :: CallKernelGen ()
+                                   }
+
+data SegGenRedSlug = SegGenRedSlug
+                     { slugOp :: GenReduceOp InKernel
+                     , slugNumSubhistos :: VName
+                     , slugSubhistos :: [SubhistosInfo]
+                     }
+
+-- | Figure out how much memory is needed per histogram, and compute
+-- some other auxiliary information.
+computeHistoUsage :: KernelSpace
+                  -> GenReduceOp InKernel
+                  -> CallKernelGen (Imp.Count Imp.Bytes, SegGenRedSlug)
+computeHistoUsage space op = do
+  let segment_dims = init $ spaceDimensions space
+      num_segments = length segment_dims
+
+  op_h <- fmap (sum . map typeSize) $ mapM lookupType $ genReduceDest op
+
+  -- Create names for the intermediate array memory blocks,
+  -- memory block sizes, arrays, and number of subhistograms.
+  num_subhistos <- dPrim "num_subhistos" int32
+  subhisto_infos <- forM (zip (genReduceDest op) (genReduceNeutral op)) $ \(dest, ne) -> do
+    dest_t <- lookupType dest
+    dest_mem <- entryArrayLocation <$> lookupArray dest
+
+    subhistos_mem <-
+      sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device")
+
+    let subhistos_shape = Shape (map snd segment_dims++[Var num_subhistos]) <>
+                          stripDims num_segments (arrayShape dest_t)
+        subhistos_membind = ArrayIn subhistos_mem $ IxFun.iota $
+                            map (primExpFromSubExp int32) $ shapeDims subhistos_shape
+    subhistos <- sArray (baseString dest ++ "_subhistos")
+                 (elemType dest_t) subhistos_shape subhistos_membind
+
+    return $ SubhistosInfo subhistos $ do
+      let unitHistoCase =
+            emit $
+            Imp.SetMem subhistos_mem (memLocationName dest_mem) $
+            Space "device"
+
+          multiHistoCase = do
+            let num_elems = foldl' (*) (Imp.var num_subhistos int32) $
+                            map (toExp' int32) $ arrayDims dest_t
+
+            let subhistos_mem_size =
+                  Imp.bytes $
+                  Imp.innerExp (Imp.elements num_elems `Imp.withElemType` elemType dest_t)
+
+            sAlloc_ subhistos_mem subhistos_mem_size $ Space "device"
+            sReplicate subhistos (Shape (map snd segment_dims ++
+                                         [Var num_subhistos, genReduceWidth op]) <>
+                                  genReduceShape op) ne
+            subhistos_t <- lookupType subhistos
+            let slice = fullSliceNum (map (toExp' int32) $ arrayDims subhistos_t) $
+                        map (unitSlice 0 . toExp' int32 . snd) segment_dims ++
+                        [DimFix 0]
+            sUpdate subhistos slice $ Var dest
+
+      sIf (Imp.var num_subhistos int32 .==. 1) unitHistoCase multiHistoCase
+
+  return (op_h, SegGenRedSlug op num_subhistos subhisto_infos)
+
+localMemLockArray :: KernelSpace -> Type
+localMemLockArray space = Array int32 (Shape [spaceGroupSize space]) NoUniqueness
+
+-- | How many bytes will be spent on lock arrays if we use a local
+-- memory implementation?
+localMemLockUsage :: KernelSpace -> [SegGenRedSlug] -> Imp.Count Imp.Bytes
+localMemLockUsage space slugs =
+  if any (isRight . atomicUpdateLocking . genReduceOp . slugOp) slugs
+  then typeSize $ localMemLockArray space
+  else 0
+
+prepareAtomicUpdateGlobal :: Maybe Locking -> [VName] -> SegGenRedSlug
+                          -> CallKernelGen (Maybe Locking,
+                                            [Imp.Exp] -> ImpM InKernel Imp.KernelOp ())
+prepareAtomicUpdateGlobal l dests slug =
+  -- We need a separate lock array if the operators are not all of a
   -- particularly simple form that permits pure atomic operations.
-  case (l, atomicUpdateLocking lam) of
-    (_, Left f) -> return (l, f dests)
-    (Just l', Right f) -> return (l, f l' dests)
+  case (l, atomicUpdateLocking $ genReduceOp $ slugOp slug) of
+    (_, Left f) -> return (l, f (Space "global") dests)
+    (Just l', Right f) -> return (l, f l' (Space "global") dests)
     (Nothing, Right f) -> do
       -- The number of locks used here is too low, but since we are
       -- currently forced to inline a huge list, I'm keeping it down
@@ -51,120 +158,105 @@
       -- A fun solution would also be to use a simple hashing
       -- algorithm to ensure good distribution of locks.
       let num_locks = 10000
+          dims = map (toExp' int32) $
+                 shapeDims (genReduceShape (slugOp slug)) ++
+                 [ Var (slugNumSubhistos slug)
+                 , genReduceWidth (slugOp slug)]
       locks <-
-        ImpGen.sStaticArray "genred_locks" (Space "device") int32 $
+        sStaticArray "genred_locks" (Space "device") int32 $
         Imp.ArrayZeros num_locks
-      let l' = Locking locks 0 1 0 ((`rem` fromIntegral num_locks) . sum)
-      return (Just l', f l' dests)
+      let l' = Locking locks 0 1 0 ((`rem` fromIntegral num_locks) . flattenIndex dims)
+      return (Just l', f l' (Space "global") dests)
 
-prepareIntermediateArrays :: [SubExp] -> Imp.Exp -> [GenReduceOp InKernel]
-                          -> CallKernelGen
-                             [(VName,
-                               [VName],
-                               [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ())]
-prepareIntermediateArrays segment_dims num_threads = fmap snd . mapAccumLM onOp Nothing
+prepareIntermediateArraysGlobal :: Imp.Exp -> [SegGenRedSlug]
+                                -> CallKernelGen
+                                   [(VName,
+                                     [VName],
+                                     [Imp.Exp] -> ImpM InKernel Imp.KernelOp ())]
+prepareIntermediateArraysGlobal num_threads = fmap snd . mapAccumLM onOp Nothing
   where
-    onOp l op = do
+    onOp l slug@(SegGenRedSlug op num_subhistos subhisto_info) = do
       -- Determining the degree of cooperation (heuristic):
       -- coop_lvl   := size of histogram (Cooperation level)
       -- num_histos := (threads / coop_lvl) (Number of histograms)
       -- threads    := min(physical_threads, segment_size)
       --
       -- Careful to avoid division by zero when genReduceWidth==0.
-      num_histos <- dPrimV "num_histos" $ num_threads `quotRoundingUp`
-                    BinOpExp (SMax Int32) 1 (ImpGen.compileSubExpOfType int32 (genReduceWidth op))
+      num_subhistos <--
+        num_threads `quotRoundingUp`
+        BinOpExp (SMax Int32) 1 (toExp' int32 (genReduceWidth op))
 
-      ImpGen.emit $ Imp.DebugPrint "num_histograms" int32 $ Imp.var num_histos int32
+      emit $ Imp.DebugPrint "Number of subhistograms in global memory" $
+        Just (int32, Imp.vi32 num_subhistos)
 
       -- Initialise sub-histograms.
       --
-      -- If num_histos is 1, then we just reuse the original
+      -- If num_subhistos is 1, then we just reuse the original
       -- destination.  The idea is to avoid a copy if we are writing a
       -- small number of values into a very large prior histogram.
-
-      dests <- forM (zip (genReduceDest op) (genReduceNeutral op)) $ \(dest, ne) -> do
-        dest_t <- lookupType dest
-        dest_mem <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray dest
-        let num_elems = foldl' (*) (Imp.var num_histos int32) $
-                        map (ImpGen.compileSubExpOfType int32) $
-                        arrayDims dest_t
-        let size = Imp.elements num_elems `Imp.withElemType` int32
-
-        (sub_mem, size') <-
-          ImpGen.sDeclareMem "subhistogram_mem" size $ Space "device"
+      dests <- forM (zip (genReduceDest op) subhisto_info) $ \(dest, info) -> do
+        dest_mem <- entryArrayLocation <$> lookupArray dest
 
-        let num_segments = length segment_dims
-            sub_shape = Shape (segment_dims++[Var num_histos]) <>
-                        stripDims num_segments (arrayShape dest_t)
-            sub_membind = ArrayIn sub_mem $ IxFun.iota $
-                          map (primExpFromSubExp int32) $ shapeDims sub_shape
-        subhisto <- sArray "genred_dest" (elemType dest_t) sub_shape sub_membind
+        sub_mem <- fmap memLocationName $
+                   entryArrayLocation <$>
+                   lookupArray (subhistosArray info)
 
         let unitHistoCase =
-              ImpGen.emit $
-              Imp.SetMem sub_mem (ImpGen.memLocationName dest_mem) $
+              emit $
+              Imp.SetMem sub_mem (memLocationName dest_mem) $
               Space "device"
 
-            multiHistoCase = do
-              ImpGen.sAlloc_ sub_mem size' $ Space "device"
-              sReplicate subhisto (Shape $ segment_dims ++ [Var num_histos, genReduceWidth op]) ne
-              subhisto_t <- lookupType subhisto
-              let slice = fullSliceNum (map (ImpGen.compileSubExpOfType int32) $ arrayDims subhisto_t) $
-                          map (unitSlice 0 . ImpGen.compileSubExpOfType int32) segment_dims ++
-                          [DimFix 0]
-              ImpGen.sUpdate subhisto slice $ Var dest
+            multiHistoCase = subhistosAlloc info
 
-        sIf (Imp.var num_histos int32 .==. 1) unitHistoCase multiHistoCase
+        sIf (Imp.var num_subhistos int32 .==. 1) unitHistoCase multiHistoCase
 
-        return subhisto
+        return $ subhistosArray info
 
-      (l', do_op) <- prepareAtomicUpdate l dests $ genReduceOp op
+      (l', do_op) <- prepareAtomicUpdateGlobal l dests slug
 
-      return (l', (num_histos, dests, do_op))
+      return (l', (num_subhistos, dests, do_op))
 
-genRedKernel :: [PatElem ExplicitMemory]
-             -> KernelSpace
-             -> [GenReduceOp InKernel]
-             -> Body InKernel
-             -> CallKernelGen [(VName, [VName])]
-genRedKernel map_pes space ops body = do
+genRedKernelGlobal :: [PatElem ExplicitMemory]
+                  -> KernelSpace
+                  -> [SegGenRedSlug]
+                  -> KernelBody InKernel
+                  -> CallKernelGen ()
+genRedKernelGlobal map_pes space slugs kbody = do
   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
   let constants = base_constants { kernelThreadActive = true }
       (space_is, space_sizes) = unzip $ spaceDimensions space
-      i32_to_i64 = ConvOpExp (SExt Int32 Int64)
-      space_sizes_64 = map (i32_to_i64 . ImpGen.compileSubExpOfType int32) space_sizes
+      space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes
       total_w_64 = product space_sizes_64
 
-  histograms <- prepareIntermediateArrays (init space_sizes) (kernelNumThreads constants) ops
+  histograms <- prepareIntermediateArraysGlobal (kernelNumThreads constants) slugs
 
   elems_per_thread_64 <- dPrimV "elems_per_thread_64" $
                          total_w_64 `quotRoundingUp`
                          ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants)
 
-  sKernel constants "seggenred" $ allThreads constants $ do
+  sKernel constants "seggenred_global" $ allThreads constants $ do
     init_constants
 
-    i <- newVName "i"
-
     -- Compute subhistogram index for each thread, per histogram.
     subhisto_inds <- forM histograms $ \(num_histograms, _, _) ->
       dPrimV "subhisto_ind" $
       kernelGlobalThreadId constants `quot`
       (kernelNumThreads constants `quotRoundingUp` Imp.var num_histograms int32)
 
-    sFor i Int64 (Imp.var elems_per_thread_64 int64) $ do
+    flat_idx <- newVName "flat_idx"
+    sFor flat_idx Int64 (Imp.var elems_per_thread_64 int64) $ do
       -- Compute the offset into the input and output.  To this a
       -- thread can add its local ID to figure out which element it is
       -- responsible for.  The calculation is done with 64-bit
       -- integers to avoid overflow, but the final segment indexes are
       -- 32 bit.
       offset <- dPrimV "offset" $
-                (i32_to_i64 (kernelGroupId constants) *
+                (i32Toi64 (kernelGroupId constants) *
                  (Imp.var elems_per_thread_64 int64 *
-                  i32_to_i64 (kernelGroupSize constants)))
-                + (Imp.var i int64 * i32_to_i64 (kernelGroupSize constants))
+                  i32Toi64 (kernelGroupSize constants)))
+                + (Imp.var flat_idx int64 * i32Toi64 (kernelGroupSize constants))
 
-      j <- dPrimV "j" $ Imp.var offset int64 + i32_to_i64 (kernelLocalThreadId constants)
+      j <- dPrimV "j" $ Imp.var offset int64 + i32Toi64 (kernelLocalThreadId constants)
 
       -- Construct segment indices.
       let setIndex v e = do dPrim_ v int32
@@ -178,98 +270,396 @@
       -- arrays.
       let input_in_bounds = Imp.var j int32 .<. total_w_64
 
-      sWhen input_in_bounds $ ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do
-        let (red_res, map_res) = splitFromEnd (length map_pes) $ bodyResult body
+      sWhen input_in_bounds $ compileStms mempty (kernelBodyStms kbody) $ do
+        let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody
 
         sComment "save map-out results" $
-          forM_ (zip map_pes map_res) $ \(pe, se) ->
-          ImpGen.copyDWIM (patElemName pe)
-          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se []
+          forM_ (zip map_pes map_res) $ \(pe, res) ->
+          copyDWIM (patElemName pe)
+          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants)
+          (kernelResultSubExp res) []
 
-        let (buckets, vs) = splitAt (length ops) red_res
-            perOp = chunks $ map (length . genReduceDest) ops
+        let (buckets, vs) = splitAt (length slugs) red_res
+            perOp = chunks $ map (length . genReduceDest . slugOp) slugs
 
         sComment "perform atomic updates" $
-          forM_ (zip5 ops histograms buckets (perOp vs) subhisto_inds) $
+          forM_ (zip5 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds) $
           \(GenReduceOp dest_w _ _ shape lam,
             (_, _, do_op), bucket, vs', subhisto_ind) -> do
 
-            let bucket' = ImpGen.compileSubExpOfType int32 bucket
-                dest_w' = ImpGen.compileSubExpOfType int32 dest_w
+            let bucket' = toExp' int32 $ kernelResultSubExp bucket
+                dest_w' = toExp' int32 dest_w
                 bucket_in_bounds = 0 .<=. bucket' .&&. bucket' .<. dest_w'
                 bucket_is = map (`Imp.var` int32) (init space_is) ++
                             [Imp.var subhisto_ind int32, bucket']
                 vs_params = takeLast (length vs') $ lambdaParams lam
 
             sWhen bucket_in_bounds $ do
-              ImpGen.dLParams $ lambdaParams lam
-              vectorLoops [] (shapeDims shape) $ \is -> do
+              dLParams $ lambdaParams lam
+              sLoopNest shape $ \is -> do
+                forM_ (zip vs_params vs') $ \(p, res) ->
+                  copyDWIM (paramName p) [] (kernelResultSubExp res) is
+                do_op (bucket_is ++ is)
+
+prepareIntermediateArraysLocal :: KernelSpace -> KernelConstants
+                               -> VName -> [SegGenRedSlug]
+                               -> CallKernelGen
+                                  [([VName],
+                                    InKernelGen ([VName],
+                                                 [Imp.Exp] -> ImpM InKernel Imp.KernelOp ()))]
+prepareIntermediateArraysLocal space constants num_subhistos_per_group =
+  fmap snd . mapAccumLM onOp Nothing
+  where
+    onOp l (SegGenRedSlug op num_subhistos subhisto_info) = do
+
+      num_subhistos <--
+        toExp' int32 (spaceNumGroups space)
+
+      emit $ Imp.DebugPrint "Number of subhistograms in global memory" $
+        Just (int32, Imp.vi32 num_subhistos)
+
+      -- Some trickery is afoot here because we need to construct a
+      -- Locking structure in the CallKernelGen monad, but the actual
+      -- initialisation of the locks array must happen on the device.
+      -- Also, we want only one locks array, no matter how many
+      -- operators need locking.
+      (l', mk_op) <-
+        case (l, atomicUpdateLocking $ genReduceOp op) of
+          (_, Left f) -> return (l, return f)
+          (Just l', Right f) -> return (l, return $ f l')
+          (Nothing, Right f) -> do
+            locks <- newVName "locks"
+            num_locks <- toExp $ spaceGroupSize space
+
+            let dims = map (toExp' int32) $
+                       Var num_subhistos_per_group :
+                       shapeDims (genReduceShape op) ++
+                       [genReduceWidth op]
+                l' = Locking locks 0 1 0 ((`rem` num_locks) . flattenIndex dims)
+                locks_t = localMemLockArray space
+
+                mk_op = do
+                  locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local"
+                  dArray locks int32 (arrayShape locks_t) $
+                    ArrayIn locks_mem $ IxFun.iota $
+                    map (primExpFromSubExp int32) $ arrayDims locks_t
+
+                  sComment "All locks start out unlocked" $
+                    copyDWIM locks [kernelLocalThreadId constants] (intConst Int32 0) []
+
+                  return $ f l'
+
+            return (Just l', mk_op)
+
+      -- Initialise local-memory sub-histograms.  These are
+      -- represented as two-dimensional arrays.
+      let init_local_subhistos = do
+            local_subhistos <-
+              forM (genReduceDest op) $ \dest -> do
+                dest_t <- lookupType dest
+
+                let sub_local_shape =
+                      Shape [Var num_subhistos_per_group] <> arrayShape dest_t
+                sAllocArray "subhistogram_local"
+                  (elemType dest_t) sub_local_shape (Space "local")
+
+            do_op <- mk_op
+
+            return (local_subhistos, do_op (Space "local") local_subhistos)
+
+      -- Initialise global-memory sub-histograms.
+      glob_subhistos <- forM subhisto_info $ \info -> do
+        subhistosAlloc info
+        return $ subhistosArray info
+
+      return (l', (glob_subhistos, init_local_subhistos))
+
+genRedKernelLocal :: VName
+                  -> [PatElem ExplicitMemory]
+                  -> KernelSpace
+                  -> [SegGenRedSlug]
+                  -> KernelBody InKernel
+                  -> CallKernelGen ()
+genRedKernelLocal num_subhistos_per_group_var map_pes space slugs kbody = do
+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
+  let (space_is, space_sizes) = unzip $ spaceDimensions space
+      segment_dims = init space_sizes
+      num_segments = length segment_dims
+      constants = base_constants { kernelThreadActive = true }
+      space_sizes_64 = map (i32Toi64 . toExp' int32) space_sizes
+      total_w_64 = product space_sizes_64
+      num_subhistos_per_group = Imp.var num_subhistos_per_group_var int32
+
+  emit $ Imp.DebugPrint "Number of local subhistograms per group" $ Just (int32, num_subhistos_per_group)
+
+  init_histograms <- prepareIntermediateArraysLocal space constants num_subhistos_per_group_var slugs
+
+  elems_per_thread_64 <- dPrimV "elems_per_thread_64" $
+                         total_w_64 `quotRoundingUp`
+                         ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants)
+
+  sKernel constants "seggenred_local" $ allThreads constants $ do
+    init_constants
+
+    histograms <- forM init_histograms $
+                  \(glob_subhistos, init_local_subhistos) -> do
+      (local_subhistos, do_op) <- init_local_subhistos
+      return (zip glob_subhistos local_subhistos, do_op)
+
+    -- Find index of local subhistograms updated by this thread.  We
+    -- try to ensure, as much as possible, that threads in the same
+    -- warp use different subhistograms, to avoid conflicts.
+    thread_local_subhisto_i <-
+      fmap (`Imp.var` int32) $ dPrimV "thread_local_subhisto_i" $
+      kernelLocalThreadId constants `rem` num_subhistos_per_group
+
+    let (red_res, map_res) = splitFromEnd (length map_pes) $
+                             map kernelResultSubExp $ kernelBodyResult kbody
+        (buckets, vs) = splitAt (length slugs) red_res
+        perOp = chunks $ map (length . genReduceDest . slugOp) slugs
+
+    let onSlugs f = forM_ (zip slugs histograms) $ \(slug, (dests, _)) -> do
+          let histo_dims =
+                map (toExp' int32) $
+                segment_dims ++
+                genReduceWidth (slugOp slug) : shapeDims (genReduceShape (slugOp slug))
+          histo_size <- fmap (`Imp.var` int32) $ dPrimV "histo_size" $
+                        product histo_dims
+          f slug dests histo_dims histo_size
+
+    let onAllHistograms f =
+          onSlugs $ \slug dests histo_dims histo_size -> do
+            let group_hists_size = num_subhistos_per_group * histo_size
+            init_per_thread <- dPrimV "init_per_thread" $
+                               group_hists_size `quotRoundingUp` kernelGroupSize constants
+
+            forM_ (zip dests (genReduceNeutral $ slugOp slug)) $ \((dest_global, dest_local), ne) -> do
+              i <- newVName "local_i"
+              sFor i Int32 (Imp.var init_per_thread int32) $ do
+                j <- fmap (`Imp.var` int32) $ dPrimV "j" $
+                     Imp.var i int32 * kernelGroupSize constants +
+                     kernelLocalThreadId constants
+                j_offset <- fmap (`Imp.var` int32) $ dPrimV "j_offset" $
+                            num_subhistos_per_group *
+                            histo_size *
+                            kernelGroupId constants + j
+
+                local_subhisto_i <- dPrimV "local_subhisto_i" $ j `quot` histo_size
+                let bucket_is = unflattenIndex histo_dims $ j `rem` histo_size
+                global_subhisto_i <- dPrimV "global_subhisto_i" $ j_offset `quot` histo_size
+
+                sWhen (j .<. group_hists_size) $
+                  f dest_local dest_global (slugOp slug) ne
+                  (Imp.var local_subhisto_i int32) (Imp.var global_subhisto_i int32)
+                  bucket_is
+
+    sComment "initialize histograms in local memory" $
+      onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i bucket_is ->
+      sComment "First subhistogram is initialised from global memory; others with neutral element." $ do
+      let global_is = take num_segments bucket_is ++
+                      [0] ++ drop num_segments bucket_is
+          local_is = local_subhisto_i : bucket_is
+      sIf (global_subhisto_i .==. 0)
+        (copyDWIM dest_local local_is (Var dest_global) global_is)
+        (sLoopNest (genReduceShape op) $ \is ->
+            copyDWIM dest_local (local_is++is) ne [])
+
+    sOp Imp.LocalBarrier
+
+    flat_idx <- newVName "flat_idx"
+    sFor flat_idx Int64 (Imp.var elems_per_thread_64 int64) $ do
+      -- Compute the offset into the input and output.  To this a
+      -- thread can add its local ID to figure out which element it is
+      -- responsible for.  The calculation is done with 64-bit
+      -- integers to avoid overflow, but the final segment indexes are
+      -- 32 bit.
+      offset <- dPrimV "offset" $
+                (i32Toi64 (kernelGroupId constants) *
+                 (Imp.var elems_per_thread_64 int64 *
+                  i32Toi64 (kernelGroupSize constants)))
+                + (Imp.var flat_idx int64 * i32Toi64 (kernelGroupSize constants))
+
+      j <- dPrimV "j" $ Imp.var offset int64 + i32Toi64 (kernelLocalThreadId constants)
+
+      -- Construct segment indices.
+      zipWithM_ dPrimV_ space_is $
+        map (ConvOpExp (SExt Int64 Int32)) . unflattenIndex space_sizes_64 $ Imp.var j int64
+
+      -- We execute the bucket function once and update each histogram serially.
+      -- We apply the bucket function if j=offset+ltid is less than
+      -- num_elements.  This also involves writing to the mapout
+      -- arrays.
+      let input_in_bounds = Imp.var j int32 .<. total_w_64
+
+      sWhen input_in_bounds $ compileStms mempty (kernelBodyStms kbody) $ do
+
+        sComment "save map-out results" $
+          forM_ (zip map_pes map_res) $ \(pe, se) ->
+          copyDWIM (patElemName pe)
+          (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se []
+
+        forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $
+          \(GenReduceOp dest_w _ _ shape lam,
+            (_, do_op), bucket, vs') -> do
+
+            let bucket' = toExp' int32 bucket
+                dest_w' = toExp' int32 dest_w
+                bucket_in_bounds = 0 .<=. bucket' .&&. bucket' .<. dest_w'
+                bucket_is = thread_local_subhisto_i :
+                            map (`Imp.var` int32) (init space_is) ++ [bucket']
+                vs_params = takeLast (length vs') $ lambdaParams lam
+
+            sComment "perform atomic updates" $
+              sWhen bucket_in_bounds $ do
+              dLParams $ lambdaParams lam
+              sLoopNest shape $ \is -> do
                 forM_ (zip vs_params vs') $ \(p, v) ->
-                  ImpGen.copyDWIM (paramName p) [] v is
+                  copyDWIM (paramName p) [] v is
                 do_op (bucket_is ++ is)
 
-  let histogramInfo (num_histos, dests, _) = (num_histos, dests)
-  return $ map histogramInfo histograms
+    sOp Imp.LocalBarrier
+    sOp Imp.GlobalBarrier
 
-vectorLoops :: [Imp.Exp] -> [SubExp]
-            -> ([Imp.Exp] -> ImpGen.ImpM lore op ())
-            -> ImpGen.ImpM lore op ()
-vectorLoops is [] f = f $ reverse is
-vectorLoops is (d:ds) f = do
-  i <- newVName "i"
-  d' <- ImpGen.compileSubExp d
-  ImpGen.sFor i Int32 d' $ vectorLoops (Imp.var i int32:is) ds f
+    sComment "Compact the multiple local memory subhistograms to a single subhistogram result" $
+      onSlugs $ \slug dests histo_dims histo_size -> do
+      bins_per_thread <- fmap (`Imp.var` int32) $ dPrimV "init_per_thread" $
+                         histo_size `quotRoundingUp` kernelGroupSize constants
 
+      i <- newVName "local_i"
+      sFor i Int32 bins_per_thread $ do
+        j <- fmap (`Imp.var` int32) $ dPrimV "j" $
+             Imp.var i int32 * kernelGroupSize constants +
+             kernelLocalThreadId constants
+        sWhen (j .<. histo_size) $ do
+          -- We are responsible for compacting the flat bin 'j', which
+          -- we immediately unflatten.
+          let bucket_is = unflattenIndex histo_dims j
+          dLParams $ lambdaParams $ genReduceOp $ slugOp slug
+          let (xparams, yparams) = splitAt (length local_dests) $
+                                   lambdaParams $ genReduceOp $ slugOp slug
+              local_dests = map snd dests
+
+          sComment "Read values from subhistogram 0." $
+            forM_ (zip xparams local_dests) $ \(xp, subhisto) ->
+            copyDWIM
+            (paramName xp) []
+            (Var subhisto) (0:bucket_is)
+
+          sComment "Accumulate based on values in other subhistograms." $ do
+            subhisto_id <- newVName "subhisto_id"
+            sFor subhisto_id Int32 (num_subhistos_per_group - 1) $ do
+              forM_ (zip yparams local_dests) $ \(yp, subhisto) ->
+                copyDWIM
+                (paramName yp) []
+                (Var subhisto) (Imp.var subhisto_id int32 + 1 : bucket_is)
+              compileBody' xparams $ lambdaBody $ genReduceOp $ slugOp slug
+
+          sComment "Put values back in subhistogram 0." $
+            forM_ (zip xparams local_dests) $ \(xp, subhisto) ->
+              copyDWIM
+              subhisto (0:bucket_is)
+              (Var $ paramName xp) []
+
+    sComment "Copy the first local histogram to global memory." $
+      onSlugs $ \_slug dests histo_dims histo_size -> do
+      write_per_thread <- dPrimV "write_per_thread" $
+                          histo_size `quotRoundingUp` kernelGroupSize constants
+
+      forM_ dests $ \(dest_global, dest_local) -> do
+        i <- newVName "local_i"
+        sFor i Int32 (Imp.var write_per_thread int32) $ do
+          j <- fmap (`Imp.var` int32) $ dPrimV "j" $
+               Imp.var i int32 * kernelGroupSize constants +
+               kernelLocalThreadId constants
+
+          sWhen (j .<. histo_size) $ do
+            let bucket_is = unflattenIndex histo_dims $ j `rem` histo_size
+                global_is = take num_segments bucket_is ++
+                            [kernelGroupId constants] ++
+                            drop num_segments bucket_is
+                local_is = 0 : bucket_is
+            copyDWIM dest_global global_is (Var dest_local) local_is
+
+-- Most of this function is not the histogram part itself, but rather
+-- figuring out whether to use a local or global memory strategy, as
+-- well as collapsing the subhistograms produced (which are always in
+-- global memory, but their number may vary).
 compileSegGenRed :: Pattern ExplicitMemory
                  -> KernelSpace
                  -> [GenReduceOp InKernel]
-                 -> Body InKernel
+                 -> KernelBody InKernel
                  -> CallKernelGen ()
-compileSegGenRed (Pattern _ pes) genred_space ops body = do
+compileSegGenRed (Pattern _ pes) space ops kbody = do
   let num_red_res = length ops + sum (map (length . genReduceNeutral) ops)
       (all_red_pes, map_pes) = splitAt num_red_res pes
 
-  infos <- genRedKernel map_pes genred_space ops body
-  let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes
+  let t = 8 * 4
+  g <- toExp $ spaceGroupSize space
+  lmax <- dPrim "lmax" int32
+  sOp $ Imp.GetSizeMax lmax Imp.SizeLocalMemory
 
-  forM_ (zip3 infos pes_per_op ops) $ \((num_histos, subhistos), red_pes, op) -> do
-    let unitHistoCase =
-          -- This is OK because the memory blocks are at least as
-          -- large as the ones we are supposed to use for the result.
-          forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do
-            pe_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$>
-                      ImpGen.lookupArray (patElemName pe)
-            subhisto_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$>
-                            ImpGen.lookupArray subhisto
-            ImpGen.emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device"
+  (op_hs, slugs) <- unzip <$> mapM (computeHistoUsage space) ops
+  h <- fmap (`Imp.var` int32) $
+       dPrimV "h" $ Imp.innerExp $ sum op_hs
+  coop <- fmap (`Imp.var` int32) $
+          dPrimV "coop" $ h `quotRoundingUp` t
 
-    sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do
-      -- For the segmented reduction, we keep the segment dimensions
-      -- unchanged.  To this, we add two dimensions: one over the number
-      -- of buckets, and one over the number of subhistograms.  This
-      -- inner dimension is the one that is collapsed in the reduction.
-      let segment_dims = init $ spaceDimensions genred_space
-          num_buckets = genReduceWidth op
+  -- Check for emptyness to avoid division-by-zero.
+  sUnless (h .==. 0) $ do
+    lh <- dPrimV "lh" $ (g * t) `quotRoundingUp` h
 
-      bucket_id <- newVName "bucket_id"
-      subhistogram_id <- newVName "subhistogram_id"
-      vector_ids <- mapM (const $ newVName "vector_id") $
-                    shapeDims $ genReduceShape op
-      gtid <- newVName $ baseString $ spaceGlobalId genred_space
-      let lam = genReduceOp op
-          segred_space =
-            genred_space
-            { spaceStructure =
-                FlatThreadSpace $
-                segment_dims ++
-                [(bucket_id, num_buckets)] ++
-                zip vector_ids (shapeDims $ genReduceShape op) ++
-                [(subhistogram_id, Var num_histos)]
-            , spaceGlobalId = gtid
-            }
+    emit $ Imp.DebugPrint "\n# SegGenRed" Nothing
+    emit $ Imp.DebugPrint "Cooperation level" $ Just (int32, coop)
+    emit $ Imp.DebugPrint "Memory per set of subhistograms" $ Just (int32, h)
+    emit $ Imp.DebugPrint "Desired group size" $ Just (int32, g)
 
-      compileSegRed' (Pattern [] red_pes) segred_space
-        Commutative lam (genReduceNeutral op) $ \red_dests _ ->
-        forM_ (zip red_dests subhistos) $ \((d, is), subhisto) ->
-          ImpGen.copyDWIM d is (Var subhisto) $ map (`Imp.var` int32) $
-          map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids
+    sIf (Imp.innerExp (localMemLockUsage space slugs) + h * Imp.vi32 lh .<=. Imp.vi32 lmax
+         .&&. coop .<=. g)
+      (genRedKernelLocal lh map_pes space slugs kbody)
+      (genRedKernelGlobal map_pes space slugs kbody)
+
+    let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes
+
+    forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do
+      let num_histos = slugNumSubhistos slug
+          subhistos = map subhistosArray $ slugSubhistos slug
+
+      let unitHistoCase =
+            -- This is OK because the memory blocks are at least as
+            -- large as the ones we are supposed to use for the result.
+            forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do
+              pe_mem <- memLocationName . entryArrayLocation <$>
+                        lookupArray (patElemName pe)
+              subhisto_mem <- memLocationName . entryArrayLocation <$>
+                              lookupArray subhisto
+              emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device"
+
+      sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do
+        -- For the segmented reduction, we keep the segment dimensions
+        -- unchanged.  To this, we add two dimensions: one over the number
+        -- of buckets, and one over the number of subhistograms.  This
+        -- inner dimension is the one that is collapsed in the reduction.
+        let num_buckets = genReduceWidth op
+
+        bucket_id <- newVName "bucket_id"
+        subhistogram_id <- newVName "subhistogram_id"
+        vector_ids <- mapM (const $ newVName "vector_id") $
+                      shapeDims $ genReduceShape op
+        gtid <- newVName $ baseString $ spaceGlobalId space
+        let segred_space =
+              space { spaceStructure =
+                        FlatThreadSpace $
+                        segment_dims ++
+                        [(bucket_id, num_buckets)] ++
+                        zip vector_ids (shapeDims $ genReduceShape op) ++
+                        [(subhistogram_id, Var num_histos)]
+                    , spaceGlobalId = gtid
+                    }
+
+        let segred_op = SegRedOp Commutative (genReduceOp op) (genReduceNeutral op) mempty
+        compileSegRed' (Pattern [] red_pes) segred_space [segred_op] $ \_ red_cont ->
+          red_cont $ flip map subhistos $ \subhisto ->
+            (Var subhisto, map (`Imp.var` int32) $
+              map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids)
+
+  where segment_dims = init $ spaceDimensions space
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs
@@ -0,0 +1,26 @@
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleContexts #-}
+module Futhark.CodeGen.ImpGen.Kernels.SegMap
+  ( compileSegMap ) where
+
+import Control.Monad.Except
+
+import Prelude hiding (quot, rem)
+
+import Futhark.Representation.ExplicitMemory
+import Futhark.CodeGen.ImpGen.Kernels.Base
+
+-- | Compile 'SegMap' instance to host-level code with calls to
+-- various kernels.
+compileSegMap :: Pattern ExplicitMemory
+              -> KernelSpace
+              -> KernelBody InKernel
+              -> CallKernelGen ()
+compileSegMap pat space kbody = do
+  (constants, init_constants) <- kernelInitialisation space
+
+  sKernel constants "segmap" $ do
+    init_constants
+    compileKernelStms constants (kernelBodyStms kbody) $
+      zipWithM_ (compileKernelResult constants) (patternElements pat) $
+      kernelBodyResult kbody
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs
@@ -53,139 +53,169 @@
 
 import Prelude hiding (quot, rem)
 
+import Futhark.Error
 import Futhark.MonadFreshNames
 import Futhark.Transform.Rename
 import Futhark.Representation.ExplicitMemory
 import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
-import qualified Futhark.CodeGen.ImpGen as ImpGen
-import Futhark.CodeGen.ImpGen ((<--),
-                               sFor, sComment, sIf, sWhen,
-                               sOp,
-                               dPrim, dPrimV)
+import Futhark.CodeGen.ImpGen
 import Futhark.CodeGen.ImpGen.Kernels.Base
 import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
+import Futhark.Util (chunks)
 import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)
 
--- | For many kernels, we may not have enough physical groups to cover
--- the logical iteration space.  Some groups thus have to perform
--- double duty; we put an outer loop to accomplish this.  The
--- advantage over just launching a bazillion threads is that the cost
--- of memory expansion should be proportional to the number of
--- *physical* threads (hardware parallelism), not the amount of
--- application parallelism.
-virtualiseGroups :: KernelConstants
-                 -> Imp.Exp
-                 -> (Imp.Exp -> ImpGen.ImpM lore op ())
-                 -> ImpGen.ImpM lore op ()
-virtualiseGroups constants required_groups m = do
-  let group_id = kernelGroupId constants
-      iterations = (required_groups - group_id) `quotRoundingUp` kernelNumGroups constants
-  i <- newVName "i"
-  sFor i Int32 iterations $ m $ group_id + Imp.var i int32 * kernelNumGroups constants
+-- | The maximum number of operators we support in a single SegRed.
+-- This limit arises out of the static allocation of counters.
+maxNumOps :: Int32
+maxNumOps = 10
 
+type DoSegBody = (KernelConstants -> ([(SubExp, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen ())
+
 -- | Compile 'SegRed' instance to host-level code with calls to
 -- various kernels.
 compileSegRed :: Pattern ExplicitMemory
               -> KernelSpace
-              -> Commutativity -> Lambda InKernel -> [SubExp]
-              -> Body InKernel
+              -> [SegRedOp InKernel]
+              -> KernelBody InKernel
               -> CallKernelGen ()
-compileSegRed pat space comm red_op nes body =
-  compileSegRed' pat space comm red_op nes $ \red_dests map_dests ->
-  ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do
-  let (red_res, map_res) = splitAt (length nes) $ bodyResult body
-  sComment "save results to be reduced" $
-    forM_ (zip red_dests red_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se []
-  sComment "save map-out results" $
-    forM_ (zip map_dests map_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se []
+compileSegRed pat space reds body =
+  compileSegRed' pat space reds $ \constants red_cont ->
+  compileKernelStms constants (kernelBodyStms body) $ do
+  let (red_res, map_res) = splitAt (segRedResults reds) $ kernelBodyResult body
 
+  sComment "save map-out results" $ do
+    let map_arrs = drop (segRedResults reds) $ patternElements pat
+    zipWithM_ (compileKernelResult constants) map_arrs map_res
+
+  red_cont $ zip (map kernelResultSubExp red_res) $ repeat []
+
 -- | Like 'compileSegRed', but where the body is a monadic action.
 compileSegRed' :: Pattern ExplicitMemory
                -> KernelSpace
-               -> Commutativity -> Lambda InKernel -> [SubExp]
-               -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())
+               -> [SegRedOp InKernel]
+               -> DoSegBody
                -> CallKernelGen ()
-compileSegRed' pat space comm red_op nes body
+compileSegRed' pat space reds body
+  | genericLength reds > maxNumOps =
+      compilerLimitationS $
+      "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."
   | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space =
-      nonsegmentedReduction pat space comm red_op nes body
+      nonsegmentedReduction pat space reds body
   | otherwise = do
       segment_size <-
-        ImpGen.compileSubExp $ last $ map snd $ spaceDimensions space
-      group_size <- ImpGen.compileSubExp $ spaceGroupSize space
+        toExp $ last $ map snd $ spaceDimensions space
+      group_size <- toExp $ spaceGroupSize space
       let use_small_segments = segment_size * 2 .<. group_size
       sIf use_small_segments
-        (smallSegmentsReduction pat space red_op nes body)
-        (largeSegmentsReduction pat space comm red_op nes body)
-
-nonsegmentedReduction :: Pattern ExplicitMemory
-                      -> KernelSpace
-                      -> Commutativity -> Lambda InKernel -> [SubExp]
-                      -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())
-                      -> CallKernelGen ()
-nonsegmentedReduction segred_pat space comm red_op nes body = do
-  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
-  let constants = base_constants { kernelThreadActive = true }
-      global_tid = kernelGlobalThreadId constants
-      (_, w) = last $ spaceDimensions space
+        (smallSegmentsReduction pat space reds body)
+        (largeSegmentsReduction pat space reds body)
 
+-- | Prepare intermediate arrays for the reduction.  Prim-typed
+-- arguments go in local memory (so we need to do the allocation of
+-- those arrays inside the kernel), while array-typed arguments go in
+-- global memory.  Allocations for the former have already been
+-- performed.  This policy is baked into how the allocations are done
+-- in ExplicitAllocations.
+intermediateArrays :: KernelSpace -> SegRedOp InKernel -> InKernelGen [VName]
+intermediateArrays space (SegRedOp _ red_op nes _) = do
   let red_op_params = lambdaParams red_op
       (red_acc_params, _) = splitAt (length nes) red_op_params
-  red_arrs <- forM red_acc_params $ \p ->
+  forM red_acc_params $ \p ->
     case paramAttr p of
       MemArray pt shape _ (ArrayIn mem _) -> do
         let shape' = Shape [spaceNumThreads space] <> shape
-        ImpGen.sArray "red_arr" pt shape' $
+        sArray "red_arr" pt shape' $
           ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'
       _ -> do
         let pt = elemType $ paramType p
             shape = Shape [spaceGroupSize space]
-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"
-
-  counter <-
-    ImpGen.sStaticArray "counter" (Space "device") int32 $
-    Imp.ArrayValues $ replicate 1 $ IntValue $ Int32Value 0
+        sAllocArray "red_arr" pt shape $ Space "local"
 
-  group_res_arrs <- forM (lambdaReturnType red_op) $ \t -> do
+-- | Arrays for storing group results.
+--
+-- The group-result arrays have an extra dimension (of size groupsize)
+-- because they are also used for keeping vectorised accumulators for
+-- first-stage reduction, if necessary.  When actually storing group
+-- results, the first index is set to 0.
+groupResultArrays :: SubExp -> SubExp -> [SegRedOp InKernel] -> CallKernelGen [[VName]]
+groupResultArrays virt_num_groups group_size reds =
+  forM reds $ \(SegRedOp _ lam _ shape) ->
+    forM (lambdaReturnType lam) $ \t -> do
     let pt = elemType t
-        shape = Shape [spaceNumGroups space] <> arrayShape t
-    ImpGen.sAllocArray "group_res_arr" pt shape $ Space "device"
+        full_shape = Shape [group_size, virt_num_groups] <> shape <> arrayShape t
+        -- Move the groupsize dimension last to ensure coalesced
+        -- memory access.
+        perm = [1..shapeRank full_shape-1] ++ [0]
+    sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm
 
-  sync_arr <- ImpGen.sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
+nonsegmentedReduction :: Pattern ExplicitMemory
+                      -> KernelSpace
+                      -> [SegRedOp InKernel]
+                      -> DoSegBody
+                      -> CallKernelGen ()
+nonsegmentedReduction segred_pat space reds body = do
+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
+  let constants = base_constants { kernelThreadActive = true }
+      global_tid = kernelGlobalThreadId constants
+      (_, w) = last $ spaceDimensions space
 
+  counter <-
+    sStaticArray "counter" (Space "device") int32 $
+    Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0
+
+  -- The group-result arrays have an extra dimension (of size
+  -- groupsize) because they are also used for keeping vectorised
+  -- accumulators for first-stage reduction, if necessary.  When
+  -- actually storing group results, the first index is set to 0.
+  reds_group_res_arrs <- groupResultArrays (spaceNumGroups space) (spaceGroupSize space) reds
+
   num_threads <- dPrimV "num_threads" $ kernelNumThreads constants
 
   sKernel constants "segred_nonseg" $ allThreads constants $ do
     init_constants
 
+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
+    reds_arrs <- mapM (intermediateArrays space) reds
+
     -- Since this is the nonsegmented case, all outer segment IDs must
     -- necessarily be 0.
     let gtids = map fst $ spaceDimensions space
     forM_ (init gtids) $ \v ->
       v <-- 0
 
-    num_elements <- Imp.elements <$> ImpGen.compileSubExp w
+    num_elements <- Imp.elements <$> toExp w
     let elems_per_thread = num_elements `quotRoundingUp` Imp.elements (kernelNumThreads constants)
 
-    (group_result_params, red_op_renamed) <-
-      reductionStageOne constants segred_pat num_elements
+    slugs <- mapM (segRedOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $
+             zip3 reds reds_arrs reds_group_res_arrs
+    reds_op_renamed <-
+      reductionStageOne constants num_elements
       global_tid elems_per_thread num_threads
-      comm red_op nes red_arrs body
+      slugs body
 
-    reductionStageTwo constants segred_pat 0 [0] 0
-      (kernelNumGroups constants) group_result_params red_acc_params red_op_renamed nes
-      1 counter sync_arr group_res_arrs red_arrs
+    let segred_pes = chunks (map (length . segRedNeutral) reds) $
+                     patternElements segred_pat
+    forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes
+           slugs reds_op_renamed [0..]) $
+      \(SegRedOp _ red_op nes _,
+        red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do
+      let red_acc_params = take (length nes) $ lambdaParams red_op
+      reductionStageTwo constants pes (kernelGroupId constants) 0 [0] 0
+        (kernelNumGroups constants) slug red_acc_params red_op_renamed nes
+        1 counter (ValueExp $ IntValue $ Int32Value i)
+        sync_arr group_res_arrs red_arrs
 
 smallSegmentsReduction :: Pattern ExplicitMemory
                        -> KernelSpace
-                       -> Lambda InKernel -> [SubExp]
-                       -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())
+                       -> [SegRedOp InKernel]
+                       -> DoSegBody
                        -> CallKernelGen ()
-smallSegmentsReduction (Pattern _ segred_pes) space red_op nes body = do
+smallSegmentsReduction (Pattern _ segred_pes) space reds body = do
   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
   let constants = base_constants { kernelThreadActive = true }
 
   let (gtids, dims) = unzip $ spaceDimensions space
-  dims' <- mapM ImpGen.compileSubExp dims
+  dims' <- mapM toExp dims
 
   let segment_size = last dims'
   -- Careful to avoid division by zero now.
@@ -196,31 +226,22 @@
       segments_per_group = kernelGroupSize constants `quot` segment_size_nonzero
       required_groups = num_segments `quotRoundingUp` segments_per_group
 
-  let red_op_params = lambdaParams red_op
-      (red_acc_params, _red_next_params) = splitAt (length nes) red_op_params
-  red_arrs <- forM red_acc_params $ \p ->
-    case paramAttr p of
-      MemArray pt shape _ (ArrayIn mem _) -> do
-        let shape' = Shape [spaceNumThreads space] <> shape
-        ImpGen.sArray "red_arr" pt shape' $
-          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'
-      _ -> do
-        let pt = elemType $ paramType p
-            shape = Shape [spaceGroupSize space]
-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"
-
-  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments
-  ImpGen.emit $ Imp.DebugPrint "segment_size" int32 segment_size
-  ImpGen.emit $ Imp.DebugPrint "segments_per_group" int32 segments_per_group
-  ImpGen.emit $ Imp.DebugPrint "required_groups" int32 required_groups
+  emit $ Imp.DebugPrint "\n# SegRed-small" Nothing
+  emit $ Imp.DebugPrint "num_segments" $ Just (int32, num_segments)
+  emit $ Imp.DebugPrint "segment_size" $ Just (int32, segment_size)
+  emit $ Imp.DebugPrint "segments_per_group" $ Just (int32, segments_per_group)
+  emit $ Imp.DebugPrint "required_groups" $ Just (int32, required_groups)
 
   sKernel constants "segred_small" $ allThreads constants $ do
     init_constants
 
+    reds_arrs <- mapM (intermediateArrays space) reds
+
     -- We probably do not have enough actual workgroups to cover the
     -- entire iteration space.  Some groups thus have to perform double
     -- duty; we put an outer loop to accomplish this.
-    virtualiseGroups constants required_groups $ \group_id' -> do
+    virtualiseGroups constants required_groups $ \group_id_var' -> do
+      let group_id' = Imp.vi32 group_id_var'
       -- Compute the 'n' input indices.  The outer 'n-1' correspond to
       -- the segment ID, and are computed from the group id.  The inner
       -- is computed from the local thread id, and may be out-of-bounds.
@@ -231,43 +252,41 @@
       zipWithM_ (<--) (init gtids) $ unflattenIndex (init dims') segment_index
       last gtids <-- index_within_segment
 
-      let toLocalMemory ses =
-            forM_ (zip red_arrs ses) $ \(arr, se) ->
-            ImpGen.copyDWIM arr [ltid] se []
+      let out_of_bounds =
+            forM_ (zip reds reds_arrs) $ \(SegRedOp _ _ nes _, red_arrs) ->
+            forM_ (zip red_arrs nes) $ \(arr, ne) ->
+            copyDWIM arr [ltid] ne []
 
           in_bounds =
-            body (zip red_arrs $ repeat [ltid])
-                 (zip (map patElemName $ drop (length nes) segred_pes) $
-                  repeat $ map (`Imp.var` int32) gtids)
+            body constants $ \red_res ->
+            sComment "save results to be reduced" $ do
+            let red_dests = zip (concat reds_arrs) $ repeat [ltid]
+            forM_ (zip red_dests red_res) $ \((d,d_is), (res, res_is)) ->
+              copyDWIM d d_is res res_is
 
       sComment "apply map function if in bounds" $
         sIf (segment_size .>. 0 .&&.
              isActive (init $ zip gtids dims) .&&.
-             ltid .<. segment_size * segments_per_group) in_bounds (toLocalMemory nes)
+             ltid .<. segment_size * segments_per_group) in_bounds out_of_bounds
 
       sOp Imp.LocalBarrier
 
-      index_i <- newVName "index_i"
-      index_j <- newVName "index_j"
       let crossesSegment from to = (to-from) .>. (to `rem` segment_size)
-          red_op' = red_op { lambdaParams = Param index_i (MemPrim int32) :
-                                            Param index_j (MemPrim int32) :
-                                            lambdaParams red_op }
-
       sWhen (segment_size .>. 0) $
         sComment "perform segmented scan to imitate reduction" $
-        groupScan constants (Just crossesSegment) (segment_size*segments_per_group) red_op' red_arrs
+        forM_ (zip reds reds_arrs) $ \(SegRedOp _ red_op _ _, red_arrs) ->
+        groupScan constants (Just crossesSegment) (segment_size*segments_per_group) red_op red_arrs
 
       sOp Imp.LocalBarrier
 
       sComment "save final values of segments" $
         sWhen (group_id' * segments_per_group + ltid .<. num_segments .&&.
                ltid .<. segments_per_group) $
-        forM_ (zip segred_pes red_arrs) $ \(pe, arr) -> do
+        forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do
         -- Figure out which segment result this thread should write...
         let flat_segment_index = group_id' * segments_per_group + ltid
             gtids' = unflattenIndex (init dims') flat_segment_index
-        ImpGen.copyDWIM (patElemName pe) gtids'
+        copyDWIM (patElemName pe) gtids'
                         (Var arr) [(ltid+1) * segment_size_nonzero - 1]
 
       -- Finally another barrier, because we will be writing to the
@@ -276,58 +295,37 @@
 
 largeSegmentsReduction :: Pattern ExplicitMemory
                        -> KernelSpace
-                       -> Commutativity -> Lambda InKernel -> [SubExp]
-                       -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())
+                       -> [SegRedOp InKernel]
+                       -> DoSegBody
                        -> CallKernelGen ()
-largeSegmentsReduction segred_pat space comm red_op nes body = do
+largeSegmentsReduction segred_pat space reds body = do
   (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
   let (gtids, dims) = unzip $ spaceDimensions space
-  dims' <- mapM ImpGen.compileSubExp dims
+  dims' <- mapM toExp dims
   let segment_size = last dims'
       num_segments = product $ init dims'
 
   let (groups_per_segment, elems_per_thread) =
         groupsPerSegmentAndElementsPerThread segment_size num_segments
         (kernelNumGroups base_constants) (kernelGroupSize base_constants)
-  num_groups <- dPrimV "num_groups" $
+  virt_num_groups <- dPrimV "virt_num_groups" $
     groups_per_segment * num_segments
 
-  num_threads <- dPrimV "num_threads" $
-    Imp.var num_groups int32 * kernelGroupSize base_constants
-
   threads_per_segment <- dPrimV "thread_per_segment" $
     groups_per_segment * kernelGroupSize base_constants
 
-  let constants = base_constants
-                  { kernelThreadActive = true
-                  , kernelNumGroups = Imp.var num_groups int32
-                  , kernelNumThreads = Imp.var num_threads int32
-                  }
-
-  ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments
-  ImpGen.emit $ Imp.DebugPrint "segment_size" int32 segment_size
-  ImpGen.emit $ Imp.DebugPrint "num_groups" int32 (Imp.var num_groups int32)
-  ImpGen.emit $ Imp.DebugPrint "group_size" int32 (kernelGroupSize constants)
-  ImpGen.emit $ Imp.DebugPrint "elems_per_thread" int32 $ Imp.innerExp elems_per_thread
-  ImpGen.emit $ Imp.DebugPrint "groups_per_segment" int32 groups_per_segment
+  let constants = base_constants { kernelThreadActive = true }
 
-  let red_op_params = lambdaParams red_op
-      (red_acc_params, _) = splitAt (length nes) red_op_params
-  red_arrs <- forM red_acc_params $ \p ->
-    case paramAttr p of
-      MemArray pt shape _ (ArrayIn mem _) -> do
-        let shape' = Shape [Var num_threads] <> shape
-        ImpGen.sArray "red_arr" pt shape' $
-          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'
-      _ -> do
-        let pt = elemType $ paramType p
-            shape = Shape [spaceGroupSize space]
-        ImpGen.sAllocArray "red_arr" pt shape $ Space "local"
+  emit $ Imp.DebugPrint "\n# SegRed-large" Nothing
+  emit $ Imp.DebugPrint "num_segments" $ Just (int32, num_segments)
+  emit $ Imp.DebugPrint "segment_size" $ Just (int32, segment_size)
+  emit $ Imp.DebugPrint "virt_num_groups" $ Just (int32, Imp.vi32 virt_num_groups)
+  emit $ Imp.DebugPrint "num_groups" $ Just (int32, kernelNumGroups constants)
+  emit $ Imp.DebugPrint "group_size" $ Just (int32, kernelGroupSize constants)
+  emit $ Imp.DebugPrint "elems_per_thread" $ Just (int32, Imp.innerExp elems_per_thread)
+  emit $ Imp.DebugPrint "groups_per_segment" $ Just (int32, groups_per_segment)
 
-  group_res_arrs <- forM (lambdaReturnType red_op) $ \t -> do
-    let pt = elemType t
-        shape = Shape [Var num_groups] <> arrayShape t
-    ImpGen.sAllocArray "group_res_arr" pt shape $ Space "device"
+  reds_group_res_arrs <- groupResultArrays (Var virt_num_groups) (spaceGroupSize space) reds
 
   -- In principle we should have a counter for every segment.  Since
   -- the number of segments is a dynamic quantity, we would have to
@@ -339,49 +337,68 @@
   -- anywhere?  There are other places in the compiler that will fail
   -- if the group count exceeds the maximum group size, which is at
   -- most 1024 anyway.
-  let num_counters = 1024
+  let num_counters = fromIntegral maxNumOps * 1024
   counter <-
-    ImpGen.sStaticArray "counter" (Space "device") int32 $
+    sStaticArray "counter" (Space "device") int32 $
     Imp.ArrayZeros num_counters
 
-  sync_arr <- ImpGen.sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
-
   sKernel constants "segred_large" $ allThreads constants $ do
     init_constants
-    let segment_gtids = init gtids
-        group_id = kernelGroupId constants
-        group_size = kernelGroupSize constants
-        flat_segment_id = group_id `quot` groups_per_segment
-        local_tid = kernelLocalThreadId constants
 
-        global_tid = kernelGlobalThreadId constants
-                     `rem` (group_size * groups_per_segment)
-        w = last dims
-        first_group_for_segment = flat_segment_id * groups_per_segment
+    reds_arrs <- mapM (intermediateArrays space) reds
+    sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
 
-    zipWithM_ (<--) segment_gtids $ unflattenIndex (init dims') flat_segment_id
-    num_elements <- Imp.elements <$> ImpGen.compileSubExp w
+    -- We probably do not have enough actual workgroups to cover the
+    -- entire iteration space.  Some groups thus have to perform double
+    -- duty; we put an outer loop to accomplish this.
+    virtualiseGroups constants (Imp.vi32 virt_num_groups) $ \group_id_var -> do
 
-    (group_result_params, red_op_renamed) <-
-      reductionStageOne constants segred_pat num_elements
-      global_tid elems_per_thread threads_per_segment
-      comm red_op nes red_arrs body
+      let segment_gtids = init gtids
+          group_id = Imp.vi32 group_id_var
+          group_size = kernelGroupSize constants
+          flat_segment_id = group_id `quot` groups_per_segment
+          local_tid = kernelLocalThreadId constants
 
-    let multiple_groups_per_segment =
-          reductionStageTwo constants segred_pat
-          flat_segment_id (map (`Imp.var` int32) segment_gtids)
-          first_group_for_segment groups_per_segment
-          group_result_params red_acc_params red_op_renamed
-          nes (fromIntegral num_counters) counter sync_arr group_res_arrs red_arrs
+          global_tid = (group_id * group_size + local_tid)
+                       `rem` (group_size * groups_per_segment)
+          w = last dims
+          first_group_for_segment = flat_segment_id * groups_per_segment
 
-        one_group_per_segment =
-          ImpGen.comment "first thread in group saves final result to memory" $
-          sWhen (local_tid .==. 0) $
-            forM_ (take (length nes) $ zip (patternNames segred_pat) group_result_params) $ \(v, p) ->
-            ImpGen.copyDWIM v (map (`Imp.var` int32) segment_gtids) (Var $ paramName p) []
+      zipWithM_ (<--) segment_gtids $ unflattenIndex (init dims') flat_segment_id
+      num_elements <- Imp.elements <$> toExp w
 
-    sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment
+      slugs <- mapM (segRedOpSlug local_tid group_id) $
+               zip3 reds reds_arrs reds_group_res_arrs
+      reds_op_renamed <-
+        reductionStageOne constants num_elements
+        global_tid elems_per_thread threads_per_segment
+        slugs body
 
+      let segred_pes = chunks (map (length . segRedNeutral) reds) $
+                       patternElements segred_pat
+
+          multiple_groups_per_segment =
+            forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes
+                   slugs reds_op_renamed [0..]) $
+            \(SegRedOp _ red_op nes _, red_arrs, group_res_arrs, pes,
+              slug, red_op_renamed, i) -> do
+              let red_acc_params = take (length nes) $ lambdaParams red_op
+              reductionStageTwo constants pes
+                group_id flat_segment_id (map (`Imp.var` int32) segment_gtids)
+                first_group_for_segment groups_per_segment
+                slug red_acc_params red_op_renamed nes
+                (fromIntegral num_counters) counter (ValueExp $ IntValue $ Int32Value i)
+                sync_arr group_res_arrs red_arrs
+
+          one_group_per_segment =
+            comment "first thread in group saves final result to memory" $
+            forM_ (zip slugs segred_pes) $ \(slug, pes) ->
+            sWhen (local_tid .==. 0) $
+              forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->
+              copyDWIM (patElemName v) (map (`Imp.var` int32) segment_gtids) (Var acc) acc_is
+
+      sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment
+
 -- Careful to avoid division by zero here.
 groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp
                                      -> (Imp.Exp, Imp.Count Imp.Elements)
@@ -392,56 +409,107 @@
         segment_size `quotRoundingUp` (group_size * groups_per_segment)
   in (groups_per_segment, Imp.elements elements_per_thread)
 
-reductionStageOne :: KernelConstants
-                  -> Pattern ExplicitMemory
-                  -> Imp.Count Imp.Elements
-                  -> Imp.Exp
-                  -> Imp.Count Imp.Elements
-                  -> VName
-                  -> Commutativity
-                  -> LambdaT InKernel
-                  -> [SubExp]
-                  -> [VName]
-                  -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ())
-                  -> InKernelGen ([LParam InKernel], Lambda InKernel)
-reductionStageOne constants (Pattern _ segred_pes) num_elements global_tid elems_per_thread threads_per_segment comm red_op nes red_arrs body = do
+-- | A SegRedOp with auxiliary information.
+data SegRedOpSlug =
+  SegRedOpSlug
+  { slugOp :: SegRedOp InKernel
+  , slugArrs :: [VName]
+    -- ^ The arrays used for computing the intra-group reduction
+    -- (either local or global memory).
+  , slugAccs :: [(VName, [Imp.Exp])]
+    -- ^ Places to store accumulator in stage 1 reduction.
+  }
 
-  let red_op_params = lambdaParams red_op
-      (red_acc_params, red_next_params) = splitAt (length nes) red_op_params
-      (gtids, _dims) = unzip $ kernelDimensions constants
+slugBody :: SegRedOpSlug -> Body InKernel
+slugBody = lambdaBody . segRedLambda . slugOp
+
+slugParams :: SegRedOpSlug -> [LParam InKernel]
+slugParams = lambdaParams . segRedLambda . slugOp
+
+slugNeutral :: SegRedOpSlug -> [SubExp]
+slugNeutral = segRedNeutral . slugOp
+
+slugShape :: SegRedOpSlug -> Shape
+slugShape = segRedShape . slugOp
+
+slugsComm :: [SegRedOpSlug] -> Commutativity
+slugsComm = mconcat . map (segRedComm . slugOp)
+
+accParams, nextParams :: SegRedOpSlug -> [LParam InKernel]
+accParams slug = take (length (slugNeutral slug)) $ slugParams slug
+nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug
+
+segRedOpSlug :: Imp.Exp -> Imp.Exp -> (SegRedOp InKernel, [VName], [VName]) -> InKernelGen SegRedOpSlug
+segRedOpSlug local_tid group_id (op, group_res_arrs, param_arrs) =
+  SegRedOpSlug op group_res_arrs <$>
+  zipWithM mkAcc (lambdaParams (segRedLambda op)) param_arrs
+  where mkAcc p param_arr
+          | Prim t <- paramType p,
+            shapeRank (segRedShape op) == 0 = do
+              acc <- dPrim (baseString (paramName p) <> "_acc") t
+              return (acc, [])
+          | otherwise =
+              return (param_arr, [local_tid, group_id])
+
+reductionStageZero :: KernelConstants
+                   -> Imp.Count Imp.Elements
+                   -> Imp.Exp
+                   -> Imp.Count Imp.Elements
+                   -> VName
+                   -> [SegRedOpSlug]
+                   -> DoSegBody
+                   -> InKernelGen ([Lambda InKernel], InKernelGen ())
+reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do
+
+  let (gtids, _dims) = unzip $ kernelDimensions constants
       gtid = last gtids
       local_tid = kernelLocalThreadId constants
 
   -- Figure out how many elements this thread should process.
   chunk_size <- dPrim "chunk_size" int32
-  let ordering = case comm of Commutative -> SplitStrided $ Var threads_per_segment
-                              Noncommutative -> SplitContiguous
+  let ordering = case slugsComm slugs of
+                   Commutative -> SplitStrided $ Var threads_per_segment
+                   Noncommutative -> SplitContiguous
   computeThreadChunkSize ordering global_tid elems_per_thread num_elements chunk_size
 
-  ImpGen.dScope Nothing $ scopeOfLParams $ lambdaParams red_op
+  dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs
 
-  forM_ (zip red_acc_params nes) $ \(p, ne) ->
-    ImpGen.copyDWIM (paramName p) [] ne []
+  sComment "neutral-initialise the accumulators" $
+    forM_ slugs $ \slug ->
+    forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->
+    sLoopNest (slugShape slug) $ \vec_is ->
+    copyDWIM acc (acc_is++vec_is) ne []
 
-  red_op_renamed <- renameLambda red_op
+  slugs_op_renamed <- mapM (renameLambda . segRedLambda . slugOp) slugs
 
-  let doTheReduction = do
-        ImpGen.comment "to reduce current chunk, first store our result to memory" $
-          forM_ (zip red_arrs red_acc_params) $ \(arr, p) ->
-          when (primType $ paramType p) $
-          ImpGen.copyDWIM arr [local_tid] (Var $ paramName p) []
+  let doTheReduction =
+        forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) ->
+        sLoopNest (slugShape slug) $ \vec_is -> do
+          comment "to reduce current chunk, first store our result in memory" $ do
+            forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->
+              copyDWIM (paramName p) [] (Var acc) (acc_is++vec_is)
 
-        sOp Imp.LocalBarrier
+            forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) ->
+              when (primType $ paramType p) $
+              copyDWIM arr [local_tid] (Var $ paramName p) []
 
-        groupReduce constants (kernelGroupSize constants) red_op_renamed red_arrs
+          sOp Imp.LocalBarrier
 
-        sOp Imp.LocalBarrier
+          groupReduce constants (kernelGroupSize constants) slug_op_renamed (slugArrs slug)
 
+          sOp Imp.LocalBarrier
+
+          sComment "first thread saves the result in accumulator" $
+            sWhen (local_tid .==. 0) $
+            forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) ->
+            copyDWIM acc (acc_is++vec_is) (Var $ paramName p) []
+
   i <- newVName "i"
   -- If this is a non-commutative reduction, each thread must run the
   -- loop the same number of iterations, because we will be performing
   -- a group-wide reduction in there.
-  let (bound, check_bounds) =
+  let comm = slugsComm slugs
+      (bound, check_bounds) =
         case comm of
           Commutative -> (Imp.var chunk_size int32, id)
           Noncommutative -> (Imp.innerExp elems_per_thread,
@@ -459,78 +527,100 @@
              (index_in_segment * Imp.innerExp elems_per_thread + Imp.var i int32) *
              kernelGroupSize constants
 
-    let red_dests = zip (map paramName red_next_params) $ repeat []
-        map_dests = zip (map patElemName $ drop (length nes) segred_pes) $
-                    repeat $ map (`Imp.var` int32) gtids
+    check_bounds $ sComment "apply map function" $
+      body constants $ \all_red_res -> do
 
-    check_bounds $ sComment "apply map function" $ do
-      body red_dests map_dests
+      let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res
 
-      sComment "apply reduction operator" $
-        ImpGen.compileBody' red_acc_params $ lambdaBody red_op
+      forM_ (zip slugs slugs_res) $ \(slug, red_res) ->
+        sLoopNest (slugShape slug) $ \vec_is -> do
+        sComment "load accumulator" $
+          forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->
+          copyDWIM (paramName p) [] (Var acc) (acc_is ++ vec_is)
+        sComment "load new values" $
+          forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) ->
+          copyDWIM (paramName p) [] res (res_is ++ vec_is)
+        sComment "apply reduction operator" $
+          compileStms mempty (bodyStms $ slugBody slug) $
+          sComment "store in accumulator" $
+          forM_ (zip
+                  (slugAccs slug)
+                  (bodyResult $ slugBody slug)) $ \((acc, acc_is), se) ->
+          copyDWIM acc (acc_is ++ vec_is) se []
 
     case comm of
       Noncommutative -> do
         doTheReduction
-        sComment "first thread takes carry-out; others neutral element" $ do
-          let carry_out =
-                forM_ (zip red_acc_params $ lambdaParams red_op_renamed) $ \(p_to, p_from) ->
-                ImpGen.copyDWIM (paramName p_to) [] (Var $ paramName p_from) []
-              reset_to_neutral =
-                forM_ (zip red_acc_params nes) $ \(p, ne) ->
-                ImpGen.copyDWIM (paramName p) [] ne []
-          sIf (local_tid .==. 0) carry_out reset_to_neutral
+        sComment "first thread keeps accumulator; others reset to neutral element" $ do
+          let reset_to_neutral =
+                forM_ slugs $ \slug ->
+                forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) ->
+                sLoopNest (slugShape slug) $ \vec_is ->
+                copyDWIM acc (acc_is++vec_is) ne []
+          sUnless (local_tid .==. 0) reset_to_neutral
       _ -> return ()
 
-  group_result_params <-
-    case comm of Noncommutative -> return red_acc_params
-                 _              -> do doTheReduction
-                                      return $ lambdaParams red_op_renamed
+  return (slugs_op_renamed, doTheReduction)
 
-  return (group_result_params, red_op_renamed)
+reductionStageOne :: KernelConstants
+                  -> Imp.Count Imp.Elements
+                  -> Imp.Exp
+                  -> Imp.Count Imp.Elements
+                  -> VName
+                  -> [SegRedOpSlug]
+                  -> DoSegBody
+                  -> InKernelGen [Lambda InKernel]
+reductionStageOne constants num_elements global_tid elems_per_thread threads_per_segment slugs body = do
+  (slugs_op_renamed, doTheReduction) <-
+    reductionStageZero constants num_elements global_tid elems_per_thread threads_per_segment slugs body
 
+  case slugsComm slugs of
+    Noncommutative ->
+      forM_ slugs $ \slug ->
+      forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) ->
+      copyDWIM (paramName p) [] (Var acc) acc_is
+    _ -> doTheReduction
+
+  return slugs_op_renamed
+
 reductionStageTwo :: KernelConstants
-                  -> Pattern ExplicitMemory
+                  -> [PatElem ExplicitMemory]
                   -> Imp.Exp
+                  -> Imp.Exp
                   -> [Imp.Exp]
                   -> Imp.Exp
-                  -> PrimExp Imp.ExpLeaf
-                  -> [LParam InKernel]
-                  -> [LParam InKernel]
-                  -> Lambda InKernel
-                  -> [SubExp]
                   -> Imp.Exp
-                  -> VName
-                  -> VName
-                  -> [VName]
-                  -> [VName]
+                  -> SegRedOpSlug
+                  -> [LParam InKernel]
+                  -> Lambda InKernel -> [SubExp]
+                  -> Imp.Exp -> VName -> Imp.Exp -> VName -> [VName] -> [VName]
                   -> InKernelGen ()
-reductionStageTwo constants segred_pat
-                  flat_segment_id segment_gtids first_group_for_segment groups_per_segment
-                  group_result_params red_acc_params
+reductionStageTwo constants segred_pes
+                  group_id flat_segment_id segment_gtids first_group_for_segment groups_per_segment
+                  slug red_acc_params
                   red_op_renamed nes
-                  num_counters counter sync_arr group_res_arrs red_arrs = do
+                  num_counters counter counter_i sync_arr group_res_arrs red_arrs = do
   let local_tid = kernelLocalThreadId constants
-      group_id = kernelGroupId constants
       group_size = kernelGroupSize constants
   old_counter <- dPrim "old_counter" int32
-  (counter_mem, _, counter_offset) <- ImpGen.fullyIndexArray counter [flat_segment_id `rem` num_counters]
-  ImpGen.comment "first thread in group saves group result to memory" $
+  (counter_mem, _, counter_offset) <- fullyIndexArray counter [counter_i * num_counters +
+                                                               flat_segment_id `rem` num_counters]
+  comment "first thread in group saves group result to global memory" $
     sWhen (local_tid .==. 0) $ do
-    forM_ (take (length nes) $ zip group_res_arrs group_result_params) $ \(v, p) ->
-      ImpGen.copyDWIM v [group_id] (Var $ paramName p) []
-    sOp Imp.MemFence
+    forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) ->
+      copyDWIM v [0, group_id] (Var acc) acc_is
+    sOp Imp.MemFenceGlobal
     -- Increment the counter, thus stating that our result is
     -- available.
-    sOp $ Imp.Atomic $ Imp.AtomicAdd old_counter counter_mem counter_offset 1
+    sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset 1
     -- Now check if we were the last group to write our result.  If
     -- so, it is our responsibility to produce the final result.
-    ImpGen.sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1
+    sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1
 
   sOp Imp.LocalBarrier
 
   is_last_group <- dPrim "is_last_group" Bool
-  ImpGen.copyDWIM is_last_group [] (Var sync_arr) [0]
+  copyDWIM is_last_group [] (Var sync_arr) [0]
   sWhen (Imp.var is_last_group Bool) $ do
     -- The final group has written its result (and it was
     -- us!), so read in all the group results and perform the
@@ -539,28 +629,28 @@
     -- with an atomic to avoid warnings about write/write
     -- races in oclgrind.
     sWhen (local_tid .==. 0) $
-      sOp $ Imp.Atomic $ Imp.AtomicAdd old_counter counter_mem counter_offset $
+      sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset $
       negate groups_per_segment
-    ImpGen.comment "read in the per-group-results" $
-      forM_ (zip4 red_acc_params red_arrs nes group_res_arrs) $
-      \(p, arr, ne, group_res_arr) -> do
-        let load_group_result =
-              ImpGen.copyDWIM (paramName p) []
-              (Var group_res_arr) [first_group_for_segment + local_tid]
-            load_neutral_element =
-              ImpGen.copyDWIM (paramName p) [] ne []
-        ImpGen.sIf (local_tid .<. groups_per_segment)
-          load_group_result load_neutral_element
-        when (primType $ paramType p) $
-          ImpGen.copyDWIM arr [local_tid] (Var $ paramName p) []
+    sLoopNest (slugShape slug) $ \vec_is -> do
+      comment "read in the per-group-results" $
+        forM_ (zip4 red_acc_params red_arrs nes group_res_arrs) $
+        \(p, arr, ne, group_res_arr) -> do
+          let load_group_result =
+                copyDWIM (paramName p) []
+                (Var group_res_arr) ([0, first_group_for_segment + local_tid] ++ vec_is)
+              load_neutral_element =
+                copyDWIM (paramName p) [] ne []
+          sIf (local_tid .<. groups_per_segment)
+            load_group_result load_neutral_element
+          when (primType $ paramType p) $
+            copyDWIM arr [local_tid] (Var $ paramName p) []
 
-    sOp Imp.LocalBarrier
+      sOp Imp.LocalBarrier
 
-    sComment "reduce the per-group results" $ do
-      groupReduce constants group_size red_op_renamed red_arrs
+      sComment "reduce the per-group results" $ do
+        groupReduce constants group_size red_op_renamed red_arrs
 
-      sComment "and back to memory with the final result" $
-        sWhen (local_tid .==. 0) $
-        forM_ (take (length nes) $ zip (patternNames segred_pat) $
-               lambdaParams red_op_renamed) $ \(v, p) ->
-        ImpGen.copyDWIM v segment_gtids (Var $ paramName p) []
+        sComment "and back to memory with the final result" $
+          sWhen (local_tid .==. 0) $
+          forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) ->
+          copyDWIM (patElemName pe) (segment_gtids++vec_is) (Var $ paramName p) []
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs
new file mode 100644
--- /dev/null
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs
@@ -0,0 +1,258 @@
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleContexts #-}
+module Futhark.CodeGen.ImpGen.Kernels.SegScan
+  ( compileSegScan )
+  where
+
+import Control.Monad.Except
+import Data.Maybe
+import Data.List
+
+import Prelude hiding (quot, rem)
+
+import Futhark.MonadFreshNames
+import Futhark.Transform.Rename
+import Futhark.Representation.ExplicitMemory
+import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
+import Futhark.CodeGen.ImpGen
+import Futhark.CodeGen.ImpGen.Kernels.Base
+import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
+import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)
+
+makeLocalArrays :: SubExp -> SubExp -> [SubExp] -> Lambda InKernel
+                -> InKernelGen [VName]
+makeLocalArrays group_size num_threads nes scan_op = do
+  let (scan_x_params, _scan_y_params) =
+        splitAt (length nes) $ lambdaParams scan_op
+  forM scan_x_params $ \p ->
+    case paramAttr p of
+      MemArray pt shape _ (ArrayIn mem _) -> do
+        let shape' = Shape [num_threads] <> shape
+        sArray "scan_arr" pt shape' $
+          ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape'
+      _ -> do
+        let pt = elemType $ paramType p
+            shape = Shape [group_size]
+        sAllocArray "scan_arr" pt shape $ Space "local"
+
+type CrossesSegment = Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp)
+
+-- | Produce partially scanned intervals; one per workgroup.
+scanStage1 :: Pattern ExplicitMemory
+           -> KernelSpace
+           -> Lambda InKernel -> [SubExp]
+           -> KernelBody InKernel
+           -> CallKernelGen (Imp.Exp, CrossesSegment)
+scanStage1 (Pattern _ pes) space scan_op nes kbody = do
+  (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return ()
+  let (gtids, dims) = unzip $ spaceDimensions space
+  dims' <- mapM toExp dims
+  let constants = base_constants { kernelThreadActive = true }
+      num_elements = product dims'
+      elems_per_thread = num_elements `quotRoundingUp` kernelNumThreads constants
+      elems_per_group = kernelGroupSize constants * elems_per_thread
+
+  -- Squirrel away a copy of the operator with unique names that we
+  -- can pass to groupScan.
+  scan_op_renamed <- renameLambda scan_op
+
+  let crossesSegment =
+        case reverse dims' of
+          segment_size : _ : _ -> Just $ \from to ->
+            (to-from) .>. (to `rem` segment_size)
+          _ -> Nothing
+
+  sKernel constants "scan_stage1" $ allThreads constants $ do
+    init_constants
+
+    local_arrs <-
+      makeLocalArrays (spaceGroupSize space) (spaceNumThreads space)
+      nes scan_op
+
+    -- The variables from scan_op will be used for the carry and such
+    -- in the big chunking loop.
+    dScope Nothing $ scopeOfLParams $ lambdaParams scan_op
+    let (scan_x_params, scan_y_params) =
+          splitAt (length nes) $ lambdaParams scan_op
+
+    forM_ (zip scan_x_params nes) $ \(p, ne) ->
+      copyDWIM (paramName p) [] ne []
+
+    j <- newVName "j"
+    sFor j Int32 elems_per_thread $ do
+      chunk_offset <- dPrimV "chunk_offset" $
+                      kernelGroupSize constants * Imp.var j int32 +
+                      kernelGroupId constants * elems_per_group
+      flat_idx <- dPrimV "flat_idx" $
+                  Imp.var chunk_offset int32 + kernelLocalThreadId constants
+      -- Construct segment indices.
+      zipWithM_ (<--) gtids $ unflattenIndex dims' $ Imp.var flat_idx int32
+
+      let in_bounds =
+            foldl1 (.&&.) $ zipWith (.<.) (map (`Imp.var` int32) gtids) dims'
+          when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do
+            let (scan_res, map_res) = splitAt (length nes) $ kernelBodyResult kbody
+            sComment "write to-scan values to parameters" $
+              forM_ (zip scan_y_params scan_res) $ \(p, se) ->
+              copyDWIM (paramName p) [] (kernelResultSubExp se) []
+            sComment "write mapped values results to global memory" $
+              forM_ (zip (drop (length nes) pes) map_res) $ \(pe, se) ->
+              copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)
+              (kernelResultSubExp se) []
+          when_out_of_bounds = forM_ (zip scan_y_params nes) $ \(p, ne) ->
+            copyDWIM (paramName p) [] ne []
+
+      sComment "threads in bounds read input; others get neutral element" $
+        sIf in_bounds when_in_bounds when_out_of_bounds
+
+      sComment "combine with carry and write to local memory" $
+        compileStms mempty (bodyStms $ lambdaBody scan_op) $
+        forM_ (zip local_arrs $ bodyResult $ lambdaBody scan_op) $ \(arr, se) ->
+          copyDWIM arr [kernelLocalThreadId constants] se []
+
+      let crossesSegment' = do
+            f <- crossesSegment
+            Just $ \from to ->
+              let from' = from + Imp.var chunk_offset int32
+                  to' = to + Imp.var chunk_offset int32
+              in f from' to'
+
+      groupScan constants crossesSegment'
+        (kernelGroupSize constants) scan_op_renamed local_arrs
+
+      sComment "threads in bounds write partial scan result" $
+        sWhen in_bounds $ forM_ (zip pes local_arrs) $ \(pe, arr) ->
+        copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)
+        (Var arr) [kernelLocalThreadId constants]
+
+      sOp Imp.LocalBarrier
+
+      let load_carry =
+            forM_ (zip local_arrs scan_x_params) $ \(arr, p) ->
+            copyDWIM (paramName p) [] (Var arr) [kernelGroupSize constants - 1]
+          load_neutral =
+            forM_ (zip nes scan_x_params) $ \(ne, p) ->
+            copyDWIM (paramName p) [] ne []
+
+      sComment "first thread reads last element as carry-in for next iteration" $
+        sWhen (kernelLocalThreadId constants .==. 0) $
+        case crossesSegment of Nothing -> load_carry
+                               Just f -> sIf (f (Imp.var chunk_offset int32 +
+                                                 kernelGroupSize constants-1)
+                                                (Imp.var chunk_offset int32 +
+                                                 kernelGroupSize constants))
+                                         load_neutral load_carry
+
+      sOp Imp.LocalBarrier
+
+  return (elems_per_group, crossesSegment)
+
+scanStage2 :: Pattern ExplicitMemory
+           -> Imp.Exp -> CrossesSegment -> KernelSpace
+           -> Lambda InKernel -> [SubExp]
+           -> CallKernelGen ()
+scanStage2 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do
+  -- A single group, with one thread for each group in stage 1.
+  group_size <- toExp $ spaceNumGroups space
+  (constants, init_constants) <-
+    kernelInitialisationSimple 1 group_size Nothing
+
+  let (gtids, dims) = unzip $ spaceDimensions space
+  dims' <- mapM toExp dims
+  let crossesSegment' = do
+        f <- crossesSegment
+        Just $ \from to ->
+          f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1)
+
+  sKernel constants "scan_stage2" $ do
+    init_constants
+
+    local_arrs <- makeLocalArrays (spaceNumGroups space) (spaceNumGroups space)
+                  nes scan_op
+
+    flat_idx <- dPrimV "flat_idx" $
+      (kernelLocalThreadId constants + 1) * elems_per_group - 1
+    -- Construct segment indices.
+    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ Imp.var flat_idx int32
+
+    let in_bounds =
+          foldl1 (.&&.) $ zipWith (.<.) (map (`Imp.var` int32) gtids) dims'
+        when_in_bounds = forM_ (zip local_arrs pes) $ \(arr, pe) ->
+          copyDWIM arr [kernelLocalThreadId constants]
+          (Var $ patElemName pe) $ map (`Imp.var` int32) gtids
+        when_out_of_bounds = forM_ (zip local_arrs nes) $ \(arr, ne) ->
+          copyDWIM arr [kernelLocalThreadId constants] ne []
+
+    sComment "threads in bound read carries; others get neutral element" $
+      sIf in_bounds when_in_bounds when_out_of_bounds
+
+    groupScan constants crossesSegment'
+      (kernelGroupSize constants) scan_op local_arrs
+
+    sComment "threads in bounds write scanned carries" $
+      sWhen in_bounds $ forM_ (zip pes local_arrs) $ \(pe, arr) ->
+      copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids)
+      (Var arr) [kernelLocalThreadId constants]
+
+scanStage3 :: Pattern ExplicitMemory
+           -> Imp.Exp -> CrossesSegment -> KernelSpace
+           -> Lambda InKernel -> [SubExp]
+           -> CallKernelGen ()
+scanStage3 (Pattern _ pes) elems_per_group crossesSegment space scan_op nes = do
+  let (gtids, dims) = unzip $ spaceDimensions space
+  dims' <- mapM toExp dims
+  (constants, init_constants) <- simpleKernelConstants (product dims') "scan"
+  sKernel constants "scan_stage3" $ do
+    init_constants
+    -- Compute our logical index.
+    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ kernelGlobalThreadId constants
+    -- Figure out which group this element was originally in.
+    orig_group <- dPrimV "orig_group" $
+                  kernelGlobalThreadId constants `quot` elems_per_group
+    -- Then the index of the carry-in of the preceding group.
+    carry_in_flat_idx <- dPrimV "carry_in_flat_idx" $
+                         Imp.var orig_group int32 * elems_per_group - 1
+    -- Figure out the logical index of the carry-in.
+    let carry_in_idx = unflattenIndex dims' $ Imp.var carry_in_flat_idx int32
+
+    -- Apply the carry if we are not in the scan results for the first
+    -- group, and are not the last element in such a group (because
+    -- then the carry was updated in stage 2), and we are not crossing
+    -- a segment boundary.
+    let crosses_segment = fromMaybe false $
+          crossesSegment <*>
+            pure (Imp.var carry_in_flat_idx int32) <*>
+            pure (kernelGlobalThreadId constants)
+        is_a_carry = kernelGlobalThreadId constants .==.
+                     (Imp.var orig_group int32 + 1) * elems_per_group - 1
+        no_carry_in = Imp.var orig_group int32 .==. 0 .||. is_a_carry .||. crosses_segment
+
+    sWhen (kernelThreadActive constants) $ sUnless no_carry_in $ do
+      dScope Nothing $ scopeOfLParams $ lambdaParams scan_op
+      let (scan_x_params, scan_y_params) =
+            splitAt (length nes) $ lambdaParams scan_op
+      forM_ (zip scan_x_params pes) $ \(p, pe) ->
+        copyDWIM (paramName p) [] (Var $ patElemName pe) carry_in_idx
+      forM_ (zip scan_y_params pes) $ \(p, pe) ->
+        copyDWIM (paramName p) [] (Var $ patElemName pe) $ map (`Imp.var` int32) gtids
+      compileBody' scan_x_params $ lambdaBody scan_op
+      forM_ (zip scan_x_params pes) $ \(p, pe) ->
+        copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids) (Var $ paramName p) []
+
+-- | Compile 'SegScan' instance to host-level code with calls to
+-- various kernels.
+compileSegScan :: Pattern ExplicitMemory
+               -> KernelSpace
+               -> Lambda InKernel -> [SubExp]
+               -> KernelBody InKernel
+               -> CallKernelGen ()
+compileSegScan pat space scan_op nes kbody = do
+  (elems_per_group, crossesSegment) <- scanStage1 pat space scan_op nes kbody
+
+  emit $ Imp.DebugPrint "\n# SegScan" Nothing
+  emit $ Imp.DebugPrint "elems_per_group" $ Just (int32, elems_per_group)
+
+  scan_op' <- renameLambda scan_op
+  scan_op'' <- renameLambda scan_op
+  scanStage2 pat elems_per_group crossesSegment space scan_op' nes
+  scanStage3 pat elems_per_group crossesSegment space scan_op'' nes
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs
@@ -48,7 +48,7 @@
       opencl_code = openClCode $ M.elems kernels
       opencl_prelude = pretty $ genPrelude target requirements
   return $ ImpOpenCL.Program opencl_code opencl_prelude kernel_names
-    (S.toList $ kernelUsedTypes requirements) sizes $
+    (S.toList $ openclUsedTypes requirements) sizes $
     ImpOpenCL.Functions (M.toList extra_funs) <> prog'
   where genPrelude TargetOpenCL = genOpenClPrelude
         genPrelude TargetCUDA = genCUDAPrelude
@@ -63,10 +63,18 @@
 pointerQuals "kernel"     = return [C.ctyquals|__kernel|]
 pointerQuals s            = fail $ "'" ++ s ++ "' is not an OpenCL kernel address space."
 
-type UsedFunctions = [(String,C.Func)] -- The ordering is important!
+newtype KernelRequirements =
+  KernelRequirements { kernelLocalMemory :: [LocalMemoryUse] }
 
+instance Semigroup KernelRequirements where
+  KernelRequirements lm1 <> KernelRequirements lm2 =
+    KernelRequirements (lm1<>lm2)
+
+instance Monoid KernelRequirements where
+  mempty = KernelRequirements mempty
+
 newtype OpenClRequirements =
-  OpenClRequirements { kernelUsedTypes :: S.Set PrimType }
+  OpenClRequirements { openclUsedTypes :: S.Set PrimType }
 
 instance Semigroup OpenClRequirements where
   OpenClRequirements ts1 <> OpenClRequirements ts2 =
@@ -104,16 +112,17 @@
 onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL
 
 onKernel target kernel = do
-  let (kernel_body, _) =
-        GenericC.runCompilerM (Functions []) inKernelOperations blankNameSource mempty $
+  let (kernel_body, requirements) =
+        GenericC.runCompilerM mempty inKernelOperations blankNameSource mempty $
         GenericC.blockScope $ GenericC.compileCode $ kernelBody kernel
 
       use_params = mapMaybe useAsParam $ kernelUses kernel
 
-      (local_memory_params, local_memory_init) =
-        unzip $
+      (local_memory_args, local_memory_params, local_memory_init) =
+        unzip3 $
         flip evalState (blankNameSource :: VNameSource) $
-        mapM (prepareLocalMemory target) $ kernelLocalMemory kernel
+        mapM (prepareLocalMemory target) $ kernelLocalMemory $
+        GenericC.compUserState requirements
 
       -- CUDA has very strict restrictions on the number of blocks
       -- permitted along the 'y' and 'z' dimensions of the grid
@@ -151,27 +160,30 @@
               , clRequirements = OpenClRequirements (typesInKernel kernel)
               }
 
-  return $ LaunchKernel name (kernelArgs kernel) num_groups group_size
+  return $ LaunchKernel name (catMaybes local_memory_args ++ kernelArgs kernel) num_groups group_size
   where name = nameToString $ kernelName kernel
         num_groups = kernelNumGroups kernel
         group_size = kernelGroupSize kernel
 
-        prepareLocalMemory TargetOpenCL (mem, Left _) = do
+        prepareLocalMemory TargetOpenCL (mem, Left size) = do
           mem_aligned <- newVName $ baseString mem ++ "_aligned"
-          return (Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|],
+          return (Just $ SharedMemoryKArg size,
+                  Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|],
                   [C.citem|__local volatile char* restrict $id:mem = $id:mem_aligned;|])
         prepareLocalMemory TargetOpenCL (mem, Right size) = do
           let size' = compilePrimExp size
-          return (Nothing,
+          return (Nothing, Nothing,
                   [C.citem|ALIGNED_LOCAL_MEMORY($id:mem, $exp:size');|])
-        prepareLocalMemory TargetCUDA (mem, Left _) = do
+        prepareLocalMemory TargetCUDA (mem, Left size) = do
           param <- newVName $ baseString mem ++ "_offset"
-          return (Just [C.cparam|uint $id:param|],
+          return (Just $ SharedMemoryKArg size,
+                  Just [C.cparam|uint $id:param|],
                   [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|])
         prepareLocalMemory TargetCUDA (mem, Right size) = do
+          -- We declare the shared memory array as int64_t to force alignment.
           let size' = compilePrimExp size
-          return (Nothing,
-                  [CUDAC.citem|__shared__ volatile char $id:mem[$exp:size'];|])
+          return (Nothing, Nothing,
+                  [CUDAC.citem|__shared__ volatile typename int64_t $id:mem[(($exp:size' + 7) & ~7)/8];|])
 
 useAsParam :: KernelUse -> Maybe C.Param
 useAsParam (ScalarUse name bt) =
@@ -384,16 +396,14 @@
           return [C.cexp|$id:(zEncodeString (pretty key))|]
 
 kernelArgs :: Kernel -> [KernelArg]
-kernelArgs kernel =
-  mapMaybe (fmap (SharedMemoryKArg . memSizeToExp) . localMemorySize)
-  (kernelLocalMemory kernel) ++
-  mapMaybe useToArg (kernelUses kernel)
-  where localMemorySize (_, Left size) = Just size
-        localMemorySize (_, Right{}) = Nothing
+kernelArgs = mapMaybe useToArg . kernelUses
+  where useToArg (MemoryUse mem)  = Just $ MemKArg mem
+        useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt
+        useToArg ConstUse{}       = Nothing
 
 --- Generating C
 
-inKernelOperations :: GenericC.Operations KernelOp UsedFunctions
+inKernelOperations :: GenericC.Operations KernelOp KernelRequirements
 inKernelOperations = GenericC.Operations
                      { GenericC.opsCompiler = kernelOps
                      , GenericC.opsMemoryType = kernelMemoryType
@@ -405,7 +415,7 @@
                      , GenericC.opsStaticArray = noStaticArrays
                      , GenericC.opsFatMemory = False
                      }
-  where kernelOps :: GenericC.OpCompiler KernelOp UsedFunctions
+  where kernelOps :: GenericC.OpCompiler KernelOp KernelRequirements
         kernelOps (GetGroupId v i) =
           GenericC.stm [C.cstm|$id:v = get_group_id($int:i);|]
         kernelOps (GetLocalId v i) =
@@ -422,74 +432,83 @@
           GenericC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|]
         kernelOps GlobalBarrier =
           GenericC.stm [C.cstm|barrier(CLK_GLOBAL_MEM_FENCE);|]
-        kernelOps MemFence =
+        kernelOps MemFenceLocal =
+          GenericC.stm [C.cstm|mem_fence_local();|]
+        kernelOps MemFenceGlobal =
           GenericC.stm [C.cstm|mem_fence_global();|]
-        kernelOps (Atomic aop) = atomicOps aop
+        kernelOps (PrivateAlloc name size) = do
+          size' <- GenericC.compileExp $ innerExp size
+          name' <- newVName $ pretty name ++ "_backing"
+          GenericC.item [C.citem|__private char $id:name'[$exp:size'];|]
+          GenericC.stm [C.cstm|$id:name = $id:name';|]
+        kernelOps (LocalAlloc name size) = do
+          name' <- newVName $ pretty name ++ "_backing"
+          GenericC.modifyUserState (<>KernelRequirements [(name', size)])
+          GenericC.stm [C.cstm|$id:name = (__local char*) $id:name';|]
+        kernelOps (Atomic space aop) = atomicOps space aop
 
-        atomicOps (AtomicAdd old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_add((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]
+        atomicCast s t = do
+          let volatile = [C.ctyquals|volatile|]
+          quals <- case s of DefaultSpace -> pointerQuals "global"
+                             Space sid    -> pointerQuals sid
+          return [C.cty|$tyquals:(volatile++quals) $ty:t|]
 
-        atomicOps (AtomicSMax old arr ind val) = do
+        doAtomic s old arr ind val op ty = do
           ind' <- GenericC.compileExp $ innerExp ind
           val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]
+          cast <- atomicCast s ty
+          GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|]
 
-        atomicOps (AtomicSMin old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]
+        atomicOps s (AtomicAdd old arr ind val) =
+          doAtomic s old arr ind val "atomic_add" [C.cty|int|]
 
-        atomicOps (AtomicUMax old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]
+        atomicOps s (AtomicSMax old arr ind val) =
+          doAtomic s old arr ind val "atomic_max" [C.cty|int|]
 
-        atomicOps (AtomicUMin old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]
+        atomicOps s (AtomicSMin old arr ind val) =
+          doAtomic s old arr ind val "atomic_min" [C.cty|int|]
 
-        atomicOps (AtomicAnd old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_and((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]
+        atomicOps s (AtomicUMax old arr ind val) =
+          doAtomic s old arr ind val "atomic_max" [C.cty|unsigned int|]
 
-        atomicOps (AtomicOr old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_or((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]
+        atomicOps s (AtomicUMin old arr ind val) =
+          doAtomic s old arr ind val "atomic_min" [C.cty|unsigned int|]
 
-        atomicOps (AtomicXor old arr ind val) = do
-          ind' <- GenericC.compileExp $ innerExp ind
-          val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_xor((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|]
+        atomicOps s (AtomicAnd old arr ind val) =
+          doAtomic s old arr ind val "atomic_and" [C.cty|unsigned int|]
 
-        atomicOps (AtomicCmpXchg old arr ind cmp val) = do
+        atomicOps s (AtomicOr old arr ind val) =
+          doAtomic s old arr ind val "atomic_or" [C.cty|unsigned int|]
+
+        atomicOps s (AtomicXor old arr ind val) =
+          doAtomic s old arr ind val "atomic_xor" [C.cty|unsigned int|]
+
+        atomicOps s (AtomicCmpXchg old arr ind cmp val) = do
           ind' <- GenericC.compileExp $ innerExp ind
           cmp' <- GenericC.compileExp cmp
           val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_cmpxchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:cmp', $exp:val');|]
+          cast <- atomicCast s [C.cty|int|]
+          GenericC.stm [C.cstm|$id:old = atomic_cmpxchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|]
 
-        atomicOps (AtomicXchg old arr ind val) = do
+        atomicOps s (AtomicXchg old arr ind val) = do
           ind' <- GenericC.compileExp $ innerExp ind
           val' <- GenericC.compileExp val
-          GenericC.stm [C.cstm|$id:old = atomic_xchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|]
+          cast <- atomicCast s [C.cty|int|]
+          GenericC.stm [C.cstm|$id:old = atomic_xchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|]
 
-        cannotAllocate :: GenericC.Allocate KernelOp UsedFunctions
+        cannotAllocate :: GenericC.Allocate KernelOp KernelRequirements
         cannotAllocate _ =
           fail "Cannot allocate memory in kernel"
 
-        cannotDeallocate :: GenericC.Deallocate KernelOp UsedFunctions
+        cannotDeallocate :: GenericC.Deallocate KernelOp KernelRequirements
         cannotDeallocate _ _ =
           fail "Cannot deallocate memory in kernel"
 
-        copyInKernel :: GenericC.Copy KernelOp UsedFunctions
+        copyInKernel :: GenericC.Copy KernelOp KernelRequirements
         copyInKernel _ _ _ _ _ _ _ =
           fail "Cannot bulk copy in kernel."
 
-        noStaticArrays :: GenericC.StaticArray KernelOp UsedFunctions
+        noStaticArrays :: GenericC.StaticArray KernelOp KernelRequirements
         noStaticArrays _ _ _ _ =
           fail "Cannot create static array in kernel."
 
@@ -499,11 +518,6 @@
 
 --- Checking requirements
 
-useToArg :: KernelUse -> Maybe KernelArg
-useToArg (MemoryUse mem)  = Just $ MemKArg mem
-useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt
-useToArg ConstUse{}       = Nothing
-
 typesInKernel :: Kernel -> S.Set PrimType
 typesInKernel kernel = typesInCode $ kernelBody kernel
 
@@ -530,7 +544,7 @@
   typesInExp e <> typesInCode c1 <> typesInCode c2
 typesInCode (Assert e _ _) = typesInExp e
 typesInCode (Comment _ c) = typesInCode c
-typesInCode (DebugPrint _ _ e) = typesInExp e
+typesInCode (DebugPrint _ v) = maybe mempty (typesInExp . snd) v
 typesInCode Op{} = mempty
 
 typesInExp :: Exp -> S.Set PrimType
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs
@@ -80,8 +80,8 @@
       , dec index_out $ v32 x_index * height + v32 y_index
 
       , If (v32 get_global_id_0 .<. input_size)
-        (Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize) t (Space "global") Nonvolatile $
-         index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize) t (Space "global") Nonvolatile)
+        (Write odata (elements $ v32 odata_offset + v32 index_out) t (Space "global") Nonvolatile $
+         index idata (elements $ v32 idata_offset + v32 index_in) t (Space "global") Nonvolatile)
         mempty
       ]
 
@@ -113,10 +113,10 @@
         in mconcat [ dec index_in $ (v32 y_index + i) * width + v32 x_index
                    , when (v32 y_index + i .<. height .&&.
                            v32 index_in .<. input_size) $
-                     Write block (bytes $ ((v32 get_local_id_1 + i) * (tile_dim+1)
-                                           + v32 get_local_id_0) * tsize)
+                     Write block (elements $ (v32 get_local_id_1 + i) * (tile_dim+1)
+                                             + v32 get_local_id_0)
                      t (Space "local") Nonvolatile $
-                     index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize)
+                     index idata (elements $ v32 idata_offset + v32 index_in)
                      t (Space "global") Nonvolatile]
       , Op LocalBarrier
       , SetScalar x_index $ v32 get_group_id_1 * tile_dim + v32 get_local_id_0
@@ -127,17 +127,16 @@
         in mconcat [ dec index_out $ (v32 y_index + i) * height + v32 x_index
                    , when (v32 y_index + i .<. width .&&.
                            v32 index_out .<. output_size) $
-                     Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize)
+                     Write odata (elements $ v32 odata_offset + v32 index_out)
                      t (Space "global") Nonvolatile $
-                     index block (bytes $ (v32 get_local_id_0 * (tile_dim+1)
-                                           +v32 get_local_id_1+i)*tsize)
+                     index block (elements $ v32 get_local_id_0 * (tile_dim+1)
+                                             + v32 get_local_id_1+i)
                      t (Space "local") Nonvolatile
                    ]
       ]
 
   where dec v e = DeclareScalar v int32 <> SetScalar v e
         v32 = flip var int32
-        tsize = LeafExp (SizeOf t) int32
         tile_dim = 2 * block_dim
 
         when a b = If a b mempty
@@ -198,19 +197,18 @@
           , dec y_index y_in_index
           , dec index_in $ v32 y_index * width + v32 x_index
           , when (v32 x_index .<. width .&&. v32 y_index .<. height .&&. v32 index_in .<. input_size) $
-            Write block (bytes $ (v32 get_local_id_1 * (block_dim+1) + v32 get_local_id_0) * tsize)
+            Write block (elements $ v32 get_local_id_1 * (block_dim+1) + v32 get_local_id_0)
             t (Space "local") Nonvolatile $
-            index idata (bytes $ (v32 idata_offset + v32 index_in) * tsize)
+            index idata (elements $ v32 idata_offset + v32 index_in)
             t (Space "global") Nonvolatile
           , Op LocalBarrier
           , SetScalar x_index x_out_index
           , SetScalar y_index y_out_index
           , dec index_out $ v32 y_index * height + v32 x_index
           , when (v32 x_index .<. height .&&. v32 y_index .<. width .&&. v32 index_out .<. output_size) $
-            Write odata (bytes $ (v32 odata_offset + v32 index_out) * tsize)
+            Write odata (elements $ v32 odata_offset + v32 index_out)
             t (Space "global") Nonvolatile $
-            index block (bytes $ (v32 get_local_id_0 * (block_dim+1)
-                                   +v32 get_local_id_1)*tsize)
+            index block (elements $ v32 get_local_id_0 * (block_dim+1) + v32 get_local_id_1)
             t (Space "local") Nonvolatile
           ]
 
@@ -218,8 +216,9 @@
                    -> Kernel
 mapTransposeKernel desc block_dim_int args t kind =
   Kernel
-  { kernelBody = mapTranspose block_dim args t kind
-  , kernelLocalMemory = [(block, Right block_size)]
+  { kernelBody = DeclareMem block (Space "local") <>
+                 Op (LocalAlloc block (Right block_size)) <>
+                 mapTranspose block_dim args t kind
   , kernelUses = uses
   , kernelNumGroups = num_groups
   , kernelGroupSize = group_size
diff --git a/src/Futhark/CodeGen/ImpGen/Sequential.hs b/src/Futhark/CodeGen/ImpGen/Sequential.hs
--- a/src/Futhark/CodeGen/ImpGen/Sequential.hs
+++ b/src/Futhark/CodeGen/ImpGen/Sequential.hs
@@ -11,7 +11,7 @@
 import Futhark.MonadFreshNames
 
 compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program)
-compileProg = ImpGen.compileProg ops Imp.DefaultSpace []
+compileProg = ImpGen.compileProg ops Imp.DefaultSpace
   where ops = ImpGen.defaultOperations opCompiler
         opCompiler :: ImpGen.OpCompiler ExplicitMemory Imp.Sequential
         opCompiler dest (Alloc e space) =
diff --git a/src/Futhark/CodeGen/SetDefaultSpace.hs b/src/Futhark/CodeGen/SetDefaultSpace.hs
--- a/src/Futhark/CodeGen/SetDefaultSpace.hs
+++ b/src/Futhark/CodeGen/SetDefaultSpace.hs
@@ -33,8 +33,8 @@
   TransparentValue $ setValueSpace space v
 
 setValueSpace :: Space -> ValueDesc -> ValueDesc
-setValueSpace space (ArrayValue mem memsize _ bt ept shape) =
-  ArrayValue mem memsize space bt ept shape
+setValueSpace space (ArrayValue mem _ bt ept shape) =
+  ArrayValue mem space bt ept shape
 setValueSpace _ (ScalarValue bt ept v) =
   ScalarValue bt ept v
 
@@ -81,8 +81,8 @@
         setArgSpace (ExpArg e) = ExpArg $ setExpSpace space e
 setBodySpace space (Assert e msg loc) =
   Assert (setExpSpace space e) msg loc
-setBodySpace space (DebugPrint s t e) =
-  DebugPrint s t (setExpSpace space e)
+setBodySpace space (DebugPrint s v) =
+  DebugPrint s $ fmap (fmap (setExpSpace space)) v
 setBodySpace _ (Op op) =
   Op op
 
diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs
--- a/src/Futhark/Construct.hs
+++ b/src/Futhark/Construct.hs
@@ -374,7 +374,7 @@
   x   <- newVName "x"
   y   <- newVName "y"
   body <- insertStmsM $ do
-    res <- letSubExp "res" $ BasicOp $ bop (Var x) (Var y)
+    res <- letSubExp "binlam_res" $ BasicOp $ bop (Var x) (Var y)
     return $ resultBody [res]
   return Lambda {
              lambdaParams     = [Param x (Prim arg_t),
diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs
--- a/src/Futhark/Internalise.hs
+++ b/src/Futhark/Internalise.hs
@@ -12,6 +12,7 @@
 
 import Control.Monad.State
 import Control.Monad.Reader
+import Data.Bitraversable
 import qualified Data.Map.Strict as M
 import qualified Data.Set as S
 import Data.List
@@ -92,7 +93,7 @@
       internaliseBody body >>=
         ensureResultExtShape asserting msg loc (map I.fromDecl rettype')
 
-    let free_in_fun = freeInBody body' `S.difference` normal_param_names
+    let free_in_fun = freeIn body' `S.difference` normal_param_names
 
     used_free_params <- forM (S.toList free_in_fun) $ \v -> do
       v_t <- lookupType v
@@ -122,12 +123,40 @@
     -- them from somewhere else.
     zeroExts ts = generaliseExtTypes ts ts
 
+allDimsFreshInType :: MonadFreshNames m => E.PatternType -> m E.PatternType
+allDimsFreshInType = bitraverse onDim pure
+  where onDim (E.NamedDim v) =
+          E.NamedDim . E.qualName <$> newVName (baseString $ E.qualLeaf v)
+        onDim _ = pure AnyDim
+
+-- | Replace all named dimensions with a fresh name, and remove all
+-- constant dimensions.  The point is to remove the constraints, but
+-- keep the names around.  We use this for constructing the entry
+-- point parameters.
+allDimsFreshInPat :: MonadFreshNames m => E.Pattern -> m E.Pattern
+allDimsFreshInPat (PatternAscription p _ _) =
+  allDimsFreshInPat p
+allDimsFreshInPat (PatternParens p _) =
+  allDimsFreshInPat p
+allDimsFreshInPat (Id v (Info t) loc) =
+  Id v <$> (Info <$> allDimsFreshInType t) <*> pure loc
+allDimsFreshInPat (TuplePattern ps loc) =
+  TuplePattern <$> mapM allDimsFreshInPat ps <*> pure loc
+allDimsFreshInPat (RecordPattern ps loc) =
+  RecordPattern <$> mapM (traverse allDimsFreshInPat) ps <*> pure loc
+allDimsFreshInPat (Wildcard (Info t) loc) =
+  Wildcard <$> (Info <$> allDimsFreshInType t) <*> pure loc
+allDimsFreshInPat (PatternLit e (Info t) loc) =
+  PatternLit e <$> (Info <$> allDimsFreshInType t) <*> pure loc
+
 generateEntryPoint :: E.ValBind -> InternaliseM ()
-generateEntryPoint (E.ValBind _ ofname retdecl (Info rettype) _ params _ _ loc) =
-  -- We remove all shape annotations, so there should be no constant
+generateEntryPoint (E.ValBind _ ofname retdecl (Info rettype) _ params _ _ loc) = do
+  -- We replace all shape annotations, so there should be no constant
   -- parameters here.
-  bindingParams [] (map E.patternNoShapeAnnotations params) $
-  \_ shapeparams params' -> do
+  params_fresh <- mapM allDimsFreshInPat params
+  let tparams = map (`E.TypeParamDim` noLoc) $ S.toList $
+                mconcat $ map E.patternDimNames params_fresh
+  bindingParams tparams params_fresh $ \_ shapeparams params' -> do
     (entry_rettype, _) <- internaliseEntryReturnType $ anyDimShapeAnnotations rettype
     let entry' = entryPoint (zip params params') (retdecl, rettype, entry_rettype)
         args = map (I.Var . I.paramName) $ concat params'
@@ -415,6 +444,7 @@
   (qfname, args, _) <- findFuncall e
   let fname = nameFromString $ pretty $ baseName $ qualLeaf qfname
       loc = srclocOf e
+      arg_desc = nameToString fname ++ "_arg"
 
   -- Some functions are magical (overloaded) and we handle that here.
   -- Note that polymorphic functions (which are not magical) are not
@@ -424,27 +454,27 @@
            internalise desc
        | Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do
            let tag ses = [ (se, I.Observe) | se <- ses ]
-           args' <- mapM (internaliseExp "arg") args
+           args' <- mapM (internaliseExp arg_desc) args
            let args'' = concatMap tag args'
            letTupExp' desc $ I.Apply fname args'' [I.Prim rettype] (Safe, loc, [])
        | otherwise -> do
-           args' <- concat <$> mapM (internaliseExp "arg") args
+           args' <- concat <$> mapM (internaliseExp arg_desc) args
            fst <$> funcall desc qfname args' loc
 
-internaliseExp desc (E.LetPat tparams pat e body _ loc) =
-  internalisePat desc tparams pat e body loc (internaliseExp desc)
+internaliseExp desc (E.LetPat pat e body _ loc) =
+  internalisePat desc pat e body loc (internaliseExp desc)
 
 internaliseExp desc (E.LetFun ofname (tparams, params, retdecl, Info rettype, body) letbody loc) = do
   internaliseValBind $ E.ValBind False ofname retdecl (Info rettype) tparams params body Nothing loc
   internaliseExp desc letbody
 
-internaliseExp desc (E.DoLoop tparams mergepat mergeexp form loopbody loc) = do
+internaliseExp desc (E.DoLoop mergepat mergeexp form loopbody loc) = do
   -- We pretend that we saw a let-binding first to ensure that the
   -- initial values for the merge parameters match their annotated
   -- sizes
   ses <- internaliseExp "loop_init" mergeexp
   t <- I.staticShapes <$> mapM I.subExpType ses
-  stmPattern tparams mergepat t $ \cm mergepat_names match -> do
+  stmPattern mergepat t $ \cm mergepat_names match -> do
     mapM_ (uncurry (internaliseDimConstant loc)) cm
     ses' <- match (srclocOf mergepat) ses
     forM_ (zip mergepat_names ses') $ \(v,se) ->
@@ -505,8 +535,8 @@
 
       i <- newVName "i"
 
-      bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat ->
-        bindingLambdaParams [] [x] (map rowType arr_ts) $ \x_cm x_params -> do
+      bindingParams [] [mergepat] $ \mergecm shapepat nested_mergepat ->
+        bindingLambdaParams [x] (map rowType arr_ts) $ \x_cm x_params -> do
           mapM_ (uncurry (internaliseDimConstant loc)) x_cm
           mapM_ (uncurry (internaliseDimConstant loc)) mergecm
           let loopvars = zip x_params arr'
@@ -520,12 +550,12 @@
               I.Prim (IntType it) -> return it
               _                   -> fail "internaliseExp DoLoop: invalid type"
 
-      bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat -> do
+      bindingParams [] [mergepat] $ \mergecm shapepat nested_mergepat -> do
         mapM_ (uncurry (internaliseDimConstant loc)) mergecm
         forLoop nested_mergepat shapepat mergeinit $ I.ForLoop i' it num_iterations' []
 
     handleForm mergeinit (E.While cond) =
-      bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat -> do
+      bindingParams [] [mergepat] $ \mergecm shapepat nested_mergepat -> do
         mergeinit_ts <- mapM subExpType mergeinit
         mapM_ (uncurry (internaliseDimConstant loc)) mergecm
         let mergepat' = concat nested_mergepat
@@ -584,7 +614,7 @@
   let pat = E.Id (E.identName name) (E.identType name) loc
       src_t = E.fromStruct <$> E.identType src
       e = E.Update (E.Var (E.qualName $ E.identName src) src_t loc) idxs ve loc
-  internaliseExp desc $ E.LetPat [] pat e body t loc
+  internaliseExp desc $ E.LetPat pat e body t loc
 
 internaliseExp desc (E.Update src slice ve loc) = do
   ves <- internaliseExp "lw_val" ve
@@ -649,7 +679,7 @@
       bFalse <- bFalseM
       letTupExp' desc =<< generateCaseIf desc e c bFalse
       where bFalseM = do
-              eLast' <- internalisePat desc [] pLast e eLast locLast internaliseBody
+              eLast' <- internalisePat desc pLast e eLast locLast internaliseBody
               foldM (\bf c' -> eBody $ return $ generateCaseIf desc e c' bf) eLast' (reverse $ init cs')
             CasePat pLast eLast locLast = last cs'
     [] -> fail $ "internaliseExp: match with no cases at: " ++ locStr loc
@@ -757,21 +787,24 @@
 
 generateCaseIf :: String -> E.Exp -> Case -> I.Body -> InternaliseM I.Exp
 generateCaseIf desc e (CasePat p eCase loc) bFail = do
-  eCase' <- internalisePat desc [] p e eCase loc internaliseBody
+  eCase' <- internalisePat desc p e eCase loc internaliseBody
   eIf cond (return eCase') (return bFail)
   where cond = BasicOp . SubExp <$> internaliseExp1 "cond" (generateCond p e)
 
-internalisePat :: String -> [TypeParamBase VName] -> E.Pattern -> E.Exp
+internalisePat :: String -> E.Pattern -> E.Exp
                -> E.Exp -> SrcLoc -> (E.Exp -> InternaliseM a) -> InternaliseM a
-internalisePat desc tparams p e body loc m = do
-  ses <- internaliseExp desc e
+internalisePat desc p e body loc m = do
+  ses <- internaliseExp desc' e
   t <- I.staticShapes <$> mapM I.subExpType ses
-  stmPattern tparams p t $ \cm pat_names match -> do
+  stmPattern p t $ \cm pat_names match -> do
     mapM_ (uncurry (internaliseDimConstant loc)) cm
     ses' <- match loc ses
     forM_ (zip pat_names ses') $ \(v,se) ->
       letBindNames_ [v] $ I.BasicOp $ I.SubExp se
     m body
+  where desc' = case S.toList $ E.patternIdents p of
+                  [v] -> baseString $ E.identName v
+                  _ -> desc
 
 internaliseSlice :: SrcLoc
                  -> [SubExp]
@@ -1174,8 +1207,8 @@
 internaliseLambda (E.Parens e _) rowtypes =
   internaliseLambda e rowtypes
 
-internaliseLambda (E.Lambda tparams params body _ (Info (_, rettype)) loc) rowtypes =
-  bindingLambdaParams tparams params rowtypes $ \pcm params' -> do
+internaliseLambda (E.Lambda params body _ (Info (_, rettype)) loc) rowtypes =
+  bindingLambdaParams params rowtypes $ \pcm params' -> do
     (rettype', rcm) <- internaliseReturnType rettype
     body' <- internaliseBody body
     mapM_ (uncurry (internaliseDimConstant loc)) $ pcm<>rcm
@@ -1283,7 +1316,7 @@
 
                       -- Check that all were equal.
                       and_lam <- binOpLambda I.LogAnd I.Bool
-                      reduce <- I.reduceSOAC Commutative and_lam [constant True]
+                      reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]]
                       all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems reduce [cmps]
                       return $ resultBody [all_equal]
 
@@ -1394,13 +1427,13 @@
       internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc)
       where reduce w red_lam nes arrs =
               I.Screma w <$>
-              I.reduceSOAC Noncommutative red_lam nes <*> pure arrs
+              I.reduceSOAC [Reduce Noncommutative red_lam nes] <*> pure arrs
 
     handle [TupLit [lam, ne, arr] _] "reduce_comm" = Just $ \desc ->
       internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc)
       where reduce w red_lam nes arrs =
               I.Screma w <$>
-              I.reduceSOAC Commutative red_lam nes <*> pure arrs
+              I.reduceSOAC [Reduce Commutative red_lam nes] <*> pure arrs
 
     handle [TupLit [lam, ne, arr] _] "scan" = Just $ \desc ->
       internaliseScanOrReduce desc "scan" reduce (lam, ne, arr, loc)
diff --git a/src/Futhark/Internalise/AccurateSizes.hs b/src/Futhark/Internalise/AccurateSizes.hs
--- a/src/Futhark/Internalise/AccurateSizes.hs
+++ b/src/Futhark/Internalise/AccurateSizes.hs
@@ -52,7 +52,7 @@
             Just s | x:xs <- S.toList s ->
                        if Var name `elem` outer_dims
                        then return x
-                       else letSubExp "size" =<< foldBinOp (SMax Int32) x xs
+                       else letSubExp (baseString name) =<< foldBinOp (SMax Int32) x xs
             _ -> return $ intConst Int32 0
 
 ensureResultShape :: MonadBinder m =>
diff --git a/src/Futhark/Internalise/Bindings.hs b/src/Futhark/Internalise/Bindings.hs
--- a/src/Futhark/Internalise/Bindings.hs
+++ b/src/Futhark/Internalise/Bindings.hs
@@ -53,15 +53,14 @@
     I.localScope (I.scopeOfFParams $ shape_params++concat valueparams) $
     substitutingVars shape_subst $ m cm shape_params $ chunks num_param_ts (concat valueparams)
 
-bindingLambdaParams :: [E.TypeParam] -> [E.Pattern] -> [I.Type]
+bindingLambdaParams :: [E.Pattern] -> [I.Type]
                     -> (ConstParams -> [I.LParam] -> InternaliseM a)
                     -> InternaliseM a
-bindingLambdaParams tparams params ts m = do
+bindingLambdaParams params ts m = do
   (params_idents, params_types) <-
     unzip . concat <$> mapM flattenPattern params
-  let bound = boundInTypes tparams
-      param_names = M.fromList [ (E.identName x, y) | (x,y) <- params_idents ]
-  (params_ts, cm) <- internaliseParamTypes bound param_names params_types
+  let param_names = M.fromList [ (E.identName x, y) | (x,y) <- params_idents ]
+  (params_ts, cm) <- internaliseParamTypes mempty param_names params_types
 
   let ascript_substs = lambdaShapeSubstitutions (concat params_ts) ts
 
@@ -143,24 +142,23 @@
 
 type MatchPattern = SrcLoc -> [I.SubExp] -> InternaliseM [I.SubExp]
 
-stmPattern :: [E.TypeParam] -> E.Pattern -> [I.ExtType]
+stmPattern :: E.Pattern -> [I.ExtType]
            -> (ConstParams -> [VName] -> MatchPattern -> InternaliseM a)
            -> InternaliseM a
-stmPattern tparams pat ts m = do
+stmPattern pat ts m = do
   (pat', pat_types) <- unzip <$> flattenPattern pat
   (ts',_) <- instantiateShapes' ts
-  (pat_types', cm) <- internaliseParamTypes (boundInTypes tparams) mempty pat_types
+  (pat_types', cm) <- internaliseParamTypes mempty mempty pat_types
   let pat_types'' = map I.fromDecl $ concat pat_types'
-      tparam_names = S.fromList $ map E.typeParamName tparams
   let addShapeStms l =
-        m cm (map I.paramName $ concat l) (matchPattern tparam_names pat_types'')
+        m cm (map I.paramName $ concat l) (matchPattern pat_types'')
   bindingFlatPattern pat' ts' addShapeStms
 
-matchPattern :: S.Set VName -> [I.ExtType] -> MatchPattern
-matchPattern tparam_names exts loc ses =
+matchPattern :: [I.ExtType] -> MatchPattern
+matchPattern exts loc ses =
   forM (zip exts ses) $ \(et, se) -> do
   se_t <- I.subExpType se
-  et' <- unExistentialise tparam_names et se_t
+  et' <- unExistentialise mempty et se_t
   ensureExtShape asserting (I.ErrorMsg [I.ErrorString "value cannot match pattern"])
     loc et' "correct_shape" se
 
diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs
--- a/src/Futhark/Internalise/Defunctionalise.hs
+++ b/src/Futhark/Internalise/Defunctionalise.hs
@@ -17,10 +17,16 @@
 import           Language.Futhark
 import           Futhark.Representation.AST.Pretty ()
 
+-- | An expression or an extended 'Lambda' (with size parameters,
+-- which AST lambdas do not support).
+data ExtExp = ExtLambda [TypeParam] [Pattern] Exp (Aliasing, StructType) SrcLoc
+            | ExtExp Exp
+  deriving (Show)
+
 -- | A static value stores additional information about the result of
 -- defunctionalization of an expression, aside from the residual expression.
 data StaticVal = Dynamic PatternType
-               | LambdaSV [VName] Pattern StructType Exp Env
+               | LambdaSV [VName] Pattern StructType ExtExp Env
                  -- ^ The 'VName's are shape parameters that are bound
                  -- by the 'Pattern'.
                | RecordSV [(Name, StaticVal)]
@@ -101,6 +107,46 @@
       | otherwise -> error $ "Variable " ++ pretty x ++ " at "
                           ++ locStr loc ++ " is out of scope."
 
+defuncFun :: [TypeParam] -> [Pattern] -> Exp -> (Aliasing, StructType) -> SrcLoc
+          -> DefM (Exp, StaticVal)
+defuncFun tparams pats e0 (closure, ret) loc = do
+  when (any isTypeParam tparams) $
+    error $ "Received a lambda with type parameters at " ++ locStr loc
+         ++ ", but the defunctionalizer expects a monomorphic input program."
+  -- Extract the first parameter of the lambda and "push" the
+  -- remaining ones (if there are any) into the body of the lambda.
+  let (dims, pat, ret', e0') = case pats of
+        [] -> error "Received a lambda with no parameters."
+        [pat'] -> (map typeParamName tparams, pat', ret, ExtExp e0)
+        (pat' : pats') ->
+          -- Split shape parameters into those that are determined by
+          -- the first pattern, and those that are determined by later
+          -- patterns.
+          let bound_by_pat = (`S.member` patternDimNames pat') . typeParamName
+              (pat_dims, rest_dims) = partition bound_by_pat tparams
+          in (map typeParamName pat_dims, pat',
+              foldFunType (map (toStruct . patternPatternType) pats') ret,
+              ExtLambda rest_dims pats' e0 (closure, ret) loc)
+
+  -- Construct a record literal that closes over the environment of
+  -- the lambda.  Closed-over 'DynamicFun's are converted to their
+  -- closure representation.
+  env <- restrictEnvTo $
+         freeVars (Lambda pats e0 Nothing (Info (closure, ret)) loc) `without`
+         mconcat (map (oneName . typeParamName) tparams)
+  let (fields, env') = unzip $ map closureFromDynamicFun $ M.toList env
+  return (RecordLit fields loc, LambdaSV dims pat ret' e0' $ M.fromList env')
+
+  where closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) =
+          let name = nameFromString $ pretty vn
+          in (RecordFieldExplicit name clsr_env noLoc, (vn, sv))
+
+        closureFromDynamicFun (vn, sv) =
+          let name = nameFromString $ pretty vn
+              tp' = typeFromSV sv
+          in (RecordFieldExplicit name
+               (Var (qualName vn) (Info tp') noLoc) noLoc, (vn, sv))
+
 -- | Defunctionalization of an expression. Returns the residual expression and
 -- the associated static value in the defunctionalization monad.
 defuncExp :: Exp -> DefM (Exp, StaticVal)
@@ -169,7 +215,7 @@
     -- can get rid of them.
     IntrinsicSV -> do
       (pats, body, tp) <- etaExpand e
-      defuncExp $ Lambda [] pats body Nothing (Info (mempty, tp)) noLoc
+      defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) noLoc
     _ -> let tp = typeFromSV sv
          in return (Var qn (Info tp) loc, sv)
 
@@ -178,23 +224,20 @@
                                return (Ascript e0' tydecl t loc, sv)
   | otherwise = defuncExp e0
 
-defuncExp (LetPat tparams pat e1 e2 _ loc) = do
-  let env_dim = envFromShapeParams tparams
-  (e1', sv1) <- localEnv env_dim $ defuncExp e1
+defuncExp (LetPat pat e1 e2 _ loc) = do
+  (e1', sv1) <- defuncExp e1
   let env  = matchPatternSV pat sv1
       pat' = updatePattern pat sv1
-  (e2', sv2) <- localEnv (env <> env_dim) $ defuncExp e2
-  return (LetPat tparams pat' e1' e2' (Info $ typeOf e2') loc, sv2)
+  (e2', sv2) <- localEnv env $ defuncExp e2
+  return (LetPat pat' e1' e2' (Info $ typeOf e2') loc, sv2)
 
-defuncExp (LetFun vn (dims, pats, _, rettype@(Info ret), e1) e2 loc) = do
-  let env_dim = envFromShapeParams dims
-  (pats', e1', sv1) <- localEnv env_dim $ defuncLet dims pats e1 rettype
-  (e2', sv2) <- extendEnv vn sv1 $ defuncExp e2
-  case pats' of
-    []  -> let t1 = combineTypeShapes (fromStruct ret) $ typeOf e1'
-           in return (LetPat dims (Id vn (Info t1) noLoc) e1' e2' (Info $ typeOf e2') loc, sv2)
-    _:_ -> let t1 = combineTypeShapes ret $ toStruct $ typeOf e1'
-           in return (LetFun vn (dims, pats', Nothing, Info t1, e1') e2' loc, sv2)
+-- Local functions are handled by rewriting them to lambdas, so that
+-- the same machinery can be re-used.
+defuncExp (LetFun vn (dims, pats, _, Info ret, e1) e2 loc) = do
+  (e1', sv1) <- defuncFun dims pats e1 (mempty, ret) loc
+  (e2', sv2) <- localEnv (M.singleton vn sv1) $ defuncExp e2
+  return (LetPat (Id vn (Info (typeOf e1')) loc) e1' e2' (Info $ typeOf e2') loc,
+          sv2)
 
 defuncExp (If e1 e2 e3 tp loc) = do
   (e1', _ ) <- defuncExp e1
@@ -216,41 +259,8 @@
   (e0', sv) <- defuncExp e0
   return (Negate e0' loc, sv)
 
-defuncExp e@(Lambda tparams pats e0 decl (Info (closure, ret)) loc) = do
-  when (any isTypeParam tparams) $
-    error $ "Received a lambda with type parameters at " ++ locStr loc
-         ++ ", but the defunctionalizer expects a monomorphic input program."
-  -- Extract the first parameter of the lambda and "push" the
-  -- remaining ones (if there are any) into the body of the lambda.
-  let (dims, pat, ret', e0') = case pats of
-        [] -> error "Received a lambda with no parameters."
-        [pat'] -> (map typeParamName tparams, pat', ret, e0)
-        (pat' : pats') ->
-          -- Split shape parameters into those that are determined by
-          -- the first pattern, and those that are determined by later
-          -- patterns.
-          let bound_by_pat = (`S.member` patternDimNames pat') . typeParamName
-              (pat_dims, rest_dims) = partition bound_by_pat tparams
-          in (map typeParamName pat_dims, pat',
-              foldFunType (map (toStruct . patternPatternType) pats') ret,
-              Lambda rest_dims pats' e0 decl (Info (closure, ret)) loc)
-
-  -- Construct a record literal that closes over the environment of
-  -- the lambda.  Closed-over 'DynamicFun's are converted to their
-  -- closure representation.
-  env <- restrictEnvTo (freeVars e)
-  let (fields, env') = unzip $ map closureFromDynamicFun $ M.toList env
-  return (RecordLit fields loc, LambdaSV dims pat ret' e0' $ M.fromList env')
-
-  where closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) =
-          let name = nameFromString $ pretty vn
-          in (RecordFieldExplicit name clsr_env noLoc, (vn, sv))
-
-        closureFromDynamicFun (vn, sv) =
-          let name = nameFromString $ pretty vn
-              tp' = typeFromSV sv
-          in (RecordFieldExplicit name
-               (Var (qualName vn) (Info tp') noLoc) noLoc, (vn, sv))
+defuncExp (Lambda pats e0 _ (Info (closure, ret)) loc) =
+  defuncFun [] pats e0 (closure, ret) loc
 
 -- Operator sections are expected to be converted to lambda-expressions
 -- by the monomorphizer, so they should no longer occur at this point.
@@ -260,8 +270,7 @@
 defuncExp ProjectSection{} = error "defuncExp: unexpected projection section."
 defuncExp IndexSection{}   = error "defuncExp: unexpected projection section."
 
-defuncExp (DoLoop tparams pat e1 form e3 loc) = do
-  let env_dim = envFromShapeParams tparams
+defuncExp (DoLoop pat e1 form e3 loc) = do
   (e1', sv1) <- defuncExp e1
   let env1 = matchPatternSV pat sv1
   (form', env2) <- case form of
@@ -269,10 +278,10 @@
                         return (For v e2', envFromIdent v)
     ForIn pat2 e2 -> do e2' <- defuncExp' e2
                         return (ForIn pat2 e2', envFromPattern pat2)
-    While e2      -> do e2' <- localEnv (env1 <> env_dim) $ defuncExp' e2
+    While e2      -> do e2' <- localEnv env1 $ defuncExp' e2
                         return (While e2', mempty)
-  (e3', sv) <- localEnv (env1 <> env2 <> env_dim) $ defuncExp e3
-  return (DoLoop tparams pat e1' form' e3' loc, sv)
+  (e3', sv) <- localEnv (env1 <> env2) $ defuncExp e3
+  return (DoLoop pat e1' form' e3' loc, sv)
   where envFromIdent (Ident vn (Info tp) _) =
           M.singleton vn $ Dynamic tp
 
@@ -351,6 +360,11 @@
 defuncExp' :: Exp -> DefM Exp
 defuncExp' = fmap fst . defuncExp
 
+defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)
+defuncExtExp (ExtExp e) = defuncExp e
+defuncExtExp (ExtLambda tparams pats e0 (closure, ret) loc) =
+  defuncFun tparams pats e0 (closure, ret) loc
+
 defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal)
 defuncCase sv (CasePat p e loc) = do
   let p'  = updatePattern p sv
@@ -369,18 +383,17 @@
 defuncSoacExp (Parens e loc) =
   Parens <$> defuncSoacExp e <*> pure loc
 
-defuncSoacExp (Lambda tparams params e0 decl tp loc) = do
-  let env_dim = envFromShapeParams tparams
-      env = foldMap envFromPattern params
-  e0' <- localEnv (env <> env_dim) $ defuncSoacExp e0
-  return $ Lambda tparams params e0' decl tp loc
+defuncSoacExp (Lambda params e0 decl tp loc) = do
+  let env = foldMap envFromPattern params
+  e0' <- localEnv env $ defuncSoacExp e0
+  return $ Lambda params e0' decl tp loc
 
 defuncSoacExp e
   | Arrow{} <- typeOf e = do
       (pats, body, tp) <- etaExpand e
       let env = foldMap envFromPattern pats
       body' <- localEnv env $ defuncExp' body
-      return $ Lambda [] pats body' Nothing (Info (mempty, tp)) noLoc
+      return $ Lambda pats body' Nothing (Info (mempty, tp)) noLoc
   | otherwise = defuncExp' e
 
 etaExpand :: Exp -> DefM ([Pattern], Exp, StructType)
@@ -417,10 +430,10 @@
           bound_by_pat = (`S.member` patternDimNames pat) . typeParamName
           (_pat_dims, rest_dims) = partition bound_by_pat dims
       (pats', body', sv) <- localEnv env $ defuncLet rest_dims pats body (Info rettype)
-      closure <- defuncExp $ Lambda dims ps body Nothing (Info (mempty, rettype)) noLoc
+      closure <- defuncFun dims ps body (mempty, rettype) noLoc
       return (pat : pats', body', DynamicFun closure sv)
   | otherwise = do
-      (e, sv) <- defuncExp $ Lambda dims ps body Nothing (Info (mempty, rettype)) noLoc
+      (e, sv) <- defuncFun dims ps body (mempty, rettype) noLoc
       return ([], e, sv)
 defuncLet _ [] body (Info rettype) = do
   (body', sv) <- defuncExp body
@@ -445,7 +458,7 @@
     LambdaSV dims pat e0_t e0 closure_env -> do
       let env' = matchPatternSV pat sv2
           env_dim = envFromDimNames dims
-      (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExp e0
+      (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExtExp e0
 
       let closure_pat = buildEnvPattern closure_env
           pat' = updatePattern pat sv2
@@ -742,8 +755,8 @@
                           foldMap freeVars incl
   Var qn (Info t) _    -> NameSet $ M.singleton (qualLeaf qn) $ uniqueness t
   Ascript e t _ _      -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t)
-  LetPat _ pat e1 e2 _ _ -> freeVars e1 <> ((names (patternDimNames pat) <> freeVars e2)
-                                            `without` patternVars pat)
+  LetPat pat e1 e2 _ _ -> freeVars e1 <> ((names (patternDimNames pat) <> freeVars e2)
+                                          `without` patternVars pat)
 
   LetFun vn (_, pats, _, _, e1) e2 _ ->
     ((freeVars e1 <> names (foldMap patternDimNames pats))
@@ -753,18 +766,17 @@
   If e1 e2 e3 _ _           -> freeVars e1 <> freeVars e2 <> freeVars e3
   Apply e1 e2 _ _ _         -> freeVars e1 <> freeVars e2
   Negate e _                -> freeVars e
-  Lambda tps pats e0 _ _ _  -> (names (foldMap patternDimNames pats) <> freeVars e0)
-                               `without` (foldMap patternVars pats <>
-                                          mconcat (map (oneName . typeParamName) tps))
+  Lambda pats e0 _ _ _      -> (names (foldMap patternDimNames pats) <> freeVars e0)
+                               `without` foldMap patternVars pats
   OpSection{}                 -> mempty
   OpSectionLeft _  _ e _ _ _  -> freeVars e
   OpSectionRight _ _ e _ _ _  -> freeVars e
   ProjectSection{}            -> mempty
   IndexSection idxs _ _       -> foldMap freeDimIndex idxs
 
-  DoLoop _ pat e1 form e3 _ -> let (e2fv, e2ident) = formVars form
-                               in freeVars e1 <> e2fv <>
-                               (freeVars e3 `without` (patternVars pat <> e2ident))
+  DoLoop pat e1 form e3 _ -> let (e2fv, e2ident) = formVars form
+                             in freeVars e1 <> e2fv <>
+                             (freeVars e3 `without` (patternVars pat <> e2ident))
     where formVars (For v e2) = (freeVars e2, ident v)
           formVars (ForIn p e2)   = (freeVars e2, patternVars p)
           formVars (While e2)     = (freeVars e2, mempty)
diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs
--- a/src/Futhark/Internalise/Lambdas.hs
+++ b/src/Futhark/Internalise/Lambdas.hs
@@ -58,15 +58,20 @@
       outer = (`setOuterSize` I.Var chunk_size)
   localScope (scopeOfLParams [chunk_param]) $ do
     argtypes <- mapM I.subExpType args
-    (params, body, rettype) <- internaliseLambda lam $ map outer argtypes
+    (orig_chunk_param : params, orig_body, rettype) <-
+      internaliseLambda lam $ I.Prim int32 : map outer argtypes
+    body <- runBodyBinder $ do
+      letBindNames_ [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size
+      return orig_body
     (rettype', inner_shapes) <- instantiateShapes' rettype
     let outer_shape = arraysSize 0 argtypes
     shapefun <- makeShapeFun (chunk_param:params) body rettype' inner_shapes
     bindMapShapes (slice0 chunk_size) [zero] inner_shapes shapefun args outer_shape
-    body' <- localScope (scopeOfLParams params) $
-             ensureResultShape asserting
-             (ErrorMsg [ErrorString "not all iterations produce same shape"])
-             (srclocOf lam) (map outer rettype') body
+    body' <- localScope (scopeOfLParams params) $ insertStmsM $ do
+      letBindNames_ [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size
+      ensureResultShape asserting
+        (ErrorMsg [ErrorString "not all iterations produce same shape"])
+        (srclocOf lam) (map outer rettype') body
     return $ I.Lambda (chunk_param:params) body' (map outer rettype')
   where slice0 chunk_size arg = do
           arg' <- letExp "arg" $ I.BasicOp $ I.SubExp arg
@@ -102,7 +107,7 @@
       let sizefun_safe =
             all (I.safeExp . I.stmExp) $ I.bodyStms $ I.lambdaBody sizefun'
           sizefun_arg_invariant =
-            not $ any (`S.member` freeInBody (I.lambdaBody sizefun')) $
+            not $ any (`S.member` freeIn (I.lambdaBody sizefun')) $
             map I.paramName $ lambdaParams sizefun'
       if sizefun_safe && sizefun_arg_invariant
         then do ses <- bodyBind $ lambdaBody sizefun'
@@ -147,9 +152,13 @@
   chunk_size <- newVName "chunk_size"
   let chunk_param = I.Param chunk_size $ I.Prim int32
       chunktypes = map (`arrayOfRow` I.Var chunk_size) rowts
-  (params, body, _) <- localScope (scopeOfLParams [chunk_param]) $
-                       internaliseLambda lam chunktypes
-  return (chunk_param:params, body)
+  localScope (scopeOfLParams [chunk_param]) $ do
+    (orig_chunk_param : params, orig_body, _) <-
+      internaliseLambda lam $ I.Prim int32 : chunktypes
+    body <- runBodyBinder $ do
+      letBindNames_ [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size
+      return orig_body
+    return (chunk_param:params, body)
 
 -- Given @k@ lambdas, this will return a lambda that returns an
 -- (k+2)-element tuple of integers.  The first element is the
diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs
--- a/src/Futhark/Internalise/Monomorphise.hs
+++ b/src/Futhark/Internalise/Monomorphise.hs
@@ -208,10 +208,10 @@
 transformExp (Ascript e tp t loc) =
   Ascript <$> transformExp e <*> pure tp <*> pure t <*> pure loc
 
-transformExp (LetPat tparams pat e1 e2 (Info t) loc) = do
+transformExp (LetPat pat e1 e2 (Info t) loc) = do
   (pat', rr) <- expandRecordPattern pat
   t' <- transformType t
-  LetPat tparams pat' <$> transformExp e1 <*>
+  LetPat pat' <$> transformExp e1 <*>
     withRecordReplacements rr (transformExp e2) <*>
     pure (Info t') <*> pure loc
 
@@ -227,10 +227,10 @@
         let (bs_local, bs_prop) = Seq.partition ((== fname) . fst) bs
         return (unfoldLetFuns (map snd $ toList bs_local) e', const bs_prop)
 
-  | otherwise =
-      transformExp $ LetPat [] (Id fname (Info ft) loc) lam e (Info $ fromStruct ret) loc
-        where lam = Lambda tparams params body Nothing (Info (mempty, ret)) loc
-              ft = foldFunType (map patternType params) $ fromStruct ret
+  | otherwise = do
+      body' <- transformExp body
+      LetFun fname (tparams, params, retdecl, Info ret, body') <$>
+        transformExp e <*> pure loc
 
 transformExp (If e1 e2 e3 tp loc) = do
   e1' <- transformExp e1
@@ -248,9 +248,9 @@
 transformExp (Negate e loc) =
   Negate <$> transformExp e <*> pure loc
 
-transformExp (Lambda tparams params e0 decl tp loc) = do
+transformExp (Lambda params e0 decl tp loc) = do
   e0' <- transformExp e0
-  return $ Lambda tparams params e0' decl tp loc
+  return $ Lambda params e0' decl tp loc
 
 transformExp (OpSection qn t loc) =
   transformExp $ Var qn t loc
@@ -273,14 +273,14 @@
 transformExp (IndexSection idxs (Info t) loc) =
   desugarIndexSection idxs t loc
 
-transformExp (DoLoop tparams pat e1 form e3 loc) = do
+transformExp (DoLoop pat e1 form e3 loc) = do
   e1' <- transformExp e1
   form' <- case form of
     For ident e2  -> For ident <$> transformExp e2
     ForIn pat2 e2 -> ForIn pat2 <$> transformExp e2
     While e2      -> While <$> transformExp e2
   e3' <- transformExp e3
-  return $ DoLoop tparams pat e1' form' e3' loc
+  return $ DoLoop pat e1' form' e3' loc
 
 transformExp (BinOp (QualName qs fname) (Info t) (e1, d1) (e2, d2) tp loc) = do
   fname' <- transformFName fname (toStructural t)
@@ -346,7 +346,7 @@
   (e2, p2) <- makeVarParam e_right $ fromStruct ytype
   let body = BinOp qn (Info t) (e1, Info xtype) (e2, Info ytype) (Info rettype) loc
       rettype' = toStruct rettype
-  return $ Lambda [] (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc
+  return $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc
 
   where makeVarParam (Just e) _ = return (e, [])
         makeVarParam Nothing argtype = do
@@ -358,7 +358,7 @@
 desugarProjectSection fields (Arrow _ _ t1 t2) loc = do
   p <- newVName "project_p"
   let body = foldl project (Var (qualName p) (Info t1) noLoc) fields
-  return $ Lambda [] [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc
+  return $ Lambda [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc
   where project e field =
           case typeOf e of
             Record fs | Just t <- M.lookup field fs ->
@@ -371,7 +371,7 @@
 desugarIndexSection idxs (Arrow _ _ t1 t2) loc = do
   p <- newVName "index_i"
   let body = Index (Var (qualName p) (Info t1) loc) idxs (Info t2) loc
-  return $ Lambda [] [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc
+  return $ Lambda [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc
 desugarIndexSection  _ t _ = error $ "desugarIndexSection: not a function type: " ++ pretty t
 
 noticeDims :: TypeBase (DimDecl VName) as -> MonoM ()
diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs
--- a/src/Futhark/Optimise/DoubleBuffer.hs
+++ b/src/Futhark/Optimise/DoubleBuffer.hs
@@ -33,12 +33,15 @@
 import           Data.Maybe
 import           Data.List
 
-import           Futhark.MonadFreshNames
+import           Futhark.Construct
 import           Futhark.Representation.AST
+import           Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp)
+import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
 import           Futhark.Representation.ExplicitMemory
                  hiding (Prog, Body, Stm, Pattern, PatElem,
                          BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType)
 import           Futhark.Pass
+import           Futhark.Util (maybeHead)
 
 doubleBuffer :: Pass ExplicitMemory ExplicitMemory
 doubleBuffer =
@@ -98,7 +101,8 @@
 -- | Bunch up all the constraints for less typing.
 type LoreConstraints lore inner =
   (ExpAttr lore ~ (), BodyAttr lore ~ (),
-   ExplicitMemorish lore, Op lore ~ MemOp inner)
+   ExplicitMemorish lore, Op lore ~ MemOp inner,
+   BinderOps lore)
 
 optimiseBody :: LoreConstraints lore inner =>
                 Body lore -> DoubleBufferM lore (Body lore)
@@ -130,7 +134,7 @@
                                       -- necessary to prevent the GHC
                                       -- 8.4 type checker from going
                                       -- nuts.
-                                      (optimiseBody x :: DoubleBufferM lore (Body lore))
+                                      optimiseBody x :: DoubleBufferM lore (Body lore)
                                   , mapOnOp = optimiseOp
                                   }
 
@@ -182,7 +186,7 @@
 
 -- | The booleans indicate whether we should also play with the
 -- initial merge values.
-data DoubleBuffer lore = BufferAlloc VName SubExp Space Bool
+data DoubleBuffer lore = BufferAlloc VName (PrimExp VName) Space Bool
                        | BufferCopy VName IxFun VName Bool
                        -- ^ First name is the memory block to copy to,
                        -- second is the name of the array copy.
@@ -211,13 +215,24 @@
             Nothing ->
               Just (Var v, True)
 
+        sizeForMem mem = maybeHead $ mapMaybe (arrayInMem . paramAttr) val_params
+          where arrayInMem (MemArray pt shape _ (ArrayIn arraymem ixfun))
+                  | IxFun.isDirect ixfun,
+                    Just (dims, b) <-
+                      mapAndUnzipM loopInvariantSize $ shapeDims shape,
+                    mem == arraymem =
+                      Just (arraySizeInBytesExp $
+                             Array pt (Shape dims) NoUniqueness,
+                            or b)
+                arrayInMem _ = Nothing
+
         buffer fparam = case paramType fparam of
-          Mem size space
-            | Just (size', b) <- loopInvariantSize size -> do
+          Mem space
+            | Just (size, b) <- sizeForMem $ paramName fparam -> do
                 -- Let us double buffer this!
                 bufname <- lift $ newVName "double_buffer_mem"
                 modify $ M.insert (paramName fparam) (bufname, b)
-                return $ BufferAlloc bufname size' space b
+                return $ BufferAlloc bufname size space b
           Array {}
             | MemArray _ _ _ (ArrayIn mem ixfun) <- paramAttr fparam -> do
                 buffered <- gets $ M.lookup mem
@@ -234,9 +249,11 @@
           -> DoubleBufferM lore ([(FParam lore, SubExp)], [Stm lore])
 allocStms merge = runWriterT . zipWithM allocation merge
   where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do
-          tell [Let (Pattern [] [PatElem name $ MemMem size space]) (defAux ()) $
-                Op $ Alloc size space]
-          if b then return (Param pname $ MemMem size space, Var name)
+          stms <- lift $ runBinder_ $ do
+            size' <- letSubExp "double_buffer_size" =<< toExp size
+            letBindNames_ [name] $ Op $ Alloc size' space
+          tell $ stmsToList stms
+          if b then return (Param pname $ MemMem space, Var name)
                else return m
         allocation (f, Var v) (BufferCopy mem _ _ b) | b = do
           v_copy <- lift $ newVName $ baseString v ++ "_double_buffer_copy"
diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs
--- a/src/Futhark/Optimise/Fusion.hs
+++ b/src/Futhark/Optimise/Fusion.hs
@@ -2,7 +2,6 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE ViewPatterns #-}
 -- | Perform horizontal and vertical fusion of SOACs.
 module Futhark.Optimise.Fusion ( fuseSOACs )
   where
@@ -438,8 +437,8 @@
       infusible_nms  = S.fromList $ filter (`S.member` infusible res) out_nms
       out_arr_nms    = case soac of
                         -- the accumulator result cannot be fused!
-                        SOAC.Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) _) _ ->
-                          drop (length scan_nes + length red_nes) out_nms
+                        SOAC.Screma _ (ScremaForm (_, scan_nes) reds _) _ ->
+                          drop (length scan_nes + redResults reds) out_nms
                         SOAC.Stream _ frm _ _ -> drop (length $ getStreamAccums frm) out_nms
                         _ -> out_nms
       to_fuse_knms1  = S.toList $ getKersWithInpArrs res (out_arr_nms++inp_nms)
@@ -451,28 +450,23 @@
                                        ++ "kernel name not found in kernels field!")
                           Just ker -> return ker
 
-  -- for each kernel get the index in the bindings where the kernel is located
-  -- and sort based on the index so that partial fusion may succeed.
+  -- For each kernel get the index in the bindings where the kernel is
+  -- located and sort based on the index so that partial fusion may
+  -- succeed.  We use the last position where one of the kernel
+  -- outputs occur.
   let bnd_nms = map (patternNames . stmPattern) rem_bnds
   kernminds <- forM to_fuse_knms $ \ker_nm -> do
     ker <- lookupKernel ker_nm
-    let out_nm  = case fsoac ker of
-                    SOAC.Stream _ frm _ _
-                      | x:_ <- drop (length $ getStreamAccums frm) $ outNames ker ->
-                        x
-                    SOAC.Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) _) _
-                      | x:_ <- drop (length scan_nes + length red_nes) $ outNames ker ->
-                        x
-                    _ -> head $ outNames ker
-    case L.findIndex (elem out_nm) bnd_nms of
-      Nothing -> return Nothing
-      Just i  -> return $ Just (ker,ker_nm,i)
+    case mapMaybe (\out_nm -> L.findIndex (elem out_nm) bnd_nms) (outNames ker) of
+      [] -> return Nothing
+      is -> return $ Just (ker,ker_nm,maximum is)
 
   scope <- askScope
   let kernminds' = L.sortBy (\(_,_,i1) (_,_,i2)->compare i1 i2) $ catMaybes kernminds
       soac_kernel = newKernel cs soac consumed out_nms scope
+
   -- now try to fuse kernels one by one (in a fold); @ok_ind@ is the index of the
-  -- kernel until which fusion succeded, and @fused_ker@ is the resulted kernel.
+  -- kernel until which fusion succeded, and @fused_ker@ is the resulting kernel.
   (_,ok_ind,_,fused_ker,_) <-
       foldM (\(cur_ok,n,prev_ind,cur_ker,ufus_nms) (ker, _ker_nm, bnd_ind) -> do
                 -- check that we still try fusion and that the intermediate
@@ -497,7 +491,7 @@
                                       -- check that consumer's lambda body does not use
                                       -- directly the produced arrays (e.g., see noFusion3.fut).
                                       Right conssoac -> return $ S.null $ S.intersection curker_outset $
-                                                                 freeInBody $ lambdaBody $ SOAC.lambda conssoac
+                                                                 freeIn $ lambdaBody $ SOAC.lambda conssoac
                                       Left _         -> return True
 
                 let interm_bnds_ok = cur_ok && consumer_ok && out_transf_ok && cons_no_out_transf &&
@@ -505,7 +499,7 @@
                                        -- (i) check that the in-between bindings do
                                        --     not use the result of current kernel OR
                                        S.null ( S.intersection curker_outset $
-                                                      freeInExp (stmExp bnd) ) ||
+                                                      freeIn (stmExp bnd) ) ||
                                        --(ii) that the pattern-binding corresponds to
                                        --     the result of the consumer kernel; in the
                                        --     latter case it means it corresponds to a
@@ -556,14 +550,17 @@
 ------------------------------------------------------------------------
 
 fusionGatherBody :: FusedRes -> Body -> FusionGM FusedRes
+fusionGatherBody fres (Body _ stms res) =
+  fusionGatherStms fres (stmsToList stms) res
 
+fusionGatherStms :: FusedRes -> [Stm] -> Result -> FusionGM FusedRes
+
 -- Some forms of do-loops can profitably be considered streamSeqs.  We
 -- are careful to ensure that the generated nested loop cannot itself
 -- be considered a stream, to avoid infinite recursion.
-fusionGatherBody fres (Body blore (stmsToList ->
-                                    Let (Pattern [] pes) bndtp
-                                    (DoLoop [] merge (ForLoop i it w loop_vars) body)
-                                    :bnds) res) | not $ null loop_vars = do
+fusionGatherStms fres (Let (Pattern [] pes) bndtp
+                       (DoLoop [] merge (ForLoop i it w loop_vars) body) : bnds) res
+  | not $ null loop_vars = do
   let (merge_params,merge_init) = unzip merge
       (loop_params,loop_arrs) = unzip loop_vars
   chunk_size <- newVName "chunk_size"
@@ -606,10 +603,10 @@
   discard <- newVName "discard"
   let discard_pe = PatElem discard $ Prim int32
 
-  fusionGatherBody fres $ Body blore
-    (oneStm (Let (Pattern [] (pes<>[discard_pe])) bndtp (Op stream))<>stmsFromList bnds) res
+  fusionGatherStms fres
+    (Let (Pattern [] (pes<>[discard_pe])) bndtp (Op stream) : bnds) res
 
-fusionGatherBody fres (Body _ (stmsToList -> (bnd@(Let pat _ e):bnds)) res) = do
+fusionGatherStms fres (bnd@(Let pat _ e):bnds) res = do
   maybesoac <- SOAC.fromExp e
   case maybesoac of
     Right soac@(SOAC.Scatter _len lam _ivs _as) -> do
@@ -626,10 +623,9 @@
       fres' <- addNamesToInfusible fres $ S.fromList $ patternNames pat
       mapLike fres' soac lam
 
-    Right soac@(SOAC.Screma _ (ScremaForm (scan_lam, scan_nes)
-                                              (_, reduce_lam, reduce_nes)
-                                              map_lam) _) ->
-      reduceLike soac [scan_lam, reduce_lam, map_lam] $ scan_nes <> reduce_nes
+    Right soac@(SOAC.Screma _ (ScremaForm (scan_lam, scan_nes) reds map_lam) _) ->
+      reduceLike soac (map redLambda reds ++ [scan_lam, map_lam]) $
+      scan_nes <> concatMap redNeutral reds
 
     Right soac@(SOAC.Stream _ form lam _) -> do
       -- a redomap does not neccessarily start a new kernel, e.g.,
@@ -644,32 +640,31 @@
 
     _ | [pe] <- patternValueElements pat,
         Just (src,trns) <- SOAC.transformFromExp (stmCerts bnd) e ->
-          bindingTransform pe src trns $ fusionGatherBody fres body
+          bindingTransform pe src trns $ fusionGatherStms fres bnds res
       | otherwise -> do
           let pat_vars = map (BasicOp . SubExp . Var) $ patternNames pat
-          bres <- gatherStmPattern pat e $ fusionGatherBody fres body
+          bres <- gatherStmPattern pat e $ fusionGatherStms fres bnds res
           bres' <- checkForUpdates bres e
           foldM fusionGatherExp bres' (e:pat_vars)
 
-  where body = mkBody (stmsFromList bnds) res
-        cs = stmCerts bnd
+  where cs = stmCerts bnd
         rem_bnds = bnd : bnds
         consumed = consumedInExp $ Alias.analyseExp e
 
         reduceLike soac lambdas nes = do
           (used_lam, lres)  <- foldM fusionGatherLam (S.empty, fres) lambdas
-          bres  <- bindingFamily pat $ fusionGatherBody lres body
+          bres  <- bindingFamily pat $ fusionGatherStms lres bnds res
           bres' <- foldM fusionGatherSubExp bres nes
           consumed' <- varsAliases consumed
           greedyFuse rem_bnds used_lam bres' (pat, cs, soac, consumed')
 
         mapLike fres' soac lambda = do
-          bres  <- bindingFamily pat $ fusionGatherBody fres' body
+          bres  <- bindingFamily pat $ fusionGatherStms fres' bnds res
           (used_lam, blres) <- fusionGatherLam (S.empty, bres) lambda
           consumed' <- varsAliases consumed
           greedyFuse rem_bnds used_lam blres (pat, cs, soac, consumed')
 
-fusionGatherBody fres (Body _ _ res) =
+fusionGatherStms fres [] res =
   foldM fusionGatherExp fres $ map (BasicOp . SubExp) res
 
 fusionGatherExp :: FusedRes -> Exp -> FusionGM FusedRes
@@ -715,8 +710,7 @@
 ---- Generic Traversal         ----
 -----------------------------------
 
-fusionGatherExp fres e =
-  addNamesToInfusible fres $ freeInExp e
+fusionGatherExp fres e = addNamesToInfusible fres $ freeIn e
 
 fusionGatherSubExp :: FusedRes -> SubExp -> FusionGM FusedRes
 fusionGatherSubExp fres (Var idd) = addVarToInfusible fres idd
@@ -828,14 +822,13 @@
 finaliseSOAC :: SOAC.SOAC SOACS -> FusionGM (SOAC.SOAC SOACS)
 finaliseSOAC new_soac =
   case new_soac of
-    SOAC.Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs -> do
+    SOAC.Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs -> do
       scan_lam' <- simplifyAndFuseInLambda scan_lam
-      red_lam' <- simplifyAndFuseInLambda red_lam
+      reds' <- forM reds $ \(Reduce comm red_lam red_nes) -> do
+        red_lam' <- simplifyAndFuseInLambda red_lam
+        return $ Reduce comm red_lam' red_nes
       map_lam' <- simplifyAndFuseInLambda map_lam
-      return $ SOAC.Screma w (ScremaForm (scan_lam', scan_nes)
-                                             (comm, red_lam', red_nes)
-                                             map_lam')
-                               arrs
+      return $ SOAC.Screma w (ScremaForm (scan_lam', scan_nes) reds' map_lam') arrs
     SOAC.Scatter w lam inps dests -> do
       lam' <- simplifyAndFuseInLambda lam
       return $ SOAC.Scatter w lam' inps dests
@@ -859,10 +852,7 @@
                   -> Binder SOACS (Futhark.SOAC SOACS)
 copyNewlyConsumed was_consumed soac =
   case soac of
-    Futhark.Screma w (Futhark.ScremaForm
-                         (scan_lam, scan_nes)
-                         (comm, reduce_lam, reduce_nes)
-                         map_lam) arrs -> do
+    Futhark.Screma w (Futhark.ScremaForm (scan_lam, scan_nes) reds map_lam) arrs -> do
       -- Copy any arrays that are consumed now, but were not in the
       -- constituents.
       arrs' <- mapM copyConsumedArr arrs
@@ -870,11 +860,12 @@
       -- lambda, and we have to substitute the name of the copy for
       -- the original.
       map_lam' <- copyFreeInLambda map_lam
+      let reds' = map (\red -> red { redLambda =
+                                       Aliases.removeLambdaAliases
+                                       (redLambda red)})
+                  reds
       return $ Futhark.Screma w
-        (Futhark.ScremaForm
-         (Aliases.removeLambdaAliases scan_lam, scan_nes)
-         (comm, Aliases.removeLambdaAliases reduce_lam, reduce_nes)
-         map_lam') arrs'
+        (Futhark.ScremaForm (Aliases.removeLambdaAliases scan_lam, scan_nes) reds' map_lam') arrs'
 
     _ -> return $ removeOpAliases soac
   where consumed = consumedInOp soac
diff --git a/src/Futhark/Optimise/Fusion/LoopKernel.hs b/src/Futhark/Optimise/Fusion/LoopKernel.hs
--- a/src/Futhark/Optimise/Fusion/LoopKernel.hs
+++ b/src/Futhark/Optimise/Fusion/LoopKernel.hs
@@ -59,7 +59,8 @@
 
 applyTransform :: SOAC.ArrayTransform -> Ident -> (BasicOp, Certificates)
 applyTransform (SOAC.Rearrange cs perm) v =
-  (Rearrange perm $ identName v, cs)
+  (Rearrange perm' $ identName v, cs)
+  where perm' = perm ++ drop (length perm) [0..arrayRank (identType v)-1]
 applyTransform (SOAC.Reshape cs shape) v =
   (Reshape shape $ identName v, cs)
 applyTransform (SOAC.ReshapeOuter cs shape) v =
@@ -208,7 +209,7 @@
                          (([], []), (p,inp):_) -> ([p], [inp])
                          ((ps_, inps_), _)     -> (ps_, inps_)
         used p = paramName p `S.member` freeVars
-        freeVars = freeInBody $ lambdaBody l
+        freeVars = freeIn $ lambdaBody l
 
 -- | Check that the consumer uses at least one output of the producer
 -- unmodified.
@@ -265,10 +266,13 @@
   case (soac_c, soac_p) of
     _ | SOAC.width soac_p /= SOAC.width soac_c -> fail "SOAC widths must match."
 
-    (SOAC.Screma _ (ScremaForm (scan_lam_c, scan_nes_c) (comm_c, red_lam_c, red_nes_c) _) _,
-     SOAC.Screma _ (ScremaForm (scan_lam_p, scan_nes_p) (comm_p, red_lam_p, red_nes_p) _) _)
-      | mapFusionOK (drop (length $ scan_nes_p++red_nes_p) outVars) ker || horizFuse -> do
-      let (res_lam', new_inp) = fuseRedomap unfus_set outVars
+    (SOAC.Screma _ (ScremaForm (scan_lam_c, scan_nes_c) reds_c _) _,
+     SOAC.Screma _ (ScremaForm (scan_lam_p, scan_nes_p) reds_p _) _)
+      | mapFusionOK (drop (length scan_nes_p+Futhark.redResults reds_p) outVars) ker
+        || horizFuse -> do
+      let red_nes_p = concatMap redNeutral reds_p
+          red_nes_c = concatMap redNeutral reds_c
+          (res_lam', new_inp) = fuseRedomap unfus_set outVars
                                             lam_p scan_nes_p red_nes_p inp_p_arr
                                             outPairs
                                             lam_c scan_nes_c red_nes_c inp_c_arr
@@ -278,14 +282,11 @@
             splitAt3 (length scan_nes_c) (length red_nes_c) $ outNames ker
           unfus_arrs  = returned_outvars \\ (soac_p_scanout++soac_p_redout)
           scan_lam'   = mergeReduceOps scan_lam_p scan_lam_c
-          red_lam'    = mergeReduceOps red_lam_p red_lam_c
       success (soac_p_scanout ++ soac_c_scanout ++
                soac_p_redout ++ soac_c_redout ++
                soac_c_mapout ++ unfus_arrs) $
-        SOAC.Screma w (ScremaForm (scan_lam', scan_nes_p++scan_nes_c)
-                                      (comm_p<>comm_c, red_lam', red_nes_p++red_nes_c)
-                                      res_lam')
-                        new_inp
+        SOAC.Screma w (ScremaForm (scan_lam', scan_nes_p++scan_nes_c) (reds_p ++ reds_c) res_lam')
+        new_inp
 
     ------------------
     -- Scatter fusion --
@@ -555,18 +556,14 @@
 
       let map_body = mkBody (oneStm $
                               Let (setPatternOuterDimTo w map_pat) (defAux ()) $
-                              Op $ Futhark.Screma w (ScremaForm (scan_fun', nes')
-                                                                    (mempty, nilFn, mempty)
-                                                                    id_map_lam) arrs') $
+                              Op $ Futhark.Screma w (ScremaForm (scan_fun', nes') [] id_map_lam) arrs') $
                             map Var $ patternNames map_pat
           map_fun' = Lambda map_params map_body map_rettype
           perm = case lambdaReturnType map_fun of
                    []  -> []
                    t:_ -> 1 : 0 : [2..arrayRank t]
 
-      return (SOAC.Screma map_w
-               (ScremaForm (nilFn, mempty) (mempty, nilFn, mempty) map_fun')
-               map_arrs',
+      return (SOAC.Screma map_w (ScremaForm (nilFn, mempty) [] map_fun') map_arrs',
               ots SOAC.|> SOAC.Rearrange map_cs perm)
 
 iswim _ _ _ =
diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs
--- a/src/Futhark/Optimise/InPlaceLowering.hs
+++ b/src/Futhark/Optimise/InPlaceLowering.hs
@@ -117,9 +117,10 @@
     updates -> do
       let updateStms = map updateStm updates
       lower <- asks lowerUpdate
+      scope <- askScope
       -- Condition (5) and (7) are assumed to be checked by
       -- lowerUpdate.
-      case lower bnd' updates of
+      case lower scope bnd' updates of
         Just lowering -> do new_bnds <- lowering
                             new_bnds' <- optimiseStms new_bnds $
                                          tell bup { forwardThese = [] }
@@ -156,13 +157,16 @@
   where optimise = identityMapper { mapOnBody = const optimiseBody
                                   }
 onKernelOp :: OnOp Kernels
-onKernelOp (HostOp (Kernel debug kspace ts kbody)) = do
+onKernelOp (HostOp op) = do
   old_scope <- askScope
   modifyNameSource $ runForwardingM lowerUpdateInKernel onKernelExp $
-    bindingScope (castScope old_scope <> scopeOfKernelSpace kspace) $ do
-    stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $
-            mapM_ seenVar $ freeIn $ kernelBodyResult kbody
-    return $ HostOp $ Kernel debug kspace ts $ kbody { kernelBodyStms = stmsFromList stms }
+    bindingScope (castScope old_scope <> scopeOfKernelSpace (kernelSpace op)) $ do
+      let mapper = identityKernelMapper { mapOnKernelKernelBody = onKernelBody }
+          onKernelBody kbody = do
+            stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $
+                    mapM_ seenVar $ freeIn $ kernelBodyResult kbody
+            return kbody { kernelBodyStms = stmsFromList stms }
+      HostOp <$> mapKernelM mapper op
 onKernelOp op = return op
 
 onKernelExp :: OnOp InKernel
diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
--- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
+++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs
@@ -13,6 +13,7 @@
 import Data.List (find)
 import Data.Maybe (mapMaybe)
 import Data.Either
+import qualified Data.Map as M
 import qualified Data.Set as S
 
 import Futhark.Representation.AST.Attributes.Aliases
@@ -38,20 +39,21 @@
 updateHasValue :: VName -> DesiredUpdate attr -> Bool
 updateHasValue name = (name==) . updateValue
 
-type LowerUpdate lore m = Stm (Aliases lore)
+type LowerUpdate lore m = Scope (Aliases lore)
+                          -> Stm (Aliases lore)
                           -> [DesiredUpdate (LetAttr (Aliases lore))]
                           -> Maybe (m [Stm (Aliases lore)])
 
 lowerUpdate :: (MonadFreshNames m, Bindable lore,
                 LetAttr lore ~ Type, CanBeAliased (Op lore)) => LowerUpdate lore m
-lowerUpdate (Let pat aux (DoLoop ctx val form body)) updates = do
-  canDo <- lowerUpdateIntoLoop updates pat ctx val body
+lowerUpdate scope (Let pat aux (DoLoop ctx val form body)) updates = do
+  canDo <- lowerUpdateIntoLoop scope updates pat ctx val form body
   Just $ do
     (prebnds, postbnds, ctxpat, valpat, ctx', val', body') <- canDo
     return $
       prebnds ++ [certify (stmAuxCerts aux) $
                   mkLet ctxpat valpat $ DoLoop ctx' val' form body'] ++ postbnds
-lowerUpdate
+lowerUpdate _
   (Let pat aux (BasicOp (SubExp (Var v))))
   [DesiredUpdate bindee_nm bindee_attr cs src is val]
   | patternNames pat == [src] =
@@ -60,11 +62,11 @@
        return [certify (stmAuxCerts aux <> cs) $
                mkLet [] [Ident bindee_nm $ typeOf bindee_attr] $
                BasicOp $ Update v is' $ Var val]
-lowerUpdate _ _ =
+lowerUpdate _ _ _ =
   Nothing
 
 lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m
-lowerUpdateKernels
+lowerUpdateKernels _
   (Let (Pattern [] [PatElem v v_attr]) aux (Op (HostOp (Kernel debug kspace ts kbody))))
   [update@(DesiredUpdate bindee_nm bindee_attr cs _src is val)]
   | v == val = do
@@ -74,7 +76,7 @@
                     mkLet [] [Ident bindee_nm $ typeOf bindee_attr] $
                     Op $ HostOp $ Kernel debug kspace ts kbody',
                    mkLet [] [Ident v $ typeOf v_attr] $ BasicOp $ Index bindee_nm is']
-lowerUpdateKernels stm updates = lowerUpdate stm updates
+lowerUpdateKernels scope stm updates = lowerUpdate scope stm updates
 
 lowerUpdateInKernel :: MonadFreshNames m => LowerUpdate InKernel m
 lowerUpdateInKernel = lowerUpdate
@@ -83,7 +85,7 @@
                       -> KernelSpace -> KernelBody (Aliases InKernel)
                       -> Maybe (KernelBody (Aliases InKernel))
 lowerUpdateIntoKernel update kspace kbody = do
-  [ThreadsReturn ThreadsInSpace se] <- Just $ kernelBodyResult kbody
+  [ThreadsReturn se] <- Just $ kernelBodyResult kbody
   is' <- mapM dimFix is
   let ret = WriteReturn (arrayDims $ snd bindee_attr) src [(is'++map Var gtids, se)]
   return kbody { kernelBodyResult = [ret] }
@@ -93,10 +95,12 @@
 lowerUpdateIntoLoop :: (Bindable lore, BinderOps lore,
                         Aliased lore, LetAttr lore ~ (als, Type),
                         MonadFreshNames m) =>
-                       [DesiredUpdate (LetAttr lore)]
+                       Scope lore
+                    -> [DesiredUpdate (LetAttr lore)]
                     -> Pattern lore
                     -> [(FParam lore, SubExp)]
                     -> [(FParam lore, SubExp)]
+                    -> LoopForm lore
                     -> Body lore
                     -> Maybe (m ([Stm lore],
                                  [Stm lore],
@@ -105,7 +109,7 @@
                                  [(FParam lore, SubExp)],
                                  [(FParam lore, SubExp)],
                                  Body lore))
-lowerUpdateIntoLoop updates pat ctx val body = do
+lowerUpdateIntoLoop scope updates pat ctx val form body = do
   -- Algorithm:
   --
   --   0) Map each result of the loop body to a corresponding in-place
@@ -141,7 +145,10 @@
     (body_res, res_bnds) <- manipulateResult in_place_map idxsubsts'
     let body' = mkBody (newbnds<>res_bnds) body_res
     return (prebnds, postbnds, ctxpat, valpat, ctx, val', body')
-  where usedInBody = freeInBody body
+  where usedInBody = S.unions $ map expandAliases $ S.toList $ freeIn body <> freeIn form
+        expandAliases v = case M.lookup v scope of
+                            Just (LetInfo attr) -> S.insert v $ aliasesOf attr
+                            _ -> S.singleton v
         resmap = zip (bodyResult body) $ patternValueIdents pat
 
         mkMerges :: (MonadFreshNames m, Bindable lore) =>
diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs
--- a/src/Futhark/Optimise/Simplify/Engine.hs
+++ b/src/Futhark/Optimise/Simplify/Engine.hs
@@ -377,7 +377,7 @@
   where block blocked (Left need) =
           (blocked <> S.fromList (provides need), Left need)
         block blocked (Right need)
-          | blocked `intersects` requires need =
+          | blocked `intersects` freeIn need =
             (blocked <> S.fromList (provides need), Left need)
           | otherwise =
             (blocked, Right need)
@@ -385,10 +385,7 @@
 provides :: Stm lore -> [VName]
 provides = patternNames . stmPattern
 
-requires :: Attributes lore => Stm lore -> Names
-requires = freeInStm
-
-expandUsage :: (Attributes lore, Aliased lore, UsageInOp (Op lore)) =>
+expandUsage :: (Attributes lore, Aliased lore) =>
                ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> UT.UsageTable
 expandUsage vtable utable bnd =
   UT.expand (`ST.lookupAliases` vtable) (usageInStm bnd <> usageThroughAliases) <>
@@ -450,7 +447,7 @@
 insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False)
 
 hasFree :: Attributes lore => Names -> BlockPred lore
-hasFree ks _ need = ks `intersects` requires need
+hasFree ks _ need = ks `intersects` freeIn need
 
 isNotSafe :: Attributes lore => BlockPred lore
 isNotSafe _ = not . safeExp . stmExp
@@ -485,7 +482,7 @@
 
 loopInvariantStm :: Attributes lore => ST.SymbolTable lore -> Stm lore -> Bool
 loopInvariantStm vtable =
-  all (`S.member` ST.availableAtClosestLoop vtable) . freeInStm
+  all (`S.member` ST.availableAtClosestLoop vtable) . freeIn
 
 hoistCommon :: SimplifiableLore lore =>
                SubExp -> IfSort
@@ -538,7 +535,7 @@
           in  sel_nms
         transClosSizes all_bnds scal_nms hoist_bnds =
           let new_bnds = filter (hasPatName scal_nms) all_bnds
-              new_nms  = mconcat $ map (freeInExp . stmExp) new_bnds
+              new_nms  = mconcat $ map (freeIn . stmExp) new_bnds
           in  if null new_bnds
               then hoist_bnds
               else transClosSizes all_bnds new_nms (new_bnds ++ hoist_bnds)
@@ -721,7 +718,6 @@
                 -- Lambdas are handled explicitly because we need to
                 -- bind their parameters.
                 , mapOnVName = simplify
-                , mapOnCertificates = simplify
                 , mapOnRetType = simplify
                 , mapOnBranchType = simplify
                 , mapOnFParam =
@@ -810,8 +806,8 @@
   simplify (Array et shape u) = do
     shape' <- simplify shape
     return $ Array et shape' u
-  simplify (Mem size space) =
-    Mem <$> simplify size <*> pure space
+  simplify (Mem space) =
+    pure $ Mem space
   simplify (Prim bt) =
     return $ Prim bt
 
diff --git a/src/Futhark/Optimise/Simplify/Lore.hs b/src/Futhark/Optimise/Simplify/Lore.hs
--- a/src/Futhark/Optimise/Simplify/Lore.hs
+++ b/src/Futhark/Optimise/Simplify/Lore.hs
@@ -42,7 +42,6 @@
 import Futhark.Transform.Rename
 import Futhark.Transform.Substitute
 import Futhark.Analysis.Rephrase
-import Futhark.Analysis.Usage (UsageInOp)
 
 data Wise lore
 
@@ -235,7 +234,7 @@
 mkWiseExpAttr pat explore e =
   (ExpWisdom
     (Names' $ consumedInExp e)
-    (Names' $ freeIn pat <> freeIn explore <> freeInExp e),
+    (Names' $ freeIn pat <> freeIn explore <> freeIn e),
    explore)
 
 instance (Bindable lore,
@@ -258,8 +257,7 @@
 
 class (AliasedOp (OpWithWisdom op),
        RangedOp (OpWithWisdom op),
-       IsOp (OpWithWisdom op),
-       UsageInOp (OpWithWisdom op)) => CanBeWise op where
+       IsOp (OpWithWisdom op)) => CanBeWise op where
   type OpWithWisdom op :: *
   removeOpWisdom :: OpWithWisdom op -> op
 
diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs
--- a/src/Futhark/Optimise/Simplify/Rules.hs
+++ b/src/Futhark/Optimise/Simplify/Rules.hs
@@ -37,7 +37,6 @@
 import Futhark.Representation.AST
 import Futhark.Representation.AST.Attributes.Aliases
 import Futhark.Construct
-import Futhark.Transform.Substitute
 import Futhark.Util
 
 topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore]
@@ -177,7 +176,7 @@
         res = bodyResult loopbody
 
         implpat = zip (patternContextElements pat) $
-                  map paramName $ loopResultContext (map fst ctx) (map fst val)
+                  map (paramName . fst) ctx
         explpat = zip (patternValueElements pat) $
                   map (paramName . fst) val
 
@@ -293,7 +292,7 @@
           (x,x_stms) <- collectStms m
           case x of
             IndexResult cs arr' slice
-              | all (not . (i `S.member`) . freeInStm) x_stms,
+              | all (not . (i `S.member`) . freeIn) x_stms,
                 DimFix (Var j) : slice' <- slice,
                 j == i, not $ i `S.member` freeIn slice -> do
                   addStms x_stms
@@ -334,17 +333,11 @@
     letBindNames_ [paramName p] $ BasicOp $ Index arr $
     DimFix (intConst Int32 0) : fullSlice (paramType p) []
 
-  (loop_body_ctx, loop_body_val) <- splitAt (length ctx) <$> (mapM asVar =<< bodyBind body)
-  let subst = M.fromList $ zip (map (paramName . fst) ctx) loop_body_ctx
-      ctx_params = substituteNames subst $ map fst ctx
-      val_params = substituteNames subst $ map fst val
-      res_context = loopResultContext ctx_params val_params
-  forM_ (zip (patternContextElements pat) res_context) $ \(pat_elem, p) ->
-    letBind_ (Pattern [] [pat_elem]) $ BasicOp $ SubExp $ Var $ paramName p
-  forM_ (zip (patternValueElements pat) loop_body_val) $ \(pat_elem, v) ->
-    letBind_ (Pattern [] [pat_elem]) $ BasicOp $ SubExp $ Var v
-  where asVar (Var v)      = return v
-        asVar (Constant v) = letExp "named" $ BasicOp $ SubExp $ Constant v
+  -- Some of the sizes in the types here might be temporarily wrong
+  -- until copy propagation fixes it up.
+  res <- bodyBind body
+  forM_ (zip (patternNames pat) res) $ \(v, se) ->
+    letBindNames_ [v] $ BasicOp $ SubExp se
 simplifKnownIterationLoop _ _ _ _ =
   cannotSimplify
 
@@ -618,7 +611,8 @@
 
       | Just inds' <- sliceIndices inds,
         Just (e, cs) <- ST.index idd inds' vtable,
-        worthInlining e ->
+        worthInlining e,
+        all (`ST.elem` vtable) (unCertificates cs) ->
         Just $ SubExpResult cs <$> (letSubExp "index_primexp" =<< toExp e)
 
     Nothing -> Nothing
diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs
--- a/src/Futhark/Optimise/TileLoops.hs
+++ b/src/Futhark/Optimise/TileLoops.hs
@@ -151,6 +151,7 @@
 
           let kspace' = kspace { spaceStructure = NestedThreadSpace gspace'
                                , spaceGroupSize = tiled_group_size
+                               , spaceNumVirtGroups = num_groups
                                , spaceNumThreads = num_threads
                                , spaceNumGroups = num_groups
                                }
@@ -246,6 +247,7 @@
 
           let kspace' = kspace { spaceGroupSize = Var inner_ldim
                                , spaceNumGroups = num_groups
+                               , spaceNumVirtGroups = num_groups
                                , spaceNumThreads = num_threads
                                , spaceStructure = structure
                                }
@@ -374,7 +376,7 @@
   foldl' add variance $ patternNames $ stmPattern bnd
   where add variance' v = M.insert v binding_variance variance'
         look variance' v = S.insert v $ M.findWithDefault mempty v variance'
-        binding_variance = mconcat $ map (look variance) $ S.toList (freeInStm bnd)
+        binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd)
 
 sufficientGroups :: MonadBinder m =>
                     [(VName, SubExp, VName, SubExp)] -> SubExp
diff --git a/src/Futhark/Optimise/TileLoops/RegTiling3D.hs b/src/Futhark/Optimise/TileLoops/RegTiling3D.hs
--- a/src/Futhark/Optimise/TileLoops/RegTiling3D.hs
+++ b/src/Futhark/Optimise/TileLoops/RegTiling3D.hs
@@ -200,12 +200,11 @@
 
       -- replace the `ThreadsInSpace` kernel return to an `InPlace` return
       -- for the z-variant kernel results
-      let ker_res_ip_tp_tab = M.fromList $ zip ker_var_res $ zip ip_out_nms $
-                                           map patElemType ker_var_patels
+      let ker_res_ip_tp_tab = M.fromList $ zip ker_var_res ip_out_nms
           (kres', kertp') = unzip $
-            zipWith (\ r tp -> case M.lookup r ker_res_ip_tp_tab of
-                                 Nothing -> (ThreadsReturn ThreadsInSpace (Var r), tp)
-                                 Just (ip_nm, ip_tp) -> (KernelInPlaceReturn ip_nm, ip_tp)
+            zipWith (\r tp -> case M.lookup r ker_res_ip_tp_tab of
+                                Nothing -> (ThreadsReturn (Var r), tp)
+                                Just (dims, arr, ivs) -> (WriteReturn dims arr ivs, tp)
                     ) ker_res_nms kertp
 
       -- finally, put everything together
@@ -263,7 +262,7 @@
             mkLet [] [Ident res_nm0 $ Prim int32] $
               BasicOp $ BinOp (Mul Int32) op1_se op2_se
 
-        retThreadInSpace (ThreadsReturn ThreadsInSpace (Var r)) = Just r
+        retThreadInSpace (ThreadsReturn (Var r)) = Just r
         retThreadInSpace _ = Nothing
 
 doRegTiling3D _ = return Nothing
@@ -328,7 +327,7 @@
     --    depend only on variables defined in the invariant stms to the z parallel dimension.
     var_nms <- concatMap (patternNames . stmPattern) var_out_stms,
     null $ S.intersection (S.fromList var_nms) $
-                          S.unions (map freeInStm var_ind_stms),
+                          S.unions (map freeIn var_ind_stms),
     -- 7. We assume (check) for simplicity that all accumulator initializers
     --     of the outer stream are invariant to the z parallel dimension.
     loop_ini_vs <- subExpVars accs_o_p,
@@ -391,7 +390,7 @@
                          (stmsFromList stms_body_i') $
                          map Var loop_ress
           myloop = DoLoop [] (zip loop_form_acc loop_inis_acc) form body_i'
-          free_in_body = freeInBody body_i'
+          free_in_body = freeIn body_i'
           elim_vars = S.fromList $ arrs_i_p ++ arrs_o_p ++
                                    map paramName arrs_i_f ++
                                    map paramName accs_o_f
@@ -468,14 +467,15 @@
 simpleStm :: Stm InKernel -> Bool
 simpleStm (Let _ _ e) = safeExp e
 
-mkScratchStm :: PatElem Kernels -> TileM (VName, Stm Kernels)
+mkScratchStm :: PatElem Kernels -> TileM (([SubExp], VName, [([SubExp], SubExp)]),
+                                          Stm Kernels)
 mkScratchStm ker_patel = do
   let (unique_arr_tp, res_arr_nm0) = (patElemType ker_patel, patElemName ker_patel)
       ptp = elemType unique_arr_tp
   scrtch_arr_nm <- newVName $ baseString res_arr_nm0 ++ "_0"
   let scratch_stm = mkLet [] [Ident scrtch_arr_nm unique_arr_tp] $
                           BasicOp $ Scratch ptp $ arrayDims unique_arr_tp
-  return (scrtch_arr_nm, scratch_stm)
+  return ((arrayDims unique_arr_tp, scrtch_arr_nm, []), scratch_stm)
 
 -- | Arguments are:
 --     1. @mm@ this is the length of z-parallel dimension divided by reg_tile
@@ -484,52 +484,50 @@
 --     4. @keres_patels@: the kernel result names tupled with the corresponding
 --                        pattern elements of the kernel statement.
 --     5. @code2_var@: the z-variant statements of the code after the stream.
---     6. @ip_arr_nms@: the "current" new names for the in-place update arrays.
+--     6. @ip_writes@: the "current" argument to a 'WriteReturn'.
 --        @unroll_code@: the current unrolled code. Both form a `foldM` accumulator.
 --     7. @k@ the "current" clone number;
 --        @loop_res_nms@ the names of the loop result corresponding to the current clone.
 --   Result:
---     1. the new name for the current in-place update result,
+--     1. the argument for the current in-place update result,
 --     2. a new if-statement is added to the unrolled-code accumulator which actually
 --        perform the in-place update.
 cloneVarCode2 :: VName -> KernelSpace -> [VName]
               -> [(VName, PatElem InKernel)] -> [Stm InKernel]
-              -> ([VName], [Stm InKernel]) -> (Int32, [VName])
-              -> TileM ([VName], [Stm InKernel])
+              -> ([([SubExp], VName, [([SubExp], SubExp)])],
+                  [Stm InKernel])
+              -> (Int32, [VName])
+              -> TileM ([([SubExp], VName, [([SubExp], SubExp)])],
+                        [Stm InKernel])
 cloneVarCode2 mm space strm_res_nms keres_patels code2_var
-              (ip_arr_nms, unroll_code) (k, loop_res_nms) = do
+              (writes, unroll_code) (k, loop_res_nms) = do
   let (ker_nms, pat_els) = unzip keres_patels
-      arr_tps = map patElemType pat_els
       root_strs = map (baseString . patElemName) pat_els
-  ip_inn_nms <- mapM (\s -> newVName $ s ++ "_inn_" ++ pretty (k+1)) root_strs
   ip_out_nms <- mapM (\s -> newVName $ s ++ "_out_" ++ pretty (k+1)) root_strs
   m <- newVName "m"
   -- make in-place update statements
   let (gidx,_) : (gidy,_) : (gidz,m_M) : rev_outer_dims = reverse $ spaceDimensions space
       (outer_dims, _) = unzip $ reverse rev_outer_dims
-      ip_stmts = map (mkInPlaceStmt (outer_dims++[m,gidy,gidx])) $
-                     zip4 ip_arr_nms ip_inn_nms ker_nms arr_tps
+      strip_dims = length $ outer_dims++[m,gidy,gidx]
+      ts = map (stripArray strip_dims . patElemType) pat_els
   -- make if
   cond_nm <- newVName "m_cond"
   let i_se = Constant $ IntValue $ Int32Value k
       m_stm = mkLet [] [Ident m $ Prim int32] $
                     BasicOp $ BinOp (Add Int32) (Var mm) i_se
       c_stm = mkCondStmt m_M m cond_nm
-      else_body = Body () mempty (map Var ip_arr_nms)
       strm_loop_tab = M.fromList $ (gidz, m) : zip strm_res_nms loop_res_nms
-      then_stms = stmsFromList $ map (substituteNames strm_loop_tab) $
-                                     code2_var ++ ip_stmts
-  then_body <- renameBody $ Body () then_stms $ map Var ip_inn_nms
-  let if_stm = mkLet [] (zipWith Ident ip_out_nms arr_tps) $
+  then_body <- renameBody $ substituteNames strm_loop_tab $
+               Body () (stmsFromList code2_var) $ map Var ker_nms
+  let else_body = Body () mempty $ map blank ts
+      if_stm = mkLet [] (zipWith Ident ip_out_nms ts) $
                      If (Var cond_nm) then_body else_body  $
-                     IfAttr (staticShapes arr_tps) IfFallback
-  return (ip_out_nms, unroll_code ++ [m_stm, c_stm, if_stm])
-  where mkInPlaceStmt :: [VName] -> (VName, VName, VName, Type)
-                      -> Stm InKernel
-        mkInPlaceStmt inds (cur_nm, new_nm, ker_nm, arr_tp) =
-          let upd_slc = map (DimFix . Var) inds
-              ipupd_exp = BasicOp $ Update cur_nm upd_slc (Var ker_nm)
-          in  mkLet [] [Ident new_nm arr_tp] ipupd_exp
+                     IfAttr (staticShapes ts) IfFallback
+      addWritePair (dims, arr, current) ker_nm =
+        (dims, arr, current ++ [(map Var $ outer_dims++[m,gidy,gidx], Var ker_nm)])
+  return (zipWith addWritePair writes ip_out_nms, unroll_code ++ [m_stm, c_stm, if_stm])
+  where blank (Prim t) = Constant $ blankPrimValue t
+        blank t = error $ "cloneVarCode2: cannot tile non-prim type " ++ pretty t
 
 helper3Stms :: VName -> SubExp -> SubExp -> Slice SubExp
              -> VName -> Stm InKernel -> TileM [Stm InKernel]
@@ -723,7 +721,7 @@
   foldl' add variance $ patternNames $ stmPattern bnd
   where add variance' v = M.insert v binding_variance variance'
         look variance' v = S.insert v $ M.findWithDefault mempty v variance'
-        binding_variance = mconcat $ map (look variance) $ S.toList (freeInStm bnd)
+        binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd)
 
 sufficientGroups :: MonadBinder m =>
                     [(VName, SubExp, VName, SubExp)] -> SubExp
diff --git a/src/Futhark/Optimise/Unstream.hs b/src/Futhark/Optimise/Unstream.hs
--- a/src/Futhark/Optimise/Unstream.hs
+++ b/src/Futhark/Optimise/Unstream.hs
@@ -39,12 +39,18 @@
   Body () <$> (stmsFromList . concat <$> mapM optimiseStm (stmsToList stms)) <*> pure res
 
 optimiseStm :: Stm Kernels -> UnstreamM [Stm Kernels]
-optimiseStm (Let pat aux (Op (HostOp (Kernel desc space ts body)))) = do
+optimiseStm (Let pat aux (Op (HostOp op))) = do
   inv <- S.fromList . M.keys <$> askScope
-  stms' <- localScope (scopeOfKernelSpace space) $
-           runBinder_ $ optimiseInKernelStms inv $ kernelBodyStms body
-  return [Let pat aux $ Op $ HostOp $
-          Kernel desc space ts $ body { kernelBodyStms = stms' }]
+
+  let mapper = identityKernelMapper { mapOnKernelKernelBody = onKernelBody }
+      onKernelBody kbody = do
+        stms' <- localScope (scopeOfKernelSpace (kernelSpace op)) $
+                 runBinder_ $ optimiseInKernelStms inv $ kernelBodyStms kbody
+        return kbody { kernelBodyStms = stms' }
+
+  op' <- mapKernelM mapper op
+  return [Let pat aux $ Op $ HostOp op']
+
 optimiseStm (Let pat aux e) =
   pure <$> (Let pat aux <$> mapExpM optimise e)
   where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody }
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
@@ -26,7 +26,7 @@
 import qualified Futhark.Representation.Kernels as Kernels
 import Futhark.Representation.Kernels.Simplify as Kernels
 import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun
-import Futhark.Pass.ExtractKernels.BlockedKernel (blockedReduction)
+import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed)
 import Futhark.Pass.ExplicitAllocations (explicitAllocationsInStms)
 import Futhark.Util.IntegralExp
 import Futhark.Util (mapAccumLM)
@@ -93,76 +93,70 @@
           S.fromList $ M.keys $ scopeOfKernelSpace kspace <>
           scopeOf (kernelBodyStms kbody)
 
-transformExp (Op (Inner (HostOp (SegRed kspace comm red_op nes ts kbody)))) = do
-  let (kbody', kbody_allocs) = extractBodyAllocations kbody
-      (red_op', red_op_allocs) = extractLambdaAllocations red_op
-      variantAlloc (Var v) = v `S.member` bound_in_kernel
-      variantAlloc _ = False
-      allocs = kbody_allocs <> red_op_allocs
-      (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs
+transformExp (Op (Inner (HostOp (SegMap kspace ts kbody)))) = do
+  (alloc_stms, (_, kbody')) <- transformScanRed kspace [] kbody
+  return (alloc_stms,
+          Op $ Inner $ HostOp $ SegMap kspace ts kbody')
 
-  (alloc_stms, alloc_offsets) <-
-    memoryRequirements kspace (bodyStms kbody) variant_allocs invariant_allocs
+transformExp (Op (Inner (HostOp (SegRed kspace reds ts kbody)))) = do
+  (alloc_stms, (lams, kbody')) <-
+    transformScanRed kspace (map segRedLambda reds) kbody
+  let reds' = zipWith (\red lam -> red { segRedLambda = lam }) reds lams
+  return (alloc_stms,
+          Op $ Inner $ HostOp $ SegRed kspace reds' ts kbody')
 
-  scope <- askScope
-  let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope
-  either compilerLimitationS pure $ runOffsetM scope' alloc_offsets $ do
-    kbody'' <- offsetMemoryInBody kbody'
-    red_op'' <- localScope (scopeOf red_op') $ offsetMemoryInLambda red_op'
+transformExp (Op (Inner (HostOp (SegScan kspace scan_op nes ts kbody)))) = do
+  (alloc_stms, (scan_op', kbody')) <-
+    transformScanRed kspace [scan_op] kbody
+  return (alloc_stms,
+          Op $ Inner $ HostOp $ SegScan kspace (head scan_op') nes ts kbody')
 
-    return (alloc_stms,
-            Op $ Inner $ HostOp $ SegRed kspace comm red_op'' nes ts kbody'')
+transformExp (Op (Inner (HostOp (SegGenRed kspace ops ts kbody)))) = do
+  (alloc_stms, (lams, kbody')) <-
+    transformScanRed kspace (map genReduceOp ops) kbody
+  let ops' = zipWith (\red lam -> red { genReduceOp = lam }) ops lams
+  return (alloc_stms,
+          Op $ Inner $ HostOp $ SegGenRed kspace ops' ts kbody')
 
-  where bound_in_kernel =
-          S.fromList $ map fst (spaceDimensions kspace) ++
-          M.keys (scopeOfKernelSpace kspace <>
-                  scopeOf (bodyStms kbody))
+transformExp e =
+  return (mempty, e)
 
-transformExp (Op (Inner (HostOp (SegGenRed kspace ops ts kbody)))) = do
-  let (kbody', kbody_allocs) = extractBodyAllocations kbody
-      (ops', ops_allocs) = unzip $ map extractGenRedOpAllocations ops
+transformScanRed :: KernelSpace
+                 -> [Lambda InKernel]
+                 -> KernelBody InKernel
+                 -> ExpandM (Stms ExplicitMemory, ([Lambda InKernel], KernelBody InKernel))
+transformScanRed kspace ops kbody = do
+  let (kbody', kbody_allocs) = extractKernelBodyAllocations kbody
+      (ops', ops_allocs) = unzip $ map extractLambdaAllocations ops
       variantAlloc (Var v) = v `S.member` bound_in_kernel
       variantAlloc _ = False
       allocs = kbody_allocs <> mconcat ops_allocs
       (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs
 
   allocsForBody variant_allocs invariant_allocs kspace kbody' $ \alloc_stms kbody'' -> do
-    ops'' <- mapM offsetMemoryInGenRedOp ops'
-
-    return (alloc_stms,
-            Op $ Inner $ HostOp $ SegGenRed kspace ops'' ts kbody'')
+    ops'' <- forM ops' $ \op' ->
+      localScope (scopeOf op') $ offsetMemoryInLambda op'
+    return (alloc_stms, (ops'', kbody''))
 
   where bound_in_kernel =
           S.fromList $ map fst (spaceDimensions kspace) ++
           M.keys (scopeOfKernelSpace kspace <>
-                  scopeOf (bodyStms kbody))
-
-        extractGenRedOpAllocations op =
-          let (lam, allocs) = extractLambdaAllocations $ genReduceOp op
-          in (op { genReduceOp = lam }, allocs)
-
-        offsetMemoryInGenRedOp op = do
-          lam <- localScope (scopeOf (genReduceOp op)) $
-                 offsetMemoryInLambda $ genReduceOp op
-          return op { genReduceOp = lam }
-
-transformExp e =
-  return (mempty, e)
+                  scopeOf (kernelBodyStms kbody))
 
 allocsForBody :: M.Map VName (SubExp, Space)
               -> M.Map VName (SubExp, Space)
               -> KernelSpace
-              -> Body InKernel
-              -> (Stms ExplicitMemory -> Body InKernel -> OffsetM b)
+              -> KernelBody InKernel
+              -> (Stms ExplicitMemory -> KernelBody InKernel -> OffsetM b)
               -> ExpandM b
 allocsForBody variant_allocs invariant_allocs kspace kbody' m = do
   (alloc_stms, alloc_offsets) <-
-    memoryRequirements kspace (bodyStms kbody') variant_allocs invariant_allocs
+    memoryRequirements kspace (kernelBodyStms kbody') variant_allocs invariant_allocs
 
   scope <- askScope
   let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope
   either compilerLimitationS pure $ runOffsetM scope' alloc_offsets $ do
-    kbody'' <- offsetMemoryInBody kbody'
+    kbody'' <- offsetMemoryInKernelBody kbody'
     m alloc_stms kbody''
 
 memoryRequirements :: KernelSpace
@@ -207,8 +201,7 @@
                          -> (Lambda InKernel,
                              M.Map VName (SubExp, Space))
 extractLambdaAllocations lam = (lam { lambdaBody = body' }, allocs)
-  where (body', allocs) = extractGenericBodyAllocations bodyStms
-                          (\stms body -> body { bodyStms = stms }) $ lambdaBody lam
+  where (body', allocs) = extractBodyAllocations $ lambdaBody lam
 
 extractGenericBodyAllocations :: (body -> Stms InKernel)
                               -> (Stms InKernel -> body -> body)
@@ -227,7 +220,8 @@
 extractThreadAllocations bnds =
   let (allocs, bnds') = mapAccumL isAlloc M.empty $ stmsToList bnds
   in (stmsFromList $ catMaybes bnds', allocs)
-  where isAlloc allocs (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) =
+  where isAlloc allocs (Let (Pattern [] [patElem]) _ (Op (Alloc size space)))
+          | space `notElem` [Space "private", Space "local"] =
           (M.insert (patElemName patElem) (size, space) allocs,
            Nothing)
 
@@ -246,18 +240,10 @@
   (alloc_bnds, rebases) <- unzip <$> mapM expand (M.toList invariant_allocs)
 
   return (mconcat alloc_bnds, mconcat rebases)
-  where expand (mem, (per_thread_size, Space "local")) = do
-          let allocpat = Pattern [] [PatElem mem $
-                                     MemMem per_thread_size $ Space "local"]
-          return (oneStm $ Let allocpat (defAux ()) $
-                   Op $ Alloc per_thread_size $ Space "local",
-                  mempty)
-
-        expand (mem, (per_thread_size, space)) = do
+  where expand (mem, (per_thread_size, space)) = do
           total_size <- newVName "total_size"
           let sizepat = Pattern [] [PatElem total_size $ MemPrim int64]
-              allocpat = Pattern [] [PatElem mem $
-                                     MemMem (Var total_size) space]
+              allocpat = Pattern [] [PatElem mem $ MemMem space]
           return (stmsFromList
                   [Let sizepat (defAux ()) $
                     BasicOp $ BinOp (Mul Int64) num_threads64 per_thread_size,
@@ -306,8 +292,7 @@
 
   return (slice_stms' <> stmsFromList alloc_bnds, mconcat rebases)
   where expand (mem, (offset, total_size, space)) = do
-          let allocpat = Pattern [] [PatElem mem $
-                                     MemMem total_size space]
+          let allocpat = Pattern [] [PatElem mem $ MemMem space]
           return (Let allocpat (defAux ()) $ Op $ Alloc total_size space,
                   M.singleton mem $ newBase offset)
 
@@ -402,7 +387,7 @@
           new_attr <- offsetMemoryInMemBound $ patElemAttr patElem
           return patElem { patElemAttr = new_attr }
         inspectCtx patElem
-          | Mem _ space <- patElemType patElem,
+          | Mem space <- patElemType patElem,
             space /= Space "local" =
               throwError $ unwords ["Cannot deal with existential memory block",
                                     pretty (patElemName patElem),
@@ -545,7 +530,6 @@
                       , mapOnOp = unAllocOp
                       , mapOnSubExp = Right
                       , mapOnVName = Right
-                      , mapOnCertificates = Right
                       }
 
 unAttr :: MemInfo d u ret -> Maybe (TypeBase (ShapeBase d) u)
@@ -583,30 +567,42 @@
       letSubExp "z" $ BasicOp $ BinOp (SMax Int64) (Var $ paramName x) (Var $ paramName y)
     return $ Lambda (xs ++ ys) (mkBody stms zs) i64s
 
+  flat_gtid_lparam <- Param <$> newVName "flat_gtid" <*> pure (Prim (IntType Int32))
+
   (size_lam', _) <- flip runBinderT inkernels_scope $ do
     params <- replicateM num_sizes $ newParam "x" (Prim int64)
     (zs, stms) <- localScope (scopeOfLParams params <>
-                              scopeOfKernelSpace kspace) $ collectStms $ do
+                              scopeOfLParams [flat_gtid_lparam]) $ collectStms $ do
+
+      -- Even though this SegRed is one-dimensional, we need to
+      -- provide indexes corresponding to the original potentially
+      -- multi-dimensional construct.
+      let (kspace_gtids, kspace_dims) = unzip $ spaceDimensions kspace
+          new_inds = unflattenIndex
+                     (map (primExpFromSubExp int32) kspace_dims)
+                     (primExpFromSubExp int32 $ Var $ paramName flat_gtid_lparam)
+      zipWithM_ letBindNames_ (map pure kspace_gtids) =<< mapM toExp new_inds
+
       mapM_ addStm kstms'
       return sizes
+
     localScope (scopeOfKernelSpace kspace) $
       Kernels.simplifyLambda kspace -- XXX, is this the right KernelSpace?
-      (Lambda mempty (Body () stms zs) i64s) []
+      (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) []
 
   ((maxes_per_thread, size_sums), slice_stms) <- flip runBinderT kernels_scope $ do
-    space_size <- letSubExp "space_size" =<<
-                  foldBinOp (Mul Int32) (intConst Int32 1)
-                  (map snd $ spaceDimensions kspace)
     num_threads_64 <- letSubExp "num_threads" $
                       BasicOp $ ConvOp (SExt Int32 Int64) $ spaceNumThreads kspace
 
     pat <- basicPattern [] <$> replicateM num_sizes
            (newIdent "max_per_thread" $ Prim int64)
 
+    thread_space_iota <- letExp "thread_space_iota" $ BasicOp $
+                         Iota (spaceNumThreads kspace) (intConst Int32 0) (intConst Int32 1) Int32
+    let red_op = SegRedOp Commutative max_lam
+                 (replicate num_sizes $ intConst Int64 0) mempty
     addStms =<<
-      blockedReduction pat space_size Commutative
-      max_lam size_lam' (spaceDimensions kspace)
-      (replicate num_sizes $ intConst Int64 0) []
+      nonSegRed pat (spaceNumThreads kspace) [red_op] size_lam' [thread_space_iota]
 
     size_sums <- forM (patternNames pat) $ \threads_max ->
       letExp "size_sum" $
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
@@ -101,7 +101,8 @@
   return $ Var v
 
 type Allocable fromlore tolore =
-  (ExplicitMemorish tolore,
+  (PrettyLore fromlore, PrettyLore tolore,
+   ExplicitMemorish tolore,
    SameScope fromlore Kernels,
    RetType fromlore ~ RetType Kernels,
    BranchType fromlore ~ BranchType Kernels,
@@ -222,20 +223,19 @@
 arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName)
 arraySizeInBytesExpM t = do
   dims <- mapM dimAllocationSize (arrayDims t)
-  let dim_prod_i32 = product $ map (primExpFromSubExp int32) dims
+  let dim_prod_i32 = product $ map (toInt64 . primExpFromSubExp int32) dims
   let elm_size_i64 = ValueExp $ IntValue $ Int64Value $ primByteSize $ elemType t
-  return $ product [ toInt64 dim_prod_i32, elm_size_i64 ]
+  return $ product [ dim_prod_i32, elm_size_i64 ]
   where toInt64 = ConvOpExp $ SExt Int32 Int64
 
 arraySizeInBytes :: Allocator lore m => Type -> m SubExp
 arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM
 
 allocForArray :: Allocator lore m =>
-                 Type -> Space -> m (SubExp, VName)
+                 Type -> Space -> m VName
 allocForArray t space = do
   size <- arraySizeInBytes t
-  m <- allocateMemory "mem" size space
-  return (size, m)
+  allocateMemory "mem" size space
 
 allocsForStm :: (Allocator lore m, ExpAttr lore ~ ()) =>
                 [Ident] -> [Ident] -> Exp lore
@@ -268,7 +268,7 @@
                        [AllocStm])
 allocsForPattern sizeidents validents rts hints = do
   let sizes' = [ PatElem size $ MemPrim int32 | size <- map identName sizeidents ]
-  (vals,(mems_and_sizes, postbnds)) <-
+  (vals, (mems, postbnds)) <-
     runWriterT $ forM (zip3 validents rts hints) $ \(ident, rt, hint) -> do
       let shape = arrayShape $ identType ident
       case rt of
@@ -276,13 +276,9 @@
           summary <- lift $ summaryForBindage (identType ident) hint
           return $ PatElem (identName ident) summary
 
-        MemMem (Free size) space ->
-          return $ PatElem (identName ident) $
-          MemMem size space
-
-        MemMem Ext{} space ->
+        MemMem space ->
           return $ PatElem (identName ident) $
-          MemMem (intConst Int32 0) space
+          MemMem space
 
         MemArray bt _ u (Just (ReturnsInBlock mem ixfun)) ->
           PatElem (identName ident) . MemArray bt shape u .
@@ -295,16 +291,15 @@
 
         MemArray bt _ u ret -> do
           let space = case ret of
-                        Just (ReturnsNewBlock mem_space _ _ _) -> mem_space
-                        _                                      -> DefaultSpace
-          (memsize,mem,(ident',ixfun)) <- lift $ memForBindee ident
-          tell ([PatElem (identName memsize) $ MemPrim int64,
-                 PatElem (identName mem)     $ MemMem (Var $ identName memsize) space],
+                        Just (ReturnsNewBlock mem_space _ _) -> mem_space
+                        _                                    -> DefaultSpace
+          (mem,(ident',ixfun)) <- lift $ memForBindee ident
+          tell ([PatElem (identName mem)     $ MemMem space],
                 [])
           return $ PatElem (identName ident') $ MemArray bt shape u $
             ArrayIn (identName mem) ixfun
 
-  return (sizes' <> mems_and_sizes,
+  return (sizes' <> mems,
           vals,
           postbnds)
   where knownShape = mapM known . shapeDims
@@ -321,10 +316,10 @@
                   -> m (MemBound NoUniqueness)
 summaryForBindage (Prim bt) _ =
   return $ MemPrim bt
-summaryForBindage (Mem size space) _ =
-  return $ MemMem size space
+summaryForBindage (Mem space) _ =
+  return $ MemMem space
 summaryForBindage t@(Array bt shape u) NoHint = do
-  (_, m) <- allocForArray t DefaultSpace
+  m <- allocForArray t DefaultSpace
   return $ directIndexFunction bt shape u m t
 summaryForBindage t (Hint ixfun space) = do
   let bt = elemType t
@@ -337,13 +332,10 @@
 memForBindee :: (MonadFreshNames m) =>
                 Ident
              -> m (Ident,
-                   Ident,
                    (Ident, IxFun))
 memForBindee ident = do
-  size <- newIdent (memname <> "_size") (Prim int64)
-  mem <- newIdent memname $ Mem (Var $ identName size) DefaultSpace
-  return (size,
-          mem,
+  mem <- newIdent memname $ Mem DefaultSpace
+  return (mem,
           (ident, IxFun.iota $ map (primExpFromSubExp int32) $ arrayDims t))
   where  memname = baseString (identName ident) <> "_mem"
          t       = identType ident
@@ -374,15 +366,13 @@
     Array bt shape u -> do
       let memname = baseString (paramName param) <> "_mem"
           ixfun = IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape
-      memsize <- lift $ newVName (memname <> "_size")
       mem <- lift $ newVName memname
-      tell [ Param memsize $ MemPrim int64
-           , Param mem $ MemMem (Var memsize) pspace]
+      tell [Param mem $ MemMem pspace]
       return param { paramAttr =  MemArray bt shape u $ ArrayIn mem ixfun }
     Prim bt ->
       return param { paramAttr = MemPrim bt }
-    Mem size space ->
-      return param { paramAttr = MemMem size space }
+    Mem space ->
+      return param { paramAttr = MemMem space }
 
 allocInMergeParams :: (Allocable fromlore tolore,
                        Allocator tolore (AllocM fromlore tolore)) =>
@@ -394,9 +384,9 @@
                        -> AllocM fromlore tolore a)
                    -> AllocM fromlore tolore a
 allocInMergeParams variant merge m = do
-  ((valparams, handle_loop_subexps), mem_and_size_params) <-
+  ((valparams, handle_loop_subexps), mem_params) <-
     runWriterT $ unzip <$> mapM allocInMergeParam merge
-  let mergeparams' = mem_and_size_params <> valparams
+  let mergeparams' = mem_params <> valparams
       summary = scopeOfFParams mergeparams'
 
       mk_loop_res ses = do
@@ -404,11 +394,11 @@
           runWriterT $ zipWithM ($) handle_loop_subexps ses
         return (memargs, valargs)
 
-  localScope summary $ m mem_and_size_params valparams mk_loop_res
+  localScope summary $ m mem_params valparams mk_loop_res
   where allocInMergeParam (mergeparam, Var v)
           | Array bt shape u <- paramDeclType mergeparam = do
               (mem, ixfun) <- lift $ lookupArraySummary v
-              Mem _ space <- lift $ lookupType mem
+              Mem space <- lift $ lookupType mem
               reuse <- asks aggressiveReuse
               if space /= Space "local" &&
                  reuse &&
@@ -448,12 +438,12 @@
 
 ensureDirectArray :: (Allocable fromlore tolore,
                       Allocator tolore (AllocM fromlore tolore)) =>
-                     Maybe Space -> VName -> AllocM fromlore tolore (SubExp, VName, SubExp)
+                     Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp)
 ensureDirectArray space_ok v = do
   (mem, ixfun) <- lookupArraySummary v
-  Mem size mem_space <- lookupType mem
+  Mem mem_space <- lookupType mem
   if IxFun.isDirect ixfun && maybe True (==mem_space) space_ok
-    then return (size, mem, Var v)
+    then return (mem, Var v)
     else needCopy (fromMaybe DefaultSpace space_ok)
   where needCopy space =
           -- We need to do a new allocation, copy 'v', and make a new
@@ -462,16 +452,16 @@
 
 allocLinearArray :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                     Space -> String -> VName
-                 -> AllocM fromlore tolore (SubExp, VName, SubExp)
+                 -> AllocM fromlore tolore (VName, SubExp)
 allocLinearArray space s v = do
   t <- lookupType v
-  (size, mem) <- allocForArray t space
+  mem <- allocForArray t space
   v' <- newIdent (s ++ "_linear") t
   let pat = Pattern [] [PatElem (identName v') $
                         directIndexFunction (elemType t) (arrayShape t)
                         NoUniqueness mem t]
   addStm $ Let pat (defAux ()) $ BasicOp $ Copy v
-  return (size, mem, Var $ identName v')
+  return (mem, Var $ identName v')
 
 funcallArgs :: (Allocable fromlore tolore,
                 Allocator tolore (AllocM fromlore tolore)) =>
@@ -488,8 +478,8 @@
                     Type -> Space -> SubExp
                  -> WriterT [SubExp] (AllocM fromlore tolore) SubExp
 linearFuncallArg Array{} space (Var v) = do
-  (size, mem, arg') <- lift $ ensureDirectArray (Just space) v
-  tell [size, Var mem]
+  (mem, arg') <- lift $ ensureDirectArray (Just space) v
+  tell [Var mem]
   return arg'
 linearFuncallArg _ _ arg =
   return arg
@@ -510,8 +500,8 @@
   where addAttr (Prim t) = return $ MemPrim t
         addAttr Mem{} = fail "memoryInRetType: too much memory"
         addAttr (Array bt shape u) = do
-          i <- get <* modify (+2)
-          return $ MemArray bt shape u $ ReturnsNewBlock DefaultSpace (i+1) (Ext i) $
+          i <- get <* modify (+1)
+          return $ MemArray bt shape u $ ReturnsNewBlock DefaultSpace i $
             IxFun.iota $ map convert $ shapeDims shape
 
         convert (Ext i) = LeafExp (Ext i) int32
@@ -536,25 +526,41 @@
   return $ Inner $ GetSizeMax size_class
 handleHostOp (CmpSizeLe key size_class x) =
   return $ Inner $ CmpSizeLe key size_class x
-handleHostOp (HostOp (Kernel desc space kernel_ts kbody)) = subInKernel $
+handleHostOp (HostOp (Kernel desc space kernel_ts kbody)) =
+  subInKernel space $
   Inner . HostOp . Kernel desc space kernel_ts <$>
   localScope (scopeOfKernelSpace space) (allocInKernelBody kbody)
 
-handleHostOp (HostOp (SegRed space comm red_op nes ts body)) = do
-  body' <- subInKernel $ localScope (scopeOfKernelSpace space) $ allocInBodyNoDirect body
-  red_op' <- allocInSegRedLambda (spaceGlobalId space) (spaceNumThreads space) red_op
-  return $ Inner $ HostOp $ SegRed space comm red_op' nes ts body'
+handleHostOp (HostOp (SegMap space ts body)) = do
+  body' <- subInKernel space $
+           localScope (scopeOfKernelSpace space) $ allocInKernelBody body
+  return $ Inner $ HostOp $ SegMap space ts body'
 
+handleHostOp (HostOp (SegRed space reds ts body)) = do
+  body' <- subInKernel space $
+           localScope (scopeOfKernelSpace space) $ allocInKernelBody body
+  reds' <- forM reds $ \(SegRedOp comm lam nes shape) -> do
+    lam' <- allocInSegRedLambda space lam
+    return $ SegRedOp comm lam' nes shape
+  return $ Inner $ HostOp $ SegRed space reds' ts body'
+
+handleHostOp (HostOp (SegScan space scan_op nes ts body)) = do
+  body' <- subInKernel space $
+           localScope (scopeOfKernelSpace space) $ allocInKernelBody body
+  scan_op' <- allocInSegRedLambda space scan_op
+  return $ Inner $ HostOp $ SegScan space scan_op' nes ts body'
+
 handleHostOp (HostOp (SegGenRed space ops ts body)) = do
-  body' <- subInKernel $ localScope (scopeOfKernelSpace space) $ allocInBodyNoDirect body
+  body' <- subInKernel space $
+           localScope (scopeOfKernelSpace space) $ allocInKernelBody body
   ops' <- forM ops $ \op -> do
-    lam <- allocInSegRedLambda (spaceGlobalId space) (spaceNumThreads space) $ genReduceOp op
+    lam <- allocInSegRedLambda space $ genReduceOp op
     return op { genReduceOp = lam }
   return $ Inner $ HostOp $ SegGenRed space ops' ts body'
 
-subInKernel :: AllocM InInKernel OutInKernel a
+subInKernel :: KernelSpace -> AllocM InInKernel OutInKernel a
             -> AllocM fromlore2 ExplicitMemory a
-subInKernel = subAllocM handleKernelExp True
+subInKernel space = subAllocM handleKernelExp True
   where handleKernelExp (Barrier se) =
           return $ Inner $ Barrier se
 
@@ -566,13 +572,13 @@
 
         handleKernelExp (GroupReduce w lam input) = do
           summaries <- mapM lookupArraySummary arrs
-          lam' <- allocInReduceLambda lam summaries
+          lam' <- allocInReduceLambda space lam summaries
           return $ Inner $ GroupReduce w lam' input
           where arrs = map snd input
 
         handleKernelExp (GroupScan w lam input) = do
           summaries <- mapM lookupArraySummary arrs
-          lam' <- allocInReduceLambda lam summaries
+          lam' <- allocInReduceLambda space lam summaries
           return $ Inner $ GroupScan w lam' input
           where arrs = map snd input
 
@@ -612,9 +618,7 @@
   case info of
     MemPrim{} -> return []
     MemMem{} -> return [] -- should not happen
-    MemArray _ _ _ (ArrayIn mem _) -> do
-      size <- lookupMemSize mem
-      return [size, Var mem]
+    MemArray _ _ _ (ArrayIn mem _) -> return [Var mem]
 
 allocInFunBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
                   [Maybe Space] -> Body fromlore -> AllocM fromlore tolore (Body tolore)
@@ -633,7 +637,7 @@
           bt <- primType <$> lookupType v
           if bt
             then return $ Var v
-            else do (_, _, v') <- ensureDirectArray space_ok v
+            else do (_, v') <- ensureDirectArray space_ok v
                     return v'
 
 allocInStms :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) =>
@@ -724,7 +728,7 @@
           v_info <- lookupMemInfo v
           case v_info of MemArray _ _ _ (ArrayIn mem _) -> do
                            mem_info <- lookupMemInfo mem
-                           case mem_info of MemMem _ space -> return $ Just space
+                           case mem_info of MemMem space -> return $ Just space
                                             _ -> return Nothing
                          _ -> return Nothing
         mkSpaceOK _ = return Nothing
@@ -733,14 +737,14 @@
 createBodyReturns ts spaces =
   evalState (zipWithM inspect ts spaces) $ S.size $ shapeContext ts
   where inspect (Array pt shape u) space = do
-          i <- get <* modify (+2)
+          i <- get <* modify (+1)
           let space' = fromMaybe DefaultSpace space
-          return $ MemArray pt shape u $ ReturnsNewBlock space' (i+1) (Ext i) $
+          return $ MemArray pt shape u $ ReturnsNewBlock space' i $
             IxFun.iota $ map convert $ shapeDims shape
         inspect (Prim pt) _ =
           return $ MemPrim pt
-        inspect (Mem size space) _ =
-          return $ MemMem (Free size) space
+        inspect (Mem space) _ =
+          return $ MemMem space
 
         convert (Ext i) = LeafExp (Ext i) int32
         convert (Free v) = Free <$> primExpFromSubExp int32 v
@@ -760,19 +764,19 @@
               in return (p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' }, a)
             Prim bt ->
               return (p { paramAttr = MemPrim bt }, a)
-            Mem size space ->
-              return (p { paramAttr = MemMem size space }, a)
+            Mem space ->
+              return (p { paramAttr = MemMem space }, a)
 
-allocInReduceLambda :: Lambda InInKernel
+allocInReduceLambda :: KernelSpace
+                    -> Lambda InInKernel
                     -> [(VName, IxFun)]
                     -> AllocM InInKernel OutInKernel (Lambda OutInKernel)
-allocInReduceLambda lam input_summaries = do
-  let (i, j_param, actual_params) =
-        partitionChunkedKernelLambdaParameters $ lambdaParams lam
-      (acc_params, arr_params) =
-        splitAt (length input_summaries) actual_params
-      this_index = LeafExp i int32
-      other_index = this_index + LeafExp (paramName j_param) int32
+allocInReduceLambda space lam input_summaries = do
+  let (acc_params, arr_params) =
+        splitAt (length input_summaries) $ lambdaParams lam
+      this_index = LeafExp (spaceGlobalId space) int32
+      other_index = this_index + primExpFromSubExp int32 (spaceNumThreads space)
+
   acc_params' <-
     allocInReduceParameters this_index $
     zip acc_params input_summaries
@@ -780,9 +784,7 @@
     allocInReduceParameters other_index $
     zip arr_params input_summaries
 
-  allocInLambda (Param i (MemPrim int32) :
-                 j_param { paramAttr = MemPrim int32 } :
-                 acc_params' ++ arr_params')
+  allocInLambda (acc_params' ++ arr_params')
     (lambdaBody lam) (lambdaReturnType lam)
 
 allocInReduceParameters :: PrimExp VName
@@ -797,20 +799,20 @@
               in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' }
             Prim bt ->
               return p { paramAttr = MemPrim bt }
-            Mem size space ->
-              return p { paramAttr = MemMem size space }
+            Mem space ->
+              return p { paramAttr = MemMem space }
 
-allocInSegRedLambda :: VName -> SubExp -> Lambda InInKernel
+allocInSegRedLambda :: KernelSpace -> Lambda InInKernel
                     -> AllocM Kernels ExplicitMemory (Lambda OutInKernel)
-allocInSegRedLambda gtid num_threads lam = do
+allocInSegRedLambda space lam = do
   let (acc_params, arr_params) =
         splitAt (length (lambdaParams lam) `div` 2) $ lambdaParams lam
-      this_index = LeafExp gtid int32
-      other_index = this_index + primExpFromSubExp int32 num_threads
+      this_index = LeafExp (spaceGlobalId space) int32
+      other_index = this_index + primExpFromSubExp int32 (spaceNumThreads space)
   (acc_params', arr_params') <-
-    allocInSegRedParameters num_threads this_index other_index acc_params arr_params
+    allocInSegRedParameters (spaceNumThreads space) this_index other_index acc_params arr_params
 
-  subInKernel $ allocInLambda (acc_params' ++ arr_params')
+  subInKernel space $ allocInLambda (acc_params' ++ arr_params')
     (lambdaBody lam) (lambdaReturnType lam)
 
 allocInSegRedParameters :: SubExp
@@ -825,7 +827,7 @@
               twice_num_threads <- letSubExp "twice_num_threads" $
                                    BasicOp $ BinOp (Mul Int32) num_threads $ intConst Int32 2
               let t = paramType x `arrayOfRow` twice_num_threads
-              (_, mem) <- allocForArray t DefaultSpace
+              mem <- allocForArray t DefaultSpace
               -- XXX: this iota ixfun is a bit inefficient; leading to uncoalesced access.
               let ixfun_base = IxFun.iota $
                                map (primExpFromSubExp int32) (arrayDims t)
@@ -838,9 +840,9 @@
             Prim bt ->
               return (x { paramAttr = MemPrim bt },
                       y { paramAttr = MemPrim bt })
-            Mem size space ->
-              return (x { paramAttr = MemMem size space },
-                      y { paramAttr = MemMem size space })
+            Mem space ->
+              return (x { paramAttr = MemMem space },
+                      y { paramAttr = MemMem space })
 
 allocInChunkedParameters :: PrimExp VName
                         -> [(LParam InInKernel, (VName, IxFun))]
@@ -853,8 +855,8 @@
               in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' }
             Prim bt ->
               return p { paramAttr = MemPrim bt }
-            Mem size space ->
-              return p { paramAttr = MemMem size space }
+            Mem space ->
+              return p { paramAttr = MemMem space }
 
 allocInLambda :: [LParam OutInKernel] -> Body InInKernel -> [Type]
               -> AllocM InInKernel OutInKernel (Lambda OutInKernel)
@@ -1024,23 +1026,34 @@
               perm_inv
   return [Hint ixfun DefaultSpace]
 
-kernelExpHints (Op (Inner (HostOp (Kernel _ space rets kbody)))) =
-  zipWithM hint rets $ kernelBodyResult kbody
-  where num_threads = spaceNumThreads space
+kernelExpHints (Op (Inner (HostOp (Kernel _ space ts kbody)))) =
+  zipWithM (mapResultHint space) ts $ kernelBodyResult kbody
 
-        spacy AllThreads = Just [num_threads]
-        spacy ThreadsInSpace = Just $ map snd $ spaceDimensions space
-        spacy _ = Nothing
+kernelExpHints (Op (Inner (HostOp (SegMap space ts body)))) =
+  zipWithM (mapResultHint space) ts $ kernelBodyResult body
 
+kernelExpHints (Op (Inner (HostOp (SegRed space reds ts body)))) =
+  (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint space) (drop num_reds ts) map_res
+  where num_reds = segRedResults reds
+        (red_res, map_res) = splitAt num_reds $ kernelBodyResult body
+
+kernelExpHints e =
+  return $ replicate (expExtTypeSize e) NoHint
+
+mapResultHint :: Allocator lore m =>
+                 KernelSpace -> Type -> KernelResult -> m ExpHint
+mapResultHint space = hint
+  where num_threads = spaceNumThreads space
+
         -- Heuristic: do not rearrange for returned arrays that are
         -- sufficiently small.
         coalesceReturnOfShape _ [] = False
         coalesceReturnOfShape bs [Constant (IntValue (Int32Value d))] = bs * d > 4
         coalesceReturnOfShape _ _ = True
 
-        hint t (ThreadsReturn threads _)
-          | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t,
-            Just space_dims <- spacy threads = do
+        hint t (ThreadsReturn _)
+          | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do
+              let space_dims = map snd $ spaceDimensions space
               t_dims <- mapM dimAllocationSize $ arrayDims t
               return $ Hint (innermost space_dims t_dims) DefaultSpace
 
@@ -1057,17 +1070,6 @@
 
         hint _ _ = return NoHint
 
-kernelExpHints (Op (Inner (HostOp (SegRed space _ _ nes ts body)))) =
-  (map (const NoHint) red_res <>) <$> zipWithM mapHint (drop (length nes) ts) map_res
-  where (red_res, map_res) = splitAt (length nes) $ bodyResult body
-
-        mapHint t _ = do
-          t_dims <- mapM dimAllocationSize $ arrayDims t
-          return $ Hint (innermost (map snd $ spaceDimensions space) t_dims) DefaultSpace
-
-kernelExpHints e =
-  return $ replicate (expExtTypeSize e) NoHint
-
 innermost :: [SubExp] -> [SubExp] -> IxFun
 innermost space_dims t_dims =
   let r = length t_dims
@@ -1089,6 +1091,17 @@
     return $ Hint ixfun $ Space "local"
   where dims = map snd cspace
         (_, ns, _) = unzip3 scatter
-
 inKernelExpHints e =
-  return $ replicate (expExtTypeSize e) NoHint
+  mapM maybePrivate =<< expExtType e
+  where maybePrivate t
+          | arrayRank t > 0,
+            Just t' <- hasStaticShape t,
+            all semiStatic $ arrayDims t' = do
+              alloc_dims <- mapM dimAllocationSize $ arrayDims t'
+              let ixfun = IxFun.iota $ map (primExpFromSubExp int32) alloc_dims
+              return $ Hint ixfun $ Space "private"
+          | otherwise =
+              return NoHint
+
+        semiStatic Constant{} = True
+        semiStatic _ = False
diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs
--- a/src/Futhark/Pass/ExtractKernels.hs
+++ b/src/Futhark/Pass/ExtractKernels.hs
@@ -185,7 +185,6 @@
 import Futhark.Pass.ExtractKernels.Distribution
 import Futhark.Pass.ExtractKernels.ISRWIM
 import Futhark.Pass.ExtractKernels.BlockedKernel
-import Futhark.Pass.ExtractKernels.Segmented
 import Futhark.Pass.ExtractKernels.Interchange
 import Futhark.Pass.ExtractKernels.Intragroup
 import Futhark.Util
@@ -255,7 +254,7 @@
 
 sequentialisedUnbalancedStm :: Stm -> DistribM (Maybe (Stms SOACS))
 sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ form _)))
-  | Just (_, _, _, lam2) <- isRedomapSOAC form,
+  | Just (_, lam2) <- isRedomapSOAC form,
     unbalancedLambda lam2, lambdaContainsParallelism lam2 = do
       types <- asksScope scopeForSOACs
       Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types
@@ -302,27 +301,21 @@
       types <- asksScope scopeForSOACs
       transformStms path =<< (stmsToList . snd <$> runBinderT (certifying cs do_iswim) types)
 
-  | Just (scan_lam, scan_nes) <- isScanSOAC form,
-    ScremaForm _ _ map_lam <- form =
-      doScan (scan_lam, scan_nes) (mempty, nilFn, mempty) map_lam
-
-  | ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam <- form,
-    not $ null scan_nes, all primType $ lambdaReturnType scan_lam,
-    not $ lambdaContainsParallelism map_lam =
-      doScan (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam
-
-  where doScan (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam = do
-          scan_lam_sequential <- Kernelise.transformLambda scan_lam
-          red_lam_sequential <- Kernelise.transformLambda red_lam
-          map_lam_sequential <- Kernelise.transformLambda map_lam
-          runBinder_ $ certifying cs $
-            blockedScan res_pat w
-            (scan_lam_sequential, scan_nes)
-            (comm, red_lam_sequential, red_nes)
-            map_lam_sequential (intConst Int32 16) [] [] arrs
+  -- We are only willing to generate code for scanomaps that do not
+  -- involve array accumulators, and do not have parallelism in their
+  -- map function.  Such cases will fall through to the
+  -- screma-splitting case, and produce an ordinary map and scan.
+  -- Hopefully, the scan then triggers the ISWIM case above (otherwise
+  -- we will still crash in code generation).
+  | Just (scan_lam, nes, map_lam) <- isScanomapSOAC form,
+    all primType $ lambdaReturnType scan_lam,
+    not $ lambdaContainsParallelism map_lam = do
+      scan_lam_sequential <- Kernelise.transformLambda scan_lam
+      map_lam_sequential <- Kernelise.transformLambda map_lam
+      segScan res_pat w w scan_lam_sequential map_lam_sequential nes arrs [] []
 
 transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs)))
-  | Just (comm, red_fun, nes) <- isReduceSOAC form,
+  | Just [Reduce comm red_fun nes] <- isReduceSOAC form,
     let comm' | commutativeLambda red_fun = Commutative
               | otherwise                 = comm,
     Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do
@@ -331,14 +324,17 @@
       transformStms path $ stmsToList bnds
 
 transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs)))
-  | Just (comm, red_lam, nes, map_lam) <- isRedomapSOAC form = do
+  | Just (reds, map_lam) <- isRedomapSOAC form = do
 
   let paralleliseOuter = runBinder_ $ do
-        red_lam_sequential <- Kernelise.transformLambda red_lam
+        red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do
+          (red_lam', nes', shape) <- determineReduceOp red_lam nes
+          let comm' | commutativeLambda red_lam = Commutative
+                    | otherwise = comm
+          return $ SegRedOp comm' red_lam' nes' shape
         map_lam_sequential <- Kernelise.transformLambda map_lam
         addStms =<<
-          (fmap (certify cs) <$>
-           nonSegRed pat w comm' red_lam_sequential map_lam_sequential nes arrs)
+          (fmap (certify cs) <$> nonSegRed pat w red_ops map_lam_sequential arrs)
 
       outerParallelBody =
         renameBody =<<
@@ -347,14 +343,14 @@
       paralleliseInner path' = do
         (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs)
         transformStms path' [certify cs mapbnd, certify cs redbnd]
+          where comm' | commutativeLambda red_lam = Commutative
+                      | otherwise = comm
+                (Reduce comm red_lam nes) = singleReduce reds
 
       innerParallelBody path' =
         renameBody =<<
         (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat)))
 
-      comm' | commutativeLambda red_lam = Commutative
-            | otherwise = comm
-
   if not $ lambdaContainsParallelism map_lam
     then paralleliseOuter
     else if incrementalFlattening then do
@@ -398,23 +394,23 @@
           let (red_pat_elems, concat_pat_elems) =
                 splitAt (length nes) $ patternValueElements pat
               red_pat = Pattern [] red_pat_elems
-              concat_pat = Pattern [] concat_pat_elems
 
-          (map_bnd, map_misc_bnds) <- blockedMap concat_pat w InOrder fold_fun_sequential nes arrs
-          let num_threads = arraysSize 0 $ patternTypes $ stmPattern map_bnd
+          ((num_threads, red_results), stms) <-
+            streamMap (map (baseString . patElemName) red_pat_elems) concat_pat_elems w
+            Noncommutative fold_fun_sequential nes arrs
 
-          reduce_soac <- reduceSOAC comm' red_fun nes
+          reduce_soac <- reduceSOAC [Reduce comm' red_fun nes]
 
-          ((map_misc_bnds<>oneStm map_bnd)<>) <$>
-            inScopeOf (map_misc_bnds<>oneStm map_bnd)
+          (stms<>) <$>
+            inScopeOf stms
             (transformStm path' $ Let red_pat aux $
-             Op (Screma num_threads reduce_soac $ patternNames $ stmPattern map_bnd))
+             Op (Screma num_threads reduce_soac red_results))
 
       | otherwise = do
           red_fun_sequential <- Kernelise.transformLambda red_fun
           fold_fun_sequential <- Kernelise.transformLambda fold_fun
           fmap (certify cs) <$>
-            blockedReductionStream pat w comm' red_fun_sequential fold_fun_sequential [] nes arrs
+            streamRed pat w comm' red_fun_sequential fold_fun_sequential nes arrs
 
     outerParallelBody path' =
       renameBody =<<
@@ -433,7 +429,7 @@
           | otherwise                               = comm
 
 transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) = do
-  -- This with-loop is too complicated for us to immediately do
+  -- This screma is too complicated for us to immediately do
   -- anything, so split it up and try again.
   scope <- asksScope scopeForSOACs
   transformStms path . map (certify cs) . stmsToList . snd =<<
@@ -834,7 +830,10 @@
       -- decide whether to just sequentialise, or exploit inner
       -- parallelism.
       let map_nesting = MapNesting pat cs w $ zip (lambdaParams lam) arrs
-          lam_res' = rearrangeShape perm lam_res
+          -- Normally the permutation is for the output pattern, but
+          -- we can't really change that, so we change the result
+          -- order instead.
+          lam_res' = rearrangeShape (rearrangeInverse perm) lam_res
           nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest
           extra_scope = targetsScope $ kernelTargets acc'
 
@@ -853,7 +852,7 @@
           localScope (scopeOfLParams (lambdaParams lam)) $ runBinder_ $
           Kernelise.transformStms lam_bnds
         let kbody = KernelBody () sequentialised_map_body $
-                    map (ThreadsReturn ThreadsInSpace) lam_res'
+                    map ThreadsReturn lam_res'
         constructKernel nest' kbody
 
       let outer_pat = loopNestingPattern $ fst nest
@@ -877,7 +876,7 @@
      if null $ kernelStms acc'
        then return acc'
        else do let kbody = Body () (kernelStms acc') res
-                   used_in_body = freeInBody kbody
+                   used_in_body = freeIn kbody
                    (used_params, used_arrs) =
                      unzip $
                      filter ((`S.member` used_in_body) . paramName . fst) $
@@ -962,7 +961,7 @@
         addStmToKernel stm acc
 
 maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc
-  | Just (comm, lam, nes) <- isReduceSOAC form,
+  | Just [Reduce comm lam nes] <- isReduceSOAC form,
     Just m <- irwim pat w comm lam $ zip nes arrs = do
       types <- asksScope scopeForSOACs
       (_, bnds) <- runBinderT (certifying cs m) types
@@ -1024,7 +1023,8 @@
 -- If the reduction cannot be distributed by itself, it will be
 -- sequentialised in the default case for this function.
 maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc
-  | Just (comm, lam, nes, map_lam) <- isRedomapSOAC form,
+  | Just (reds, map_lam) <- isRedomapSOAC form,
+    Reduce comm lam nes <- singleReduce reds,
     isIdentityLambda map_lam || incrementalFlattening =
   distributeSingleStm acc bnd >>= \case
     Just (kernels, res, nest, acc')
@@ -1136,8 +1136,8 @@
       addStmToKernel stm acc
 
   where segmentedConcat nest =
-          isSegmentedOp nest [0] w [] mempty mempty [] (x:xs) $
-          \pat _ _ _ _ _ _ (x':xs') _ ->
+          isSegmentedOp nest [0] w mempty mempty [] (x:xs) $
+          \pat _ _ _ _ (x':xs') _ ->
             let d' = d + length (snd nest) + 1
             in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w
 
@@ -1154,7 +1154,7 @@
       | res == map Var (patternNames $ stmPattern bnd),
         (outer, inners) <- nest,
         [(arr_p, arr)] <- loopNestingParamsAndArrs outer,
-        boundInKernelNest nest `S.intersection` freeInStm bnd
+        boundInKernelNest nest `S.intersection` freeIn bnd
         == S.singleton (paramName arr_p) -> do
           addKernels kernels
           let outerpat = loopNestingPattern $ fst nest
@@ -1241,7 +1241,7 @@
   -- good enough for flatKernel to work.
   let nest' = pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam)
               (MapNesting scatter_pat cs scatter_w $ zip (lambdaParams lam) ivs) nest
-  (nest_bnds, w, ispace, kernel_inps, _rets) <- flatKernel nest'
+  (nest_bnds, w, ispace, kernel_inps) <- flatKernel nest'
 
   let (as_ws, as_ns, as) = unzip3 dests
 
@@ -1290,7 +1290,7 @@
   -- We replicate some of the checking done by 'isSegmentedOp', but
   -- things are different because a GenReduce is not a reduction or
   -- scan.
-  (nest_stms, _, ispace, inputs, _rets) <- flatKernel nest
+  (nest_stms, _, ispace, inputs) <- flatKernel nest
   let orig_pat = Pattern [] $ rearrangeShape perm $
                  patternValueElements $ loopNestingPattern $ fst nest
 
@@ -1313,17 +1313,35 @@
                 -> Certificates -> SubExp -> [SOAC.GenReduceOp SOACS]
                 -> InKernelLambda -> [VName]
                 -> DistribM KernelsStms
-genReduceKernel orig_pat ispace inputs cs genred_w ops lam arrs = do
+genReduceKernel orig_pat ispace inputs cs genred_w ops lam arrs = runBinder_ $ do
   ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> do
-    let (shape, op') = isVectorMap op
-    Out.GenReduceOp num_bins dests nes shape <$> Kernelise.transformLambda op'
+    (op', nes', shape) <- determineReduceOp op nes
+    return $ Out.GenReduceOp num_bins dests nes' shape op'
 
   let isDest = flip elem $ concatMap Out.genReduceDest ops'
       inputs' = filter (not . isDest . kernelInputArray) inputs
 
-  k_stms <- blockedGenReduce orig_pat genred_w ispace inputs' ops' lam arrs
+  certifying cs $
+    addStms =<< segGenRed orig_pat genred_w ispace inputs' ops' lam arrs
 
-  return $ certify cs <$> k_stms
+determineReduceOp :: (MonadBinder m, Lore m ~ Out.Kernels) =>
+                     Lambda -> [SubExp] -> m (Out.Lambda Out.InKernel, [SubExp], Shape)
+determineReduceOp lam nes =
+  -- FIXME? We are assuming that the accumulator is a replicate, and
+  -- we fish out its value in a gross way.
+  case mapM subExpVar nes of
+    Just ne_vs' -> do
+      let (shape, lam') = isVectorMap lam
+      nes' <- forM ne_vs' $ \ne_v -> do
+        ne_v_t <- lookupType ne_v
+        letSubExp "genred_ne" $
+          BasicOp $ Index ne_v $ fullSlice ne_v_t $
+          replicate (shapeRank shape) $ DimFix $ intConst Int32 0
+      lam'' <- Kernelise.transformLambda lam'
+      return (lam'', nes', shape)
+    Nothing -> do
+      lam' <- Kernelise.transformLambda lam
+      return (lam', nes, mempty)
 
 isVectorMap :: Lambda -> (Shape, Lambda)
 isVectorMap lam
@@ -1343,17 +1361,9 @@
                         -> [SubExp] -> [VName]
                         -> KernelM (Maybe KernelsStms)
 segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs =
-  isSegmentedOp nest perm segment_size
-  (lambdaReturnType map_lam) (freeInLambda lam) (freeInLambda map_lam) nes arrs $
-  \pat flat_pat _num_segments total_num_elements ispace inps nes' _ arrs' -> do
-    regularSegmentedScan segment_size flat_pat total_num_elements
-      lam map_lam ispace inps nes' arrs'
-
-    forM_ (zip (patternValueElements pat) (patternNames flat_pat)) $
-      \(dst_pat_elem, flat) -> do
-        let ident = patElemIdent dst_pat_elem
-            dims = arrayDims $ identType ident
-        addStm $ mkLet [] [ident] $ BasicOp $ Reshape (map DimNew dims) flat
+  isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $
+  \pat total_num_elements ispace inps nes' _ _ ->
+    addStms =<< segScan pat total_num_elements segment_size lam map_lam nes' arrs ispace inps
 
 regularSegmentedRedomapKernel :: KernelNest
                               -> [Int]
@@ -1361,27 +1371,24 @@
                               -> InKernelLambda -> InKernelLambda -> [SubExp] -> [VName]
                               -> KernelM (Maybe KernelsStms)
 regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs =
-  isSegmentedOp nest perm segment_size
-    (lambdaReturnType map_lam) (freeInLambda lam) (freeInLambda map_lam) nes arrs $
-    \pat _flat_pat _num_segments total_num_elements ispace inps nes' _ _ ->
-      addStms =<< segRed pat total_num_elements segment_size comm lam map_lam nes' arrs ispace inps
+  isSegmentedOp nest perm segment_size (freeIn lam) (freeIn map_lam) nes arrs $
+    \pat total_num_elements ispace inps nes' _ _ -> do
+      let red_op = SegRedOp comm lam nes' mempty
+      addStms =<< segRed pat total_num_elements segment_size [red_op] map_lam arrs ispace inps
 
 isSegmentedOp :: KernelNest
               -> [Int]
               -> SubExp
-              -> [Type]
               -> Names -> Names
               -> [SubExp] -> [VName]
               -> (Pattern
-                  -> Pattern
                   -> SubExp
-                  -> SubExp
                   -> [(VName, SubExp)]
                   -> [KernelInput]
                   -> [SubExp] -> [VName]  -> [VName]
                   -> Binder Out.Kernels ())
               -> KernelM (Maybe KernelsStms)
-isSegmentedOp nest perm segment_size ret free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do
+isSegmentedOp nest perm segment_size free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do
   -- We must verify that array inputs to the operation are inputs to
   -- the outermost loop nesting or free in the loop nest.  Nothing
   -- free in the op may be bound by the nest.  Furthermore, the
@@ -1392,7 +1399,7 @@
 
   let bound_by_nest = boundInKernelNest nest
 
-  (pre_bnds, nesting_size, ispace, kernel_inps, _rets) <- flatKernel nest
+  (pre_bnds, nesting_size, ispace, kernel_inps) <- flatKernel nest
 
   unless (S.null $ free_in_op `S.intersection` bound_by_nest) $
     fail "Non-fold lambda uses nest-bound parameters."
@@ -1446,13 +1453,8 @@
 
     let pat = Pattern [] $ rearrangeShape perm $
               patternValueElements $ loopNestingPattern $ fst nest
-        flatPatElem pat_elem t = do
-          let t' = arrayOfRow t total_num_elements
-          name <- newVName $ baseString (patElemName pat_elem) ++ "_flat"
-          return $ PatElem name t'
-    flat_pat <- Pattern [] <$> zipWithM flatPatElem (patternValueElements pat) ret
 
-    m pat flat_pat nesting_size total_num_elements ispace kernel_inps nes' nested_arrs arrs'
+    m pat total_num_elements ispace kernel_inps nes' nested_arrs arrs'
 
   where replicateMissing ispace inp = do
           t <- lookupType $ kernelInputArray inp
diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
--- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
+++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs
@@ -1,33 +1,27 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE TypeFamilies #-}
 module Futhark.Pass.ExtractKernels.BlockedKernel
-       ( blockedReduction
-       , blockedReductionStream
-       , blockedGenReduce
-       , blockedMap
-       , blockedScan
-
-       , segRed
+       ( segRed
        , nonSegRed
+       , segScan
+       , segGenRed
 
+       , streamRed
+       , streamMap
+
        , mapKernel
-       , mapKernelFromBody
        , KernelInput(..)
        , readKernelInput
 
-       -- Helper functions shared with at least Segmented.hs
-       , kerneliseLambda
        , newKernelSpace
-       , chunkLambda
-       , splitArrays
        , getSize
        )
        where
 
 import Control.Monad
+import Control.Monad.Writer
 import Data.Maybe
 import Data.List
-import qualified Data.Set as S
 
 import Prelude hiding (quot)
 
@@ -39,12 +33,6 @@
 import Futhark.MonadFreshNames
 import Futhark.Tools
 import Futhark.Transform.Rename
-import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise
-import Futhark.Representation.AST.Attributes.Aliases
-import qualified Futhark.Analysis.Alias as Alias
-import qualified Futhark.Representation.SOACS.SOAC as SOAC
-import Futhark.Util
-import Futhark.Util.IntegralExp
 
 getSize :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) =>
            String -> SizeClass -> m SubExp
@@ -52,216 +40,6 @@
   size_key <- nameFromString . pretty <$> newVName desc
   letSubExp desc $ Op $ GetSize size_key size_class
 
-blockedReductionStream :: (MonadFreshNames m, HasScope Kernels m) =>
-                          Pattern Kernels
-                       -> SubExp
-                       -> Commutativity
-                       -> Lambda InKernel -> Lambda InKernel
-                       -> [(VName, SubExp)] -> [SubExp] -> [VName]
-                       -> m (Stms Kernels)
-blockedReductionStream pat w comm reduce_lam fold_lam ispace nes arrs = runBinder_ $ do
-  (max_step_one_num_groups, step_one_size) <- blockedKernelSize =<< asIntS Int64 w
-
-  let one = constant (1 :: Int32)
-      num_chunks = kernelWorkgroups step_one_size
-
-  let (acc_idents, arr_idents) = splitAt (length nes) $ patternIdents pat
-  step_one_pat <- basicPattern [] <$>
-                  ((++) <$>
-                   mapM (mkIntermediateIdent num_chunks) acc_idents <*>
-                   pure arr_idents)
-  let (_fold_chunk_param, _fold_acc_params, _fold_inp_params) =
-        partitionChunkedFoldParameters (length nes) $ lambdaParams fold_lam
-
-  fold_lam' <- kerneliseLambda nes fold_lam
-
-  my_index <- newVName "my_index"
-  other_index <- newVName "other_index"
-  let my_index_param = Param my_index (Prim int32)
-      other_index_param = Param other_index (Prim int32)
-      reduce_lam' = reduce_lam { lambdaParams = my_index_param :
-                                                other_index_param :
-                                                lambdaParams reduce_lam
-                               }
-      params_to_arrs = zip (map paramName $ drop 1 $ lambdaParams fold_lam') arrs
-      consumedArray v = fromMaybe v $ lookup v params_to_arrs
-      consumed_in_fold =
-        S.map consumedArray $ consumedByLambda $ Alias.analyseLambda fold_lam
-
-  arrs_copies <- forM arrs $ \arr ->
-    if arr `S.member` consumed_in_fold then
-      letExp (baseString arr <> "_copy") $ BasicOp $ Copy arr
-    else return arr
-
-  step_one <- chunkedReduceKernel w step_one_size comm reduce_lam' fold_lam'
-              ispace nes arrs_copies
-  addStm =<< renameStm (Let step_one_pat (defAux ()) $ Op $ HostOp step_one)
-
-  step_two_pat <- basicPattern [] <$>
-                  mapM (mkIntermediateIdent $ constant (1 :: Int32)) acc_idents
-
-  let step_two_size = KernelSize one max_step_one_num_groups one num_chunks max_step_one_num_groups
-
-  step_two <- reduceKernel step_two_size reduce_lam' nes $ take (length nes) $ patternNames step_one_pat
-
-  addStm $ Let step_two_pat (defAux ()) $ Op $ HostOp step_two
-
-  forM_ (zip (patternIdents step_two_pat) (patternIdents pat)) $ \(arr, x) ->
-    addStm $ mkLet [] [x] $ BasicOp $ Index (identName arr) $
-    fullSlice (identType arr) [DimFix $ constant (0 :: Int32)]
-  where mkIntermediateIdent chunk_size ident =
-          newIdent (baseString $ identName ident) $
-          arrayOfRow (identType ident) chunk_size
-
-chunkedReduceKernel :: (MonadBinder m, Lore m ~ Kernels) =>
-                       SubExp
-                    -> KernelSize
-                    -> Commutativity
-                    -> Lambda InKernel -> Lambda InKernel
-                    -> [(VName, SubExp)] -> [SubExp] -> [VName]
-                    -> m (Kernel InKernel)
-chunkedReduceKernel w step_one_size comm reduce_lam' fold_lam' ispace nes arrs = do
-  let ordering = case comm of Commutative -> Disorder
-                              Noncommutative -> InOrder
-      group_size = kernelWorkgroupSize step_one_size
-      num_nonconcat = length nes
-
-  space <- newKernelSpace (kernelWorkgroups step_one_size, group_size, kernelNumThreads step_one_size) $ FlatThreadSpace ispace
-  ((chunk_red_pes, chunk_map_pes), chunk_and_fold) <-
-    runBinder $ blockedPerThread (spaceGlobalId space)
-    w step_one_size ordering fold_lam' num_nonconcat arrs
-  let red_ts = map patElemType chunk_red_pes
-      map_ts = map (rowType . patElemType) chunk_map_pes
-      ts = red_ts ++ map_ts
-      ordering' =
-        case ordering of InOrder -> SplitContiguous
-                         Disorder -> SplitStrided $ kernelNumThreads step_one_size
-
-  chunk_red_pes' <- forM red_ts $ \red_t -> do
-    pe_name <- newVName "chunk_fold_red"
-    return $ PatElem pe_name $ red_t `arrayOfRow` group_size
-  combine_reds <- forM (zip chunk_red_pes' chunk_red_pes) $ \(pe', pe) -> do
-    combine_id <- newVName "combine_id"
-    return $ Let (Pattern [] [pe']) (defAux ()) $ Op $
-      Combine (combineSpace [(combine_id, group_size)]) [patElemType pe] [] $
-      Body () mempty [Var $ patElemName pe]
-
-  final_red_pes <- forM (lambdaReturnType reduce_lam') $ \t -> do
-    pe_name <- newVName "final_result"
-    return $ PatElem pe_name t
-  let reduce_chunk = Let (Pattern [] final_red_pes) (defAux ()) $ Op $
-                     GroupReduce group_size reduce_lam' $
-                     zip nes $ map patElemName chunk_red_pes'
-
-  red_rets <- forM final_red_pes $ \pe ->
-    return $ ThreadsReturn OneResultPerGroup $ Var $ patElemName pe
-  elems_per_thread <- asIntS Int32 $ kernelElementsPerThread step_one_size
-  map_rets <- forM chunk_map_pes $ \pe ->
-    return $ ConcatReturns ordering' w elems_per_thread Nothing $ patElemName pe
-  let rets = red_rets ++ map_rets
-
-  return $ Kernel (KernelDebugHints "chunked_reduce" [("input size", w)]) space ts $
-    KernelBody () (chunk_and_fold<>stmsFromList combine_reds<>oneStm reduce_chunk) rets
-
-reduceKernel :: (MonadBinder m, Lore m ~ Kernels) =>
-                KernelSize
-             -> Lambda InKernel
-             -> [SubExp]
-             -> [VName]
-             -> m (Kernel InKernel)
-reduceKernel step_two_size reduce_lam' nes arrs = do
-  let group_size = kernelWorkgroupSize step_two_size
-      red_ts = lambdaReturnType reduce_lam'
-  space <- newKernelSpace (kernelWorkgroups step_two_size, group_size, kernelNumThreads step_two_size) $
-           FlatThreadSpace []
-  let thread_id = spaceGlobalId space
-
-  (rets, kstms) <- runBinder $ localScope (scopeOfKernelSpace space) $ do
-    in_bounds <- letSubExp "in_bounds" $ BasicOp $ CmpOp (CmpSlt Int32)
-                 (Var $ spaceLocalId space)
-                 (kernelTotalElements step_two_size)
-
-    combine_body <- runBodyBinder $
-      fmap resultBody $ forM (zip arrs nes) $ \(arr, ne) -> do
-        arr_t <- lookupType arr
-        letSubExp "elem" =<<
-          eIf (eSubExp in_bounds)
-          (eBody [pure $ BasicOp $ Index arr $
-                  fullSlice arr_t [DimFix (Var thread_id)]])
-          (resultBodyM [ne])
-
-    combine_pat <- fmap (Pattern []) $ forM (zip arrs red_ts) $ \(arr, red_t) -> do
-      arr' <- newVName $ baseString arr ++ "_combined"
-      return $ PatElem arr' $ red_t `arrayOfRow` group_size
-
-    combine_id <- newVName "combine_id"
-    letBind_ combine_pat $
-      Op $ Combine (combineSpace [(combine_id, group_size)])
-      (map rowType $ patternTypes combine_pat) [] combine_body
-
-    let arrs' = patternNames combine_pat
-
-    final_res_pes <- forM (lambdaReturnType reduce_lam') $ \t -> do
-      pe_name <- newVName "final_result"
-      return $ PatElem pe_name t
-    letBind_ (Pattern [] final_res_pes) $
-      Op $ GroupReduce group_size reduce_lam' $ zip nes arrs'
-
-    forM final_res_pes $ \pe ->
-      return $ ThreadsReturn OneResultPerGroup $ Var $ patElemName pe
-
-  return $ Kernel (KernelDebugHints "reduce" []) space (lambdaReturnType reduce_lam')  $
-    KernelBody () kstms rets
-
--- | Requires a fold lambda that includes accumulator parameters.
-chunkLambda :: (MonadFreshNames m, HasScope Kernels m) =>
-               Pattern Kernels -> [SubExp] -> Lambda InKernel -> m (Lambda InKernel)
-chunkLambda pat nes fold_lam = do
-  chunk_size <- newVName "chunk_size"
-
-  let arr_idents = drop (length nes) $ patternIdents pat
-      (fold_acc_params, fold_arr_params) =
-        splitAt (length nes) $ lambdaParams fold_lam
-      chunk_size_param = Param chunk_size (Prim int32)
-  arr_chunk_params <- mapM (mkArrChunkParam $ Var chunk_size) fold_arr_params
-
-  map_arr_params <- forM arr_idents $ \arr ->
-    newParam (baseString (identName arr) <> "_in") $
-    setOuterSize (identType arr) (Var chunk_size)
-
-  fold_acc_params' <- forM fold_acc_params $ \p ->
-    newParam (baseString $ paramName p) $ paramType p
-
-  let seq_rt =
-        let (acc_ts, arr_ts) =
-              splitAt (length nes) $ lambdaReturnType fold_lam
-        in acc_ts ++ map (`arrayOfRow` Var chunk_size) arr_ts
-
-      res_idents = zipWith Ident (patternValueNames pat) seq_rt
-
-      param_scope =
-        scopeOfLParams $ fold_acc_params' ++ arr_chunk_params ++ map_arr_params
-
-  seq_loop_stms <-
-    runBinder_ $ localScope param_scope $
-    Kernelise.groupStreamMapAccumL
-    (patternElements (basicPattern [] res_idents))
-    (Var chunk_size) fold_lam (map (Var . paramName) fold_acc_params')
-    (map paramName arr_chunk_params)
-
-  let seq_body = mkBody seq_loop_stms $ map (Var . identName) res_idents
-
-  return Lambda { lambdaParams = chunk_size_param :
-                                 fold_acc_params' ++
-                                 arr_chunk_params ++
-                                 map_arr_params
-                , lambdaReturnType = seq_rt
-                , lambdaBody = seq_body
-                }
-  where mkArrChunkParam chunk_size arr_param =
-          newParam (baseString (paramName arr_param) <> "_chunk") $
-            arrayOfRow (paramType arr_param) chunk_size
-
 -- | Given a chunked fold lambda that takes its initial accumulator
 -- value as parameters, bind those parameters to the neutral element
 -- instead.
@@ -285,42 +63,60 @@
                               fold_inp_params
              }
 
-segRed :: (MonadFreshNames m, HasScope Kernels m) =>
-          Pattern Kernels
-       -> SubExp
-       -> SubExp -- segment size
-       -> Commutativity
-       -> Lambda InKernel -> Lambda InKernel
-       -> [SubExp] -> [VName]
-       -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this reduction
-       -> [KernelInput]     -- inps = inputs that can be looked up by using the gtids from ispace
-       -> m (Stms Kernels)
-segRed pat total_num_elements w comm reduce_lam map_lam nes arrs ispace inps = runBinder_ $ do
+prepareRedOrScan :: (MonadBinder m, Lore m ~ Kernels) =>
+                    SubExp -> SubExp
+                 -> LambdaT InKernel
+                 -> [VName] -> [(VName, SubExp)] -> [KernelInput]
+                 -> m (KernelSpace, KernelBody InKernel)
+prepareRedOrScan total_num_elements w map_lam arrs ispace inps = do
   (_, KernelSize num_groups group_size _ _ num_threads) <- blockedKernelSize =<< asIntS Int64 total_num_elements
   gtid <- newVName "gtid"
-  kspace <- newKernelSpace (num_groups, group_size, num_threads) $ FlatThreadSpace $
-            ispace ++ [(gtid, w)]
-  body <- runBodyBinder $ localScope (scopeOfKernelSpace kspace) $ do
+  kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $
+            FlatThreadSpace $ ispace ++ [(gtid, w)]
+  body <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $
+          localScope (scopeOfKernelSpace kspace) $ do
     mapM_ (addStm <=< readKernelInput) inps
     forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do
       arr_t <- lookupType arr
       letBindNames_ [paramName p] $
         BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid]
-    return $ lambdaBody map_lam
+    map ThreadsReturn <$> bodyBind (lambdaBody map_lam)
 
+  return (kspace, body)
+
+segRed :: (MonadFreshNames m, HasScope Kernels m) =>
+          Pattern Kernels
+       -> SubExp
+       -> SubExp -- segment size
+       -> [SegRedOp InKernel]
+       -> Lambda InKernel
+       -> [VName]
+       -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this reduction
+       -> [KernelInput]     -- inps = inputs that can be looked up by using the gtids from ispace
+       -> m (Stms Kernels)
+segRed pat total_num_elements w ops map_lam arrs ispace inps = runBinder_ $ do
+  (kspace, kbody) <- prepareRedOrScan total_num_elements w map_lam arrs ispace inps
   letBind_ pat $ Op $ HostOp $
-    SegRed kspace comm reduce_lam nes (lambdaReturnType map_lam) body
+    SegRed kspace ops (lambdaReturnType map_lam) kbody
 
-nonSegRed :: (MonadFreshNames m, HasScope Kernels m) =>
-             Pattern Kernels
-          -> SubExp
-          -> Commutativity
-          -> Lambda InKernel
-          -> Lambda InKernel
-          -> [SubExp]
-          -> [VName]
-          -> m (Stms Kernels)
-nonSegRed pat w comm red_lam map_lam nes arrs = runBinder_ $ do
+segScan :: (MonadFreshNames m, HasScope Kernels m) =>
+           Pattern Kernels
+        -> SubExp
+        -> SubExp -- segment size
+        -> Lambda InKernel -> Lambda InKernel
+        -> [SubExp] -> [VName]
+        -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this scan
+        -> [KernelInput]     -- inps = inputs that can be looked up by using the gtids from ispace
+        -> m (Stms Kernels)
+segScan pat total_num_elements w scan_lam map_lam nes arrs ispace inps = runBinder_ $ do
+  (kspace, kbody) <- prepareRedOrScan total_num_elements w map_lam arrs ispace inps
+  letBind_ pat $ Op $ HostOp $
+    SegScan kspace scan_lam nes (lambdaReturnType map_lam) kbody
+
+dummyDim :: (MonadFreshNames m, MonadBinder m) =>
+            Pattern Kernels
+         -> m (Pattern Kernels, [(VName, SubExp)], m ())
+dummyDim pat = do
   -- We add a unit-size segment on top to ensure that the result
   -- of the SegRed is an array, which we then immediately index.
   -- This is useful in the case that the value is used on the
@@ -330,42 +126,115 @@
   let addDummyDim t = t `arrayOfRow` intConst Int32 1
   pat' <- fmap addDummyDim <$> renamePattern pat
   dummy <- newVName "dummy"
-  addStms =<<
-    segRed pat' w w comm red_lam map_lam nes arrs [(dummy, intConst Int32 1)] []
+  let ispace = [(dummy, intConst Int32 1)]
 
-  forM_ (zip (patternNames pat') (patternNames pat)) $ \(from, to) -> do
-    from_t <- lookupType from
-    letBindNames_ [to] $ BasicOp $ Index from $ fullSlice from_t [DimFix $ intConst Int32 0]
+  return (pat', ispace,
+          forM_ (zip (patternNames pat') (patternNames pat)) $ \(from, to) -> do
+             from_t <- lookupType from
+             letBindNames_ [to] $ BasicOp $ Index from $
+               fullSlice from_t [DimFix $ intConst Int32 0])
 
-blockedReduction :: (MonadFreshNames m, HasScope Kernels m) =>
-                    Pattern Kernels
-                 -> SubExp
-                 -> Commutativity
-                 -> Lambda InKernel -> Lambda InKernel
-                 -> [(VName, SubExp)] -> [SubExp] -> [VName]
-                 -> m (Stms Kernels)
-blockedReduction pat w comm reduce_lam map_lam ispace nes arrs = runBinder_ $ do
-  fold_lam <- SOAC.composeLambda SOAC.nilFn reduce_lam map_lam
-  fold_lam' <- chunkLambda pat nes fold_lam
+nonSegRed :: (MonadFreshNames m, HasScope Kernels m) =>
+             Pattern Kernels
+          -> SubExp
+          -> [SegRedOp InKernel]
+          -> Lambda InKernel
+          -> [VName]
+          -> m (Stms Kernels)
+nonSegRed pat w ops map_lam arrs = runBinder_ $ do
+  (pat', ispace, read_dummy) <- dummyDim pat
+  addStms =<< segRed pat' w w ops map_lam arrs ispace []
+  read_dummy
 
-  let arr_idents = drop (length nes) $ patternIdents pat
-  map_out_arrs <- forM arr_idents $ \(Ident name t) ->
-    letExp (baseString name <> "_out_in") $
-    BasicOp $ Scratch (elemType t) (arrayDims t)
+prepareStream :: (MonadBinder m, Lore m ~ Kernels) =>
+                 KernelSize
+              -> [(VName, SubExp)]
+              -> SubExp
+              -> Commutativity
+              -> Lambda InKernel
+              -> [SubExp]
+              -> [VName]
+              -> m (KernelSpace, [Type], KernelBody InKernel)
+prepareStream size ispace w comm fold_lam nes arrs = do
+  let (KernelSize num_groups group_size elems_per_thread _ num_threads) = size
+  let (ordering, split_ordering) =
+        case comm of Commutative -> (Disorder, SplitStrided num_threads)
+                     Noncommutative -> (InOrder, SplitContiguous)
 
-  addStms =<<
-    blockedReductionStream pat w comm reduce_lam fold_lam'
-    ispace nes (arrs ++ map_out_arrs)
+  fold_lam' <- kerneliseLambda nes fold_lam
 
-blockedGenReduce :: (MonadFreshNames m, HasScope Kernels m) =>
-                    Pattern Kernels
-                 -> SubExp
-                 -> [(VName,SubExp)] -- ^ Segment indexes and sizes.
-                 -> [KernelInput]
-                 -> [GenReduceOp InKernel]
-                 -> Lambda InKernel -> [VName]
-                 -> m (Stms Kernels)
-blockedGenReduce pat arr_w ispace inps ops lam arrs = runBinder_ $ do
+  elems_per_thread_32 <- asIntS Int32 elems_per_thread
+
+  gtid <- newVName "gtid"
+  kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $
+            FlatThreadSpace $ ispace ++ [(gtid, num_threads)]
+  kbody <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $
+           localScope (scopeOfKernelSpace kspace) $ do
+    (chunk_red_pes, chunk_map_pes) <-
+      blockedPerThread gtid w size ordering fold_lam' (length nes) arrs
+    let concatReturns pe =
+          ConcatReturns split_ordering w elems_per_thread_32 Nothing $ patElemName pe
+    return (map (ThreadsReturn . Var . patElemName) chunk_red_pes ++
+            map concatReturns chunk_map_pes)
+
+  let (redout_ts, mapout_ts) = splitAt (length nes) $ lambdaReturnType fold_lam
+      ts = redout_ts ++ map rowType mapout_ts
+
+  return (kspace, ts, kbody)
+
+streamRed :: (MonadFreshNames m, HasScope Kernels m) =>
+             Pattern Kernels
+          -> SubExp
+          -> Commutativity
+          -> Lambda InKernel -> Lambda InKernel
+          -> [SubExp] -> [VName]
+          -> m (Stms Kernels)
+streamRed pat w comm red_lam fold_lam nes arrs = runBinder_ $ do
+  -- The strategy here is to rephrase the stream reduction as a
+  -- non-segmented SegRed that does explicit chunking within its body.
+  -- First, figure out how many threads to use for this.
+  (_, size) <- blockedKernelSize =<< asIntS Int64 w
+
+  let (redout_pes, mapout_pes) = splitAt (length nes) $ patternElements pat
+  (redout_pat, ispace, read_dummy) <- dummyDim $ Pattern [] redout_pes
+  let pat' = Pattern [] $ patternElements redout_pat ++ mapout_pes
+
+  (kspace, ts, kbody) <- prepareStream size ispace w comm fold_lam nes arrs
+
+  letBind_ pat' $ Op $ HostOp $ SegRed kspace
+    [SegRedOp comm red_lam nes mempty] ts kbody
+
+  read_dummy
+
+-- Similar to streamRed, but without the last reduction.
+streamMap :: (MonadFreshNames m, HasScope Kernels m) =>
+              [String] -> [PatElem Kernels] -> SubExp
+           -> Commutativity -> Lambda InKernel -> [SubExp] -> [VName]
+           -> m ((SubExp, [VName]), Stms Kernels)
+streamMap out_desc mapout_pes w comm fold_lam nes arrs = runBinder $ do
+  (_, size) <- blockedKernelSize =<< asIntS Int64 w
+
+  (kspace, ts, kbody) <- prepareStream size [] w comm fold_lam nes arrs
+
+  let redout_ts = take (length nes) ts
+
+  redout_pes <- forM (zip out_desc redout_ts) $ \(desc, t) ->
+    PatElem <$> newVName desc <*> pure (t `arrayOfRow` spaceNumThreads kspace)
+
+  let pat = Pattern [] $ redout_pes ++ mapout_pes
+  letBind_ pat $ Op $ HostOp $ SegMap kspace ts kbody
+
+  return (spaceNumThreads kspace, map patElemName redout_pes)
+
+segGenRed :: (MonadFreshNames m, HasScope Kernels m) =>
+             Pattern Kernels
+          -> SubExp
+          -> [(VName,SubExp)] -- ^ Segment indexes and sizes.
+          -> [KernelInput]
+          -> [GenReduceOp InKernel]
+          -> Lambda InKernel -> [VName]
+          -> m (Stms Kernels)
+segGenRed pat arr_w ispace inps ops lam arrs = runBinder_ $ do
   let (_, segment_sizes) = unzip ispace
   arr_w_64 <- letSubExp "arr_w_64" =<< eConvOp (SExt Int32 Int64) (toExp arr_w)
   segment_sizes_64 <- mapM (letSubExp "segment_size_64" <=< eConvOp (SExt Int32 Int64) . toExp) segment_sizes
@@ -374,55 +243,19 @@
     blockedKernelSize total_w
 
   gtid <- newVName "gtid"
-  kspace <- newKernelSpace (num_groups, group_size, num_threads) $
+  kspace <- newKernelSpace (num_groups, group_size, num_threads, num_groups) $
             FlatThreadSpace $ ispace ++ [(gtid, arr_w)]
 
-  body <- runBodyBinder $ localScope (scopeOfKernelSpace kspace) $ do
+  kbody <- fmap (uncurry (flip $ KernelBody ())) $ runBinder $
+          localScope (scopeOfKernelSpace kspace) $ do
     mapM_ (addStm <=< readKernelInput) inps
     forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do
       arr_t <- lookupType arr
       letBindNames_ [paramName p] $
         BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid]
-    return $ lambdaBody lam
-
-  letBind_ pat $ Op $ HostOp $ SegGenRed kspace ops (lambdaReturnType lam) body
-
-blockedMap :: (MonadFreshNames m, HasScope Kernels m) =>
-              Pattern Kernels -> SubExp
-           -> StreamOrd -> Lambda InKernel -> [SubExp] -> [VName]
-           -> m (Stm Kernels, Stms Kernels)
-blockedMap concat_pat w ordering lam nes arrs = runBinder $ do
-  (_, kernel_size) <- blockedKernelSize =<< asIntS Int64 w
-  let num_nonconcat = length (lambdaReturnType lam) - patternSize concat_pat
-      num_groups = kernelWorkgroups kernel_size
-      group_size = kernelWorkgroupSize kernel_size
-      num_threads = kernelNumThreads kernel_size
-      ordering' =
-        case ordering of InOrder -> SplitContiguous
-                         Disorder -> SplitStrided $ kernelNumThreads kernel_size
-
-  space <- newKernelSpace (num_groups, group_size, num_threads) (FlatThreadSpace [])
-  lam' <- kerneliseLambda nes lam
-  ((chunk_red_pes, chunk_map_pes), chunk_and_fold) <- runBinder $
-    blockedPerThread (spaceGlobalId space) w kernel_size ordering lam' num_nonconcat arrs
-
-  nonconcat_pat <-
-    fmap (Pattern []) $ forM (take num_nonconcat $ lambdaReturnType lam) $ \t -> do
-      name <- newVName "nonconcat"
-      return $ PatElem name $ t `arrayOfRow` num_threads
-
-  let pat = nonconcat_pat <> concat_pat
-      ts = map patElemType chunk_red_pes ++
-           map (rowType . patElemType) chunk_map_pes
-
-  nonconcat_rets <- forM chunk_red_pes $ \pe ->
-    return $ ThreadsReturn AllThreads $ Var $ patElemName pe
-  elems_per_thread <- asIntS Int32 $ kernelElementsPerThread kernel_size
-  concat_rets <- forM chunk_map_pes $ \pe ->
-    return $ ConcatReturns ordering' w elems_per_thread Nothing $ patElemName pe
+    map ThreadsReturn <$> bodyBind (lambdaBody lam)
 
-  return $ Let pat (defAux ()) $ Op $ HostOp $ Kernel (KernelDebugHints "chunked_map" []) space ts $
-    KernelBody () chunk_and_fold $ nonconcat_rets ++ concat_rets
+  letBind_ pat $ Op $ HostOp $ SegGenRed kspace ops (lambdaReturnType lam) kbody
 
 blockedPerThread :: (MonadBinder m, Lore m ~ InKernel) =>
                     VName -> SubExp -> KernelSize -> StreamOrd -> Lambda InKernel
@@ -532,439 +365,58 @@
   return (max_num_groups,
           KernelSize num_groups' group_size per_thread_elements w num_threads')
 
--- First stage scan kernel.
-scanKernel1 :: (MonadBinder m, Lore m ~ Kernels) =>
-               SubExp -> KernelSize
-            -> SOAC.Scan InKernel
-            -> SOAC.Reduce InKernel
-            -> Lambda InKernel -> [VName]
-            -> m (Kernel InKernel)
-scanKernel1 w scan_sizes (scan_lam, scan_nes) (_comm, red_lam, red_nes) foldlam arrs = do
-  num_elements <- asIntS Int32 $ kernelTotalElements scan_sizes
-
-  let (scan_ts, red_ts, map_ts) =
-        splitAt3 (length scan_nes) (length red_nes) $ lambdaReturnType foldlam
-      (_, foldlam_acc_params, _) =
-        partitionChunkedFoldParameters (length scan_nes + length red_nes) $ lambdaParams foldlam
-
-  -- Scratch arrays for scanout and mapout parts.
-  (scanout_arrs, scanout_arr_params, scanout_arr_ts) <-
-    unzip3 <$> mapM (mkOutArray "scanout") scan_ts
-  (mapout_arrs, mapout_arr_params, mapout_arr_ts) <-
-    unzip3 <$> mapM (mkOutArray "scanout") map_ts
-
-  last_thread <- letSubExp "last_thread" $ BasicOp $
-                 BinOp (Sub Int32) group_size (constant (1::Int32))
-  kspace <- newKernelSpace (num_groups, group_size, num_threads) $ FlatThreadSpace []
-  let lid = spaceLocalId kspace
-
-  (res, stms) <- runBinder $ localScope (scopeOfKernelSpace kspace) $ do
-    -- We create a loop that moves in group_size chunks over the input.
-    num_iterations <- letSubExp "num_iterations" =<<
-                      eDivRoundingUp Int32 (eSubExp w) (eSubExp num_threads)
-
-    -- The merge parameters are the scanout arrays, the reduction
-    -- results, the mapout arrays, and the (renamed) scan accumulator
-    -- parameters of foldlam (which function as carries).  We do not
-    -- need to keep accumulator parameters/carries for the reduction,
-    -- because the reduction result suffices.
-    (acc_params, nes') <- unzip <$> zipWithM mkAccMergeParam foldlam_acc_params
-                          (scan_nes ++ red_nes)
-    let (scan_acc_params, red_acc_params) =
-          splitAt (length scan_nes) acc_params
-        (scan_nes', red_nes') =
-          splitAt (length scan_nes) nes'
-    let merge = zip scanout_arr_params (map Var scanout_arrs) ++
-                zip red_acc_params red_nes' ++
-                zip mapout_arr_params (map Var mapout_arrs) ++
-                zip scan_acc_params scan_nes'
-    i <- newVName "i"
-    let form = ForLoop i Int32 num_iterations []
-
-    loop_body <- runBodyBinder $ localScope (scopeOfFParams (map fst merge) <>
-                                             scopeOf form) $ do
-      -- Compute the offset into the input and output.  To this a
-      -- thread can add its local ID to figure out which element it is
-      -- responsible for.
-      offset <- letSubExp "offset" =<<
-                eBinOp (Add Int32)
-                (eBinOp (Mul Int32)
-                 (eSubExp $ Var $ spaceGroupId kspace)
-                 (pure $ BasicOp $ BinOp (Mul Int32) num_iterations group_size))
-                (pure $ BasicOp $ BinOp (Mul Int32) (Var i) group_size)
-
-      -- Now we apply the fold function if j=offset+lid is less than
-      -- num_elements.  This also involves writing to the mapout
-      -- arrays.
-      j <- letSubExp "j" $ BasicOp $ BinOp (Add Int32) offset (Var lid)
-      let in_bounds = pure $ BasicOp $ CmpOp (CmpSlt Int32) j num_elements
-
-          in_bounds_fold_branch = do
-            -- Read array input.
-            arr_elems <- forM arrs $ \arr -> do
-              arr_t <- lookupType arr
-              let slice = fullSlice arr_t [DimFix j]
-              letSubExp (baseString arr ++ "_elem") $ BasicOp $ Index arr slice
-
-            -- Apply the body of the fold function.
-            fold_res <-
-              eLambda foldlam $ map eSubExp $ j : map (Var . paramName) acc_params ++ arr_elems
-
-            -- Scatter the to_map parts to the mapout arrays using
-            -- in-place updates, and return the to_scan parts.
-            let (to_scan, to_red, to_map) = splitAt3 (length scan_nes) (length red_nes) fold_res
-            mapout_arrs' <- forM (zip to_map mapout_arr_params) $ \(se,arr) -> do
-              let slice = fullSlice (paramType arr) [DimFix j]
-              letInPlace "mapout" (paramName arr) slice $ BasicOp $ SubExp se
-            return $ resultBody $ to_scan ++ to_red ++ map Var mapout_arrs'
-
-          not_in_bounds_fold_branch = return $ resultBody $ map (Var . paramName) $
-                                      scan_acc_params ++ red_acc_params ++ mapout_arr_params
-
-      (to_scan_res, to_red_res, mapout_arrs') <-
-        fmap (splitAt3 (length scan_nes) (length red_nes)) . letTupExp "foldres" =<<
-        eIf in_bounds in_bounds_fold_branch not_in_bounds_fold_branch
-
-      (scanned_arrs, scanout_arrs') <-
-        doScan j kspace in_bounds scanout_arr_params to_scan_res
-
-      new_scan_carries <-
-        resetCarries "scan" lid scan_acc_params scan_nes' $ runBodyBinder $ do
-          carries <- forM scanned_arrs $ \arr -> do
-            arr_t <- lookupType arr
-            let slice = fullSlice arr_t [DimFix last_thread]
-            letSubExp "carry" $ BasicOp $ Index arr slice
-          return $ resultBody carries
-
-      red_res <- doReduce to_red_res
-
-      new_red_carries <- resetCarries "red" lid red_acc_params red_nes' $
-                         return $ resultBody $ map Var red_res
-
-      -- HACK
-      new_scan_carries' <- letTupExp "new_carry_sync" $ Op $ Barrier $ map Var new_scan_carries
-      return $ resultBody $ map Var $
-        scanout_arrs' ++ new_red_carries ++ mapout_arrs' ++ new_scan_carries'
-
-    result <- letTupExp "result" $ DoLoop [] merge form loop_body
-    let (scanout_result, red_result, mapout_result, scan_carry_result) =
-          splitAt4 (length scan_ts) (length red_ts) (length mapout_arrs) result
-    return (map KernelInPlaceReturn scanout_result ++
-            map (ThreadsReturn OneResultPerGroup . Var) scan_carry_result ++
-            map (ThreadsReturn OneResultPerGroup . Var) red_result ++
-            map KernelInPlaceReturn mapout_result)
-
-  let kts = scanout_arr_ts ++ scan_ts ++ red_ts ++ mapout_arr_ts
-      kbody = KernelBody () stms res
-
-  return $ Kernel (KernelDebugHints "scan1" []) kspace kts kbody
-  where num_groups = kernelWorkgroups scan_sizes
-        group_size = kernelWorkgroupSize scan_sizes
-        num_threads = kernelNumThreads scan_sizes
-        consumed_in_foldlam = consumedInBody $ lambdaBody $ Alias.analyseLambda foldlam
-
-        mkOutArray desc t = do
-          let arr_t = t `arrayOfRow` w
-          arr <- letExp desc $ BasicOp $ Scratch (elemType arr_t) (arrayDims arr_t)
-          pname <- newVName $ desc++"param"
-          return (arr, Param pname $ toDecl arr_t Unique, arr_t)
-
-        mkAccMergeParam (Param pname ptype) se = do
-          pname' <- newVName $ baseString pname ++ "_merge"
-          -- We have to copy the initial merge parameter (the neutral
-          -- element) if it is consumed inside the lambda.
-          case se of
-            Var v | pname `S.member` consumed_in_foldlam -> do
-                      se' <- letSubExp "scan_ne_copy" $ BasicOp $ Copy v
-                      return (Param pname' $ toDecl ptype Unique,
-                              se')
-            _ -> return (Param pname' $ toDecl ptype Nonunique,
-                         se)
-
-        doScan j kspace in_bounds scanout_arr_params to_scan_res = do
-          let lid = spaceLocalId kspace
-              scan_ts = map (rowType . paramType) scanout_arr_params
-          -- Create an array of per-thread fold results and scan it.
-          combine_id <- newVName "combine_id"
-          to_scan_arrs <- letTupExp "combined" $
-                          Op $ Combine (combineSpace [(combine_id, group_size)]) scan_ts [] $
-                          Body () mempty $ map Var to_scan_res
-          scanned_arrs <- letTupExp "scanned" $
-                          Op $ GroupScan group_size scan_lam $ zip scan_nes to_scan_arrs
-
-          -- If we are in bounds, we write scanned_arrs[lid] to scanout[j].
-          let in_bounds_scan_branch = do
-                -- Read scanned_arrs[j].
-                arr_elems <- forM scanned_arrs $ \arr -> do
-                  arr_t <- lookupType arr
-                  let slice = fullSlice arr_t [DimFix $ Var lid]
-                  letSubExp (baseString arr ++ "_elem") $ BasicOp $ Index arr slice
-
-                -- Scatter the to_map parts to the scanout arrays using
-                -- in-place updates.
-                scanout_arrs' <- forM (zip arr_elems scanout_arr_params) $ \(se,p) -> do
-                  let slice = fullSlice (paramType p) [DimFix j]
-                  letInPlace "mapout" (paramName p) slice $ BasicOp $ SubExp se
-                return $ resultBody $ map Var scanout_arrs'
-
-              not_in_bounds_scan_branch =
-                return $ resultBody $ map (Var . paramName) scanout_arr_params
-
-          scanres <- letTupExp "scanres" =<<
-                     eIf in_bounds in_bounds_scan_branch not_in_bounds_scan_branch
-          return (scanned_arrs, scanres)
-
-        doReduce to_red_res = do
-          red_ts <- mapM lookupType to_red_res
-
-          -- Create an array of per-thread fold results and reduce it.
-          combine_id <- newVName "combine_id"
-          to_red_arrs <- letTupExp "combined" $
-                         Op $ Combine (combineSpace [(combine_id, group_size)]) red_ts [] $
-                         Body () mempty $ map Var to_red_res
-          letTupExp "reduced" $
-            Op $ GroupReduce group_size red_lam $ zip red_nes to_red_arrs
-
-        resetCarries what lid acc_params nes mk_read_res = do
-          -- All threads but the first in the group reset the accumulator
-          -- to the neutral element.  The first resets it to the carry-out
-          -- of the scan or reduction.
-          is_first_thread <- letSubExp "is_first_thread" $ BasicOp $
-                             CmpOp (CmpEq int32) (Var lid) (constant (0::Int32))
-
-          read_res <- mk_read_res
-
-          reset_carry_outs <- runBodyBinder $ do
-            carries <- forM (zip acc_params nes) $ \(p, se) ->
-              case se of
-                Var v | unique $ declTypeOf p ->
-                        letSubExp "reset_acc_copy" $ BasicOp $ Copy v
-                _ -> return se
-            return $ resultBody carries
-
-          letTupExp ("new_" ++ what ++ "_carry") $
-            If is_first_thread read_res reset_carry_outs $
-            ifCommon $ map paramType acc_params
-
--- Second stage scan kernel with no fold part.
-scanKernel2 :: (MonadBinder m, Lore m ~ Kernels) =>
-               KernelSize
-            -> Lambda InKernel
-            -> [(SubExp,VName)]
-            -> m (Kernel InKernel)
-scanKernel2 scan_sizes lam input = do
-  let (nes, arrs) = unzip input
-      scan_ts = lambdaReturnType lam
-
-  kspace <- newKernelSpace (kernelWorkgroups scan_sizes,
-                            group_size,
-                            kernelNumThreads scan_sizes) (FlatThreadSpace [])
-  (res, stms) <- runBinder $ localScope (scopeOfKernelSpace kspace) $ do
-    -- Create an array of the elements we are to scan.
-    let indexMine cid arr = do
-          arr_t <- lookupType arr
-          let slice = fullSlice arr_t [DimFix $ Var cid]
-          letSubExp (baseString arr <> "_elem") $ BasicOp $ Index arr slice
-    combine_id <- newVName "combine_id"
-    read_elements <- runBodyBinder $ resultBody <$> mapM (indexMine combine_id) arrs
-    to_scan_arrs <- letTupExp "combined" $
-                    Op $ Combine (combineSpace [(combine_id, group_size)]) scan_ts [] read_elements
-    scanned_arrs <- letTupExp "scanned" $
-                    Op $ GroupScan group_size lam $ zip nes to_scan_arrs
-
-    -- Each thread returns scanned_arrs[i].
-    res_elems <- mapM (indexMine $ spaceLocalId kspace) scanned_arrs
-    return $ map (ThreadsReturn AllThreads) res_elems
-
-  return $ Kernel (KernelDebugHints "scan2" []) kspace (lambdaReturnType lam) $ KernelBody () stms res
-  where group_size = kernelWorkgroupSize scan_sizes
-
--- | The 'VName's returned are the names of variables bound to the
--- carry-out of the last thread.  You can ignore them if you don't
--- need them.
-blockedScan :: (MonadBinder m, Lore m ~ Kernels) =>
-               Pattern Kernels -> SubExp
-            -> SOAC.Scan InKernel
-            -> SOAC.Reduce InKernel
-            -> Lambda InKernel -> SubExp -> [(VName, SubExp)] -> [KernelInput]
-            -> [VName]
-            -> m [VName]
-blockedScan pat w (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam segment_size ispace inps arrs = do
-  foldlam <- SOAC.composeLambda scan_lam red_lam map_lam
-
-  (_, first_scan_size) <- blockedKernelSize =<< asIntS Int64 w
-  my_index <- newVName "my_index"
-  other_index <- newVName "other_index"
-  let num_groups = kernelWorkgroups first_scan_size
-      group_size = kernelWorkgroupSize first_scan_size
-      num_threads = kernelNumThreads first_scan_size
-      my_index_param = Param my_index (Prim int32)
-      other_index_param = Param other_index (Prim int32)
-
-  let foldlam_scope = scopeOfLParams $ my_index_param : lambdaParams foldlam
-      bindIndex i v = letBindNames_ [i] =<< toExp v
-  compute_segments <- runBinder_ $ localScope foldlam_scope $
-                      zipWithM_ bindIndex (map fst ispace) $
-                      unflattenIndex (map (primExpFromSubExp int32 . snd) ispace)
-                      (LeafExp (paramName my_index_param) int32 `quot`
-                       primExpFromSubExp int32 segment_size)
-  read_inps <- stmsFromList <$> mapM readKernelInput inps
-  first_scan_foldlam <- renameLambda
-    foldlam { lambdaParams = my_index_param :
-                             lambdaParams foldlam
-            , lambdaBody = insertStms (compute_segments<>read_inps) $
-                           lambdaBody foldlam
-            }
-  first_scan_lam <- renameLambda
-    scan_lam { lambdaParams = my_index_param :
-                              other_index_param :
-                              lambdaParams scan_lam
-        }
-  first_scan_red_lam <- renameLambda
-    red_lam { lambdaParams = my_index_param :
-                             other_index_param :
-                             lambdaParams red_lam
-            }
-
-  let (scan_idents, red_idents, arr_idents) =
-        splitAt3 (length scan_nes) (length red_nes) $ patternIdents pat
-      final_res_pat = Pattern [] $ take (length scan_nes) $ patternValueElements pat
-  first_scan_pat <- basicPattern [] . concat <$>
-    sequence [mapM (mkIntermediateIdent "seq_scanned" [w]) scan_idents,
-              mapM (mkIntermediateIdent "scan_carry_out" [num_groups]) scan_idents,
-              mapM (mkIntermediateIdent "red_carry_out" [num_groups]) red_idents,
-              pure arr_idents]
-
-  addStm . Let first_scan_pat (defAux ()) . Op . HostOp =<< scanKernel1 w first_scan_size
-    (first_scan_lam, scan_nes)
-    (comm, first_scan_red_lam, red_nes)
-    first_scan_foldlam arrs
-
-  let (sequentially_scanned, group_carry_out, group_red_res, _) =
-        splitAt4 (length scan_nes) (length scan_nes) (length red_nes) $ patternNames first_scan_pat
-
-  let second_scan_size = KernelSize one num_groups one num_groups num_groups
-  unless (null group_red_res) $ do
-    second_stage_red_lam <- renameLambda first_scan_red_lam
-    red_res <- letTupExp "red_res" . Op . HostOp =<<
-               reduceKernel second_scan_size second_stage_red_lam red_nes group_red_res
-    forM_ (zip red_idents red_res) $ \(dest, arr) -> do
-      arr_t <- lookupType arr
-      addStm $ mkLet [] [dest] $ BasicOp $ Index arr $
-        fullSlice arr_t [DimFix $ constant (0 :: Int32)]
-
-  second_scan_lam <- renameLambda first_scan_lam
-
-  group_carry_out_scanned <-
-    letTupExp "group_carry_out_scanned" . Op . HostOp =<<
-    scanKernel2 second_scan_size
-    second_scan_lam (zip scan_nes group_carry_out)
-
-  last_group <- letSubExp "last_group" $ BasicOp $ BinOp (Sub Int32) num_groups one
-  carries <- forM group_carry_out_scanned $ \carry_outs -> do
-    arr_t <- lookupType carry_outs
-    letExp "carry_out" $ BasicOp $ Index carry_outs $ fullSlice arr_t [DimFix last_group]
-
-  scan_lam''' <- renameLambda scan_lam
-  j <- newVName "j"
-  let (acc_params, arr_params) =
-        splitAt (length scan_nes) $ lambdaParams scan_lam'''
-      result_map_input =
-        zipWith (mkKernelInput [Var j]) arr_params sequentially_scanned
-
-  chunks_per_group <- letSubExp "chunks_per_group" =<<
-    eDivRoundingUp Int32 (eSubExp w) (eSubExp num_threads)
-  elems_per_group <- letSubExp "elements_per_group" $
-    BasicOp $ BinOp (Mul Int32) chunks_per_group group_size
-
-  result_map_body <- runBodyBinder $ localScope (scopeOfLParams $ map kernelInputParam result_map_input) $ do
-    group_id <-
-      letSubExp "group_id" $
-      BasicOp $ BinOp (SQuot Int32) (Var j) elems_per_group
-    let do_nothing =
-          pure $ resultBody $ map (Var . paramName) arr_params
-        add_carry_in = runBodyBinder $ do
-          forM_ (zip acc_params group_carry_out_scanned) $ \(p, arr) -> do
-            carry_in_index <-
-              letSubExp "carry_in_index" $
-              BasicOp $ BinOp (Sub Int32) group_id one
-            arr_t <- lookupType arr
-            letBindNames_ [paramName p] $
-              BasicOp $ Index arr $ fullSlice arr_t [DimFix carry_in_index]
-          return $ lambdaBody scan_lam'''
-    group_lasts <-
-      letTupExp "final_result" =<<
-        eIf (eCmpOp (CmpEq int32) (eSubExp zero) (eSubExp group_id))
-        do_nothing
-        add_carry_in
-    return $ resultBody $ map Var group_lasts
-
-  (mapk_bnds, mapk) <- mapKernelFromBody w (FlatThreadSpace [(j, w)]) result_map_input
-                       (lambdaReturnType scan_lam) result_map_body
-  addStms mapk_bnds
-  letBind_ final_res_pat $ Op $ HostOp mapk
-
-  return carries
-  where one = constant (1 :: Int32)
-        zero = constant (0 :: Int32)
-
-        mkIntermediateIdent desc shape ident =
-          newIdent (baseString (identName ident) ++ "_" ++ desc) $
-          arrayOf (rowType $ identType ident) (Shape shape) NoUniqueness
-
-        mkKernelInput indices p arr = KernelInput { kernelInputName = paramName p
-                                                  , kernelInputType = paramType p
-                                                  , kernelInputArray = arr
-                                                  , kernelInputIndices = indices
-                                                  }
+createsArrays :: KernelBody InKernel -> Bool
+createsArrays = getAny . execWriter . mapM_ onStm . kernelBodyStms
+  where onStm stm = do
+          when (any (not . primType) $ patternTypes $ stmPattern stm) $ tell $ Any True
+          walkExpM walker $ stmExp stm
+        walker = identityWalker { walkOnBody = mapM_ onStm . bodyStms }
 
 mapKernelSkeleton :: (HasScope Kernels m, MonadFreshNames m) =>
-                     SubExp -> SpaceStructure -> [KernelInput]
+                     SubExp -> SpaceStructure -> [KernelInput] -> Bool
                   -> m (KernelSpace,
                         Stms Kernels,
                         Stms InKernel)
-mapKernelSkeleton w ispace inputs = do
-  ((group_size, num_threads, num_groups), ksize_bnds) <-
-    runBinder $ numThreadsAndGroups w
+mapKernelSkeleton w ispace inputs creates_arrays = do
+  ((group_size, num_threads, num_groups, virt_groups), ksize_bnds) <- runBinder $
+    -- If the kernel creates arrays internally (meaning it will
+    -- require memory expansion), we want to truncate the amount of
+    -- threads.  Otherwise, have at it!  This is a bit of a hack - in
+    -- principle, we should make this decision later, when we have a
+    -- clearer idea of what is happening inside the kernel.
+    if not creates_arrays then do
+      group_size <- getSize "group_size" SizeGroup
+      num_groups <- letSubExp "num_groups" =<< eDivRoundingUp Int32
+                    (eSubExp w) (eSubExp group_size)
+      num_threads <- letSubExp "num_threads" $
+        BasicOp $ BinOp (Mul Int32) num_groups group_size
+      return (group_size, num_threads, num_groups, num_groups)
 
+      else do
+      (_, ksize) <- blockedKernelSize =<< asIntS Int64 w
+      virt_groups <- letSubExp "virt_groups" =<< eDivRoundingUp Int32
+                     (eSubExp w) (eSubExp (kernelWorkgroupSize ksize))
+      return (kernelWorkgroupSize ksize, kernelNumThreads ksize,
+              kernelWorkgroups ksize, virt_groups)
+
   read_input_bnds <- stmsFromList <$> mapM readKernelInput inputs
 
-  let ksize = (num_groups, group_size, num_threads)
+  let ksize = (num_groups, group_size, num_threads, virt_groups)
 
   space <- newKernelSpace ksize ispace
   return (space, ksize_bnds, read_input_bnds)
 
--- Given the desired minium number of threads, compute the group size,
--- number of groups and total number of threads.
-numThreadsAndGroups :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) =>
-                       SubExp -> m (SubExp, SubExp, SubExp)
-numThreadsAndGroups w = do
-  group_size <- getSize "group_size" SizeGroup
-  num_groups <- letSubExp "num_groups" =<< eDivRoundingUp Int32
-    (eSubExp w) (eSubExp group_size)
-  num_threads <- letSubExp "num_threads" $
-    BasicOp $ BinOp (Mul Int32) num_groups group_size
-  return (group_size, num_threads, num_groups)
-
 mapKernel :: (HasScope Kernels m, MonadFreshNames m) =>
              SubExp -> SpaceStructure -> [KernelInput]
           -> [Type] -> KernelBody InKernel
           -> m (Stms Kernels, Kernel InKernel)
-mapKernel w ispace inputs rts (KernelBody () kstms krets) = do
-  (space, ksize_bnds, read_input_bnds) <- mapKernelSkeleton w ispace inputs
+mapKernel w ispace inputs rts kbody@(KernelBody () kstms krets) = do
+  (space, ksize_bnds, read_input_bnds) <- mapKernelSkeleton w ispace inputs $
+                                          createsArrays kbody
 
   let kbody' = KernelBody () (read_input_bnds <> kstms) krets
   return (ksize_bnds, Kernel (KernelDebugHints "map" []) space rts kbody')
 
-mapKernelFromBody :: (HasScope Kernels m, MonadFreshNames m) =>
-                     SubExp -> SpaceStructure -> [KernelInput]
-                  -> [Type] -> Body InKernel
-                  -> m (Stms Kernels, Kernel InKernel)
-mapKernelFromBody w ispace inputs rts body =
-  mapKernel w ispace inputs rts kbody
-  where kbody = KernelBody () (bodyStms body) krets
-        krets = map (ThreadsReturn ThreadsInSpace) $ bodyResult body
-
 data KernelInput = KernelInput { kernelInputName :: VName
                                , kernelInputType :: Type
                                , kernelInputArray :: VName
@@ -972,9 +424,6 @@
                                }
                  deriving (Show)
 
-kernelInputParam :: KernelInput -> Param Type
-kernelInputParam p = Param (kernelInputName p) (kernelInputType p)
-
 readKernelInput :: (HasScope scope m, Monad m) =>
                    KernelInput -> m (Stm InKernel)
 readKernelInput inp = do
@@ -985,8 +434,8 @@
     fullSlice arr_t $ map DimFix $ kernelInputIndices inp
 
 newKernelSpace :: MonadFreshNames m =>
-                  (SubExp,SubExp,SubExp) -> SpaceStructure -> m KernelSpace
-newKernelSpace (num_groups, group_size, num_threads) ispace =
+                  (SubExp,SubExp,SubExp,SubExp) -> SpaceStructure -> m KernelSpace
+newKernelSpace (num_groups, group_size, num_threads, virt_groups) ispace =
   KernelSpace
   <$> newVName "global_tid"
   <*> newVName "local_tid"
@@ -994,4 +443,5 @@
   <*> pure num_threads
   <*> pure num_groups
   <*> pure group_size
+  <*> pure virt_groups
   <*> pure ispace
diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs
--- a/src/Futhark/Pass/ExtractKernels/Distribution.hs
+++ b/src/Futhark/Pass/ExtractKernels/Distribution.hs
@@ -243,9 +243,11 @@
                    KernelNest -> KernelBody InKernel
                 -> m (Stms Kernels, SubExp, Stm Kernels)
 constructKernel kernel_nest inner_body = do
-  (w_bnds, w, ispace, inps, rts) <- flatKernel kernel_nest
+  (w_bnds, w, ispace, inps) <- flatKernel kernel_nest
   let used_inps = filter inputIsUsed inps
       cs = loopNestingCertificates first_nest
+      pat = loopNestingPattern first_nest
+      rts = map (stripArray (length ispace)) $ patternTypes pat
 
   (ksize_bnds, k) <- inScopeOf w_bnds $
     mapKernel w (FlatThreadSpace ispace) used_inps rts inner_body
@@ -253,7 +255,7 @@
   let kbnds = w_bnds <> ksize_bnds
   return (kbnds,
           w,
-          Let (loopNestingPattern first_nest) (StmAux cs ()) $ Op $ HostOp k)
+          Let pat (StmAux cs ()) $ Op $ HostOp k)
   where
     first_nest = fst kernel_nest
     inputIsUsed input = kernelInputName input `S.member`
@@ -268,26 +270,22 @@
 --
 --  (2) The index space.
 --
---  (3) The kernel inputs - not that some of these may be unused.
---
---  (4) The per-thread return type.
+--  (3) The kernel inputs - note that some of these may be unused.
 flatKernel :: MonadFreshNames m =>
               KernelNest
            -> m (Stms Kernels,
                  SubExp,
                  [(VName, SubExp)],
-                 [KernelInput],
-                 [Type])
-flatKernel (MapNesting pat _ nesting_w params_and_arrs, []) = do
+                 [KernelInput])
+flatKernel (MapNesting _ _ nesting_w params_and_arrs, []) = do
   i <- newVName "gtid"
   let inps = [ KernelInput pname ptype arr [Var i] |
                (Param pname ptype, arr) <- params_and_arrs ]
-  return (mempty, nesting_w, [(i,nesting_w)], inps,
-          map rowType $ patternTypes pat)
+  return (mempty, nesting_w, [(i,nesting_w)], inps)
 
 flatKernel (MapNesting _ _ nesting_w params_and_arrs, nest : nests) = do
   i <- newVName "gtid"
-  (w_bnds, w, ispace, inps, returns) <- flatKernel (nest, nests)
+  (w_bnds, w, ispace, inps) <- flatKernel (nest, nests)
 
   w' <- newVName "nesting_size"
   let w_bnd = mkLet [] [Ident w' $ Prim int32] $
@@ -304,7 +302,7 @@
             inp
 
   return (w_bnds <> oneStm w_bnd, Var w', (i, nesting_w) : ispace,
-          extra_inps i <> inps', returns)
+          extra_inps i <> inps')
   where extra_inps i =
           [ KernelInput pname ptype arr [Var i] |
             (Param pname ptype, arr) <- params_and_arrs ]
@@ -329,7 +327,7 @@
         removeIdentityMappingGeneral bound_by_stms inner_pat inner_res
   in (DistributionBody
       { distributionTarget = Targets (inner_pat', inner_res') targets
-      , distributionFreeInBody = fold (fmap freeInStm stms) `S.difference` bound_by_stms
+      , distributionFreeInBody = fold (fmap freeIn stms) `S.difference` bound_by_stms
       , distributionIdentityMap = inner_identity_map
       , distributionExpandTarget = inner_expand_target
       },
@@ -527,8 +525,7 @@
     Nothing ->
       return Nothing
   where (dist_body, inner_body_res) = distributionBodyFromStms targets stms
-        inner_body = KernelBody () stms $
-                     map (ThreadsReturn ThreadsInSpace) inner_body_res
+        inner_body = KernelBody () stms $ map ThreadsReturn inner_body_res
 
 tryDistributeStm :: (MonadFreshNames m, HasScope t m, Attributes lore) =>
                     Nestings -> Targets -> Stm lore
diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs
--- a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs
+++ b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs
@@ -54,7 +54,7 @@
                   patternValueIdents res_pat
 
       addStm $ Let res_pat' (StmAux map_cs ()) $ Op $ Screma map_w
-        (ScremaForm (nilFn, mempty) (mempty, nilFn, mempty) map_fun') map_arrs'
+        (ScremaForm (nilFn, mempty) [] map_fun') map_arrs'
 
       forM_ (zip (patternValueIdents res_pat)
                  (patternValueIdents res_pat')) $ \(to, from) -> do
@@ -100,7 +100,7 @@
       map_body <-
         case irwim red_pat w comm red_fun' red_input' of
           Nothing -> do
-            reduce_soac <- reduceSOAC comm red_fun' $ map fst red_input'
+            reduce_soac <- reduceSOAC [Reduce comm red_fun' $ map fst red_input']
             return $ mkBody (oneStm $ Let red_pat (defAux ()) $
                               Op $ Screma w reduce_soac $ map snd red_input') $
               map Var $ patternNames map_pat
diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs
--- a/src/Futhark/Pass/ExtractKernels/Interchange.hs
+++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs
@@ -71,7 +71,7 @@
     return $
       SeqLoop [0..patternSize pat-1] pat' merge_expanded form $
       mkBody (pre_copy_bnds<>oneStm map_bnd) res
-  where free_in_body = freeInBody body
+  where free_in_body = freeIn body
 
         copyOrRemoveParam (param, arr)
           | not (paramName param `S.member` free_in_body) =
diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs
--- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs
+++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs
@@ -39,7 +39,7 @@
                       -> m (Maybe ((SubExp, SubExp), SubExp,
                                    Out.Stms Out.Kernels, Out.Stms Out.Kernels))
 intraGroupParallelise knest lam = runMaybeT $ do
-  (w_stms, w, ispace, inps, rts) <- lift $ flatKernel knest
+  (w_stms, w, ispace, inps) <- lift $ flatKernel knest
   let num_groups = w
       body = lambdaBody lam
 
@@ -75,7 +75,7 @@
                        (eSubExp intra_avail_par)
                   else foldBinOp' (SMax Int32) ws_min
 
-    let inputIsUsed input = kernelInputName input `S.member` freeInBody body
+    let inputIsUsed input = kernelInputName input `S.member` freeIn body
         used_inps = filter inputIsUsed inps
 
     addStms w_stms
@@ -83,7 +83,7 @@
     num_threads <- letSubExp "num_threads" $
                    BasicOp $ BinOp (Mul Int32) num_groups group_size
 
-    let ksize = (num_groups, group_size, num_threads)
+    let ksize = (num_groups, group_size, num_threads, num_groups)
 
     kspace <- newKernelSpace ksize $ FlatThreadSpace $ ispace ++ [(ltid,group_size)]
 
@@ -103,7 +103,8 @@
         return $ PatElem name t'
   flat_pat <- lift $ Pattern [] <$> mapM flatPatElem (patternValueElements nested_pat)
 
-  let kstm = Let flat_pat (StmAux cs ()) $ Op $ HostOp $
+  let rts = map rowType $ patternTypes flat_pat
+      kstm = Let flat_pat (StmAux cs ()) $ Op $ HostOp $
              Kernel (KernelDebugHints "map_intra_group" []) kspace rts kbody'
       reshapeStm nested_pe flat_pe =
         Let (Pattern [] [nested_pe]) (StmAux cs ()) $
@@ -153,13 +154,18 @@
       groupInvariant Constant{} = True
 
   case e of
-    DoLoop ctx val (ForLoop i it bound inps) loopbody
-      | groupInvariant bound ->
-          localScope (scopeOf form) $
+    -- Cosmin hack: previously, only for loops were supported,
+    --              and only if `groupInvariant bound` holds;
+    --              Let's see what can possibly go wrong if we
+    --              completely generalize this (?)
+    DoLoop ctx val form loopbody ->
+          localScope (scopeOf form') $
           localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do
           loopbody' <- intraGroupBody loopbody
-          letBind_ pat $ DoLoop ctx val form loopbody'
-              where form = ForLoop i it bound inps
+          letBind_ pat $ DoLoop ctx val form' loopbody'
+              where form' = case form of
+                              ForLoop i it bound inps -> ForLoop i it bound inps
+                              WhileLoop cond          -> WhileLoop cond
 
     If cond tbody fbody ifattr
       | groupInvariant cond -> do
@@ -189,38 +195,21 @@
 
       scanfun' <- Kernelise.transformLambda scanfun
 
-      -- A GroupScan lambda needs two more parameters.
-      my_index <- newVName "my_index"
-      offset <- newVName "offset"
-      let my_index_param = Param my_index (Prim int32)
-          offset_param = Param offset (Prim int32)
-          scanfun'' = scanfun' { lambdaParams = my_index_param :
-                                                offset_param :
-                                                lambdaParams scanfun'
-                               }
       letBind_ (Pattern [] scan_pes) $
-        Op $ Out.GroupScan w scanfun'' $ zip nes scan_input
+        Op $ Out.GroupScan w scanfun' $ zip nes scan_input
       parallelMin [w]
 
     Op (Screma w form arrs)
-      | Just (_, redfun, nes, foldfun) <- isRedomapSOAC form -> do
+      | Just (reds, map_lam) <- isRedomapSOAC form,
+        Reduce _ red_lam nes <- singleReduce reds -> do
       let (red_pes, map_pes) =
             splitAt (length nes) $ patternElements pat
-      red_input <- procInput ltid (Pattern [] map_pes) w foldfun nes arrs
+      red_input <- procInput ltid (Pattern [] map_pes) w map_lam nes arrs
 
-      redfun' <- Kernelise.transformLambda redfun
+      red_lam' <- Kernelise.transformLambda red_lam
 
-      -- A GroupReduce lambda needs two more parameters.
-      my_index <- newVName "my_index"
-      offset <- newVName "offset"
-      let my_index_param = Param my_index (Prim int32)
-          offset_param = Param offset (Prim int32)
-          redfun'' = redfun' { lambdaParams = my_index_param :
-                                              offset_param :
-                                              lambdaParams redfun'
-                               }
       letBind_ (Pattern [] red_pes) $
-        Op $ Out.GroupReduce w redfun'' $ zip nes red_input
+        Op $ Out.GroupReduce w red_lam' $ zip nes red_input
       parallelMin [w]
 
     Op (Stream w (Sequential accs) lam arrs)
@@ -321,4 +310,4 @@
   (min_ws, avail_ws, kstms) <- runIntraGroupM (Env ltid deps group_variant) $
                  mapM_ intraGroupStm $ bodyStms body
   return (min_ws, avail_ws,
-          KernelBody () kstms $ map (ThreadsReturn OneResultPerGroup) $ bodyResult body)
+          KernelBody () kstms $ map GroupsReturn $ bodyResult body)
diff --git a/src/Futhark/Pass/ExtractKernels/Kernelise.hs b/src/Futhark/Pass/ExtractKernels/Kernelise.hs
--- a/src/Futhark/Pass/ExtractKernels/Kernelise.hs
+++ b/src/Futhark/Pass/ExtractKernels/Kernelise.hs
@@ -35,7 +35,8 @@
 
 transformStm (Let pat aux (Op (Screma w form arrs)))
   -- No map-out part
-  | Just (_, red_lam, nes, map_lam) <- isRedomapSOAC form,
+  | Just (reds, map_lam) <- isRedomapSOAC form,
+    Reduce _ red_lam nes <- singleReduce reds,
     patternSize pat == length nes = do
 
   fold_lam <- composeLambda nilFn red_lam map_lam
diff --git a/src/Futhark/Pass/ExtractKernels/Segmented.hs b/src/Futhark/Pass/ExtractKernels/Segmented.hs
deleted file mode 100644
--- a/src/Futhark/Pass/ExtractKernels/Segmented.hs
+++ /dev/null
@@ -1,89 +0,0 @@
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE TypeFamilies #-}
--- | Segmented scan.
-module Futhark.Pass.ExtractKernels.Segmented
-       ( regularSegmentedScan
-       )
-       where
-
-import Control.Monad
-import qualified Data.Map.Strict as M
-
-import Futhark.Transform.Rename
-import Futhark.Representation.Kernels
-import Futhark.Representation.SOACS.SOAC (nilFn)
-import Futhark.MonadFreshNames
-import Futhark.Tools hiding (false)
-import Futhark.Pass.ExtractKernels.BlockedKernel
-
-addFlagToLambda :: (MonadBinder m, Lore m ~ Kernels) =>
-                   [SubExp] -> Lambda InKernel -> m (Lambda InKernel)
-addFlagToLambda nes lam = do
-  let num_accs = length nes
-  x_flag <- newVName "x_flag"
-  y_flag <- newVName "y_flag"
-  let x_flag_param = Param x_flag $ Prim Bool
-      y_flag_param = Param y_flag $ Prim Bool
-      (x_params, y_params) = splitAt num_accs $ lambdaParams lam
-      params = [x_flag_param] ++ x_params ++ [y_flag_param] ++ y_params
-
-  body <- runBodyBinder $ localScope (scopeOfLParams params) $ do
-    new_flag <- letSubExp "new_flag" $
-                BasicOp $ BinOp LogOr (Var x_flag) (Var y_flag)
-    lhs <- fmap (map Var) $ letTupExp "seg_lhs" $ If (Var y_flag)
-           (resultBody nes)
-           (resultBody $ map (Var . paramName) x_params) $
-           ifCommon $ map paramType x_params
-    let rhs = map (Var . paramName) y_params
-
-    lam' <- renameLambda lam -- avoid shadowing
-    res <- eLambda lam' $ map eSubExp $ lhs ++ rhs
-
-    return $ resultBody $ new_flag : res
-
-  return Lambda { lambdaParams = params
-                , lambdaBody = body
-                , lambdaReturnType = Prim Bool : lambdaReturnType lam
-                }
-
-regularSegmentedScan :: (MonadBinder m, Lore m ~ Kernels) =>
-                        SubExp
-                     -> Pattern Kernels
-                     -> SubExp
-                     -> Lambda InKernel
-                     -> Lambda InKernel
-                     -> [(VName, SubExp)] -> [KernelInput]
-                     -> [SubExp] -> [VName]
-                     -> m ()
-regularSegmentedScan segment_size pat w lam map_lam ispace inps nes arrs = do
-  flags_i <- newVName "flags_i"
-
-  unused_flag_array <- newVName "unused_flag_array"
-  flags_body <-
-    runBodyBinder $ localScope (M.singleton flags_i $ IndexInfo Int32) $ do
-      segment_index <- letSubExp "segment_index" $
-                       BasicOp $ BinOp (SRem Int32) (Var flags_i) segment_size
-      start_of_segment <- letSubExp "start_of_segment" $
-                          BasicOp $ CmpOp (CmpEq int32) segment_index zero
-      let flag = start_of_segment
-      return $ resultBody [flag]
-  (mapk_bnds, mapk) <- mapKernelFromBody w (FlatThreadSpace [(flags_i, w)]) [] [Prim Bool] flags_body
-  addStms mapk_bnds
-  flags <- letExp "flags" $ Op $ HostOp mapk
-
-  lam' <- addFlagToLambda nes lam
-
-  flag_p <- newParam "flag" $ Prim Bool
-  let map_lam' = map_lam { lambdaParams = flag_p : lambdaParams map_lam
-                         , lambdaBody = (lambdaBody map_lam)
-                           { bodyResult = Var (paramName flag_p) : bodyResult (lambdaBody map_lam) }
-                         , lambdaReturnType = Prim Bool : lambdaReturnType map_lam
-                         }
-
-  let pat' = pat { patternValueElements = PatElem unused_flag_array
-                                          (arrayOf (Prim Bool) (Shape [w]) NoUniqueness) :
-                                          patternValueElements pat
-                 }
-  void $ blockedScan pat' w (lam', false:nes) (Commutative, nilFn, mempty) map_lam' segment_size ispace inps (flags:arrs)
-  where zero = constant (0 :: Int32)
-        false = constant False
diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs
--- a/src/Futhark/Pass/KernelBabysitting.hs
+++ b/src/Futhark/Pass/KernelBabysitting.hs
@@ -53,12 +53,16 @@
 nonlinearInMemory :: VName -> ExpMap -> Maybe (Maybe [Int])
 nonlinearInMemory name m =
   case M.lookup name m of
+    Just (Let _ _ (BasicOp (Opaque (Var arr)))) -> nonlinearInMemory arr m
     Just (Let _ _ (BasicOp (Rearrange perm _))) -> Just $ Just $ rearrangeInverse perm
     Just (Let _ _ (BasicOp (Reshape _ arr))) -> nonlinearInMemory arr m
     Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm
     Just (Let pat _ (Op (HostOp (Kernel _ _ ts _)))) ->
       nonlinear =<< find ((==name) . patElemName . fst)
       (zip (patternElements pat) ts)
+    Just (Let pat _ (Op (HostOp (SegMap _ ts _)))) ->
+      nonlinear =<< find ((==name) . patElemName . fst)
+      (zip (patternElements pat) ts)
     _ -> Nothing
   where nonlinear (pe, t)
           | inner_r <- arrayRank t, inner_r > 0 = do
@@ -68,31 +72,14 @@
 
 transformStm :: ExpMap -> Stm Kernels -> BabysitM ExpMap
 
-transformStm expmap (Let pat aux ke@(Op (HostOp (Kernel desc space ts kbody)))) = do
-  -- Go spelunking for accesses to arrays that are defined outside the
-  -- kernel body and where the indices are kernel thread indices.
-  scope <- askScope
-  let thread_gids = map fst $ spaceDimensions space
-      thread_local = S.fromList $ spaceGlobalId space : spaceLocalId space : thread_gids
-      free_ker_vars = freeInExp ke `S.difference` getKerVariantIds space
-  kbody'' <- evalStateT (traverseKernelBodyArrayIndexes
-                         free_ker_vars
-                         thread_local
-                         (castScope scope <> scopeOfKernelSpace space)
-                         (ensureCoalescedAccess expmap (spaceDimensions space) num_threads)
-                         kbody)
-             mempty
-
-  let bnd' = Let pat aux $ Op $ HostOp $ Kernel desc space ts kbody''
-  addStm bnd'
-  return $ M.fromList [ (name, bnd') | name <- patternNames pat ] <> expmap
-  where num_threads = spaceNumThreads space
-        getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ (FlatThreadSpace strct)) =
-            let (gids, _) = unzip strct
-            in  S.fromList $ [glb_id, loc_id, grp_id] ++ gids
-        getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ (NestedThreadSpace strct)) =
-            let (gids, _, lids, _) = unzip4 strct
-            in  S.fromList $ [glb_id, loc_id, grp_id] ++ gids ++ lids
+transformStm expmap (Let pat aux (Op (HostOp op))) = do
+  let mapper = identityKernelMapper { mapOnKernelKernelBody =
+                                        transformKernelBody expmap (kernelSpace op)
+                                    }
+  op' <- mapKernelM mapper op
+  let stm' = Let pat aux $ Op $ HostOp op'
+  addStm stm'
+  return $ M.fromList [ (name, stm') | name <- patternNames pat ] <> expmap
 
 transformStm expmap (Let pat aux e) = do
   e' <- mapExpM (transform expmap) e
@@ -104,6 +91,30 @@
 transform expmap =
   identityMapper { mapOnBody = \scope -> localScope scope . transformBody expmap }
 
+transformKernelBody :: ExpMap -> KernelSpace -> KernelBody InKernel
+                    -> BabysitM (KernelBody InKernel)
+transformKernelBody expmap space kbody = do
+  -- Go spelunking for accesses to arrays that are defined outside the
+  -- kernel body and where the indices are kernel thread indices.
+  scope <- askScope
+  let thread_gids = map fst $ spaceDimensions space
+      thread_local = S.fromList $ spaceGlobalId space : spaceLocalId space : thread_gids
+      free_ker_vars = freeIn kbody `S.difference` getKerVariantIds space
+  evalStateT (traverseKernelBodyArrayIndexes
+              free_ker_vars
+              thread_local
+              (castScope scope <> scopeOfKernelSpace space)
+              (ensureCoalescedAccess expmap (spaceDimensions space) num_threads)
+              kbody)
+    mempty
+  where num_threads = spaceNumThreads space
+        getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ _ (FlatThreadSpace strct)) =
+            let (gids, _) = unzip strct
+            in  S.fromList $ [glb_id, loc_id, grp_id] ++ gids
+        getKerVariantIds (KernelSpace glb_id loc_id grp_id _ _ _ _ (NestedThreadSpace strct)) =
+            let (gids, _, lids, _) = unzip4 strct
+            in  S.fromList $ [glb_id, loc_id, grp_id] ++ gids ++ lids
+
 type ArrayIndexTransform m =
   Names ->
   (VName -> Bool) ->           -- thread local?
@@ -451,4 +462,4 @@
   foldl' add variance $ patternNames $ stmPattern bnd
   where add variance' v = M.insert v binding_variance variance'
         look variance' v = S.insert v $ M.findWithDefault mempty v variance'
-        binding_variance = mconcat $ map (look variance) $ S.toList (freeInStm bnd)
+        binding_variance = mconcat $ map (look variance) $ S.toList (freeIn bnd)
diff --git a/src/Futhark/Pass/ResolveAssertions.hs b/src/Futhark/Pass/ResolveAssertions.hs
--- a/src/Futhark/Pass/ResolveAssertions.hs
+++ b/src/Futhark/Pass/ResolveAssertions.hs
@@ -23,7 +23,6 @@
 import Futhark.Representation.SOACS (SOACS)
 import qualified Futhark.Representation.SOACS.Simplify as Simplify
 import qualified Futhark.Optimise.Simplify as Simplify
-import Futhark.Optimise.Simplify.Rules
 
 import Prelude
 
@@ -33,7 +32,7 @@
   "resolve assertions"
   "Try to statically resolve bounds checks and similar." $
   Simplify.simplifyProg Simplify.simpleSOACS rulebook Simplify.noExtraHoistBlockers
-  where rulebook = standardRules <> ruleBook [ RuleBasicOp simplifyScalExp ] []
+  where rulebook = Simplify.soacRules <> ruleBook [ RuleBasicOp simplifyScalExp ] []
 
 simplifyScalExp :: BinderOps lore => TopDownRuleBasicOp lore
 simplifyScalExp vtable pat _ e = do
diff --git a/src/Futhark/Representation/AST/Attributes/Names.hs b/src/Futhark/Representation/AST/Attributes/Names.hs
--- a/src/Futhark/Representation/AST/Attributes/Names.hs
+++ b/src/Futhark/Representation/AST/Attributes/Names.hs
@@ -10,10 +10,6 @@
          , Names
          -- * Specialised Functions
          , freeInStmsAndRes
-         , freeInBody
-         , freeInExp
-         , freeInStm
-         , freeInLambda
          -- * Bound Names
          , boundInBody
          , boundByStm
@@ -43,9 +39,8 @@
               Walker lore (Writer Names)
 freeWalker = identityWalker {
                walkOnSubExp = tell . freeIn
-             , walkOnBody = tell . freeInBody
+             , walkOnBody = tell . freeIn
              , walkOnVName = tell . S.singleton
-             , walkOnCertificates = tell . freeIn
              , walkOnOp = tell . freeIn
              }
 
@@ -60,68 +55,9 @@
                      FreeAttr (ExpAttr lore)) =>
                     Stms lore -> Result -> Names
 freeInStmsAndRes stms res =
-  (freeIn res `mappend` fold (fmap freeInStm stms))
+  (freeIn res `mappend` fold (fmap freeIn stms))
   `S.difference` boundByStms stms
 
--- | Return the set of variable names that are free in the given body.
-freeInBody :: (FreeAttr (ExpAttr lore),
-               FreeAttr (BodyAttr lore),
-               FreeIn (FParamAttr lore),
-               FreeIn (LParamAttr lore),
-               FreeIn (LetAttr lore),
-               FreeIn (Op lore)) =>
-              Body lore -> Names
-freeInBody (Body attr stms res) =
-  precomputed attr $ freeIn attr <> freeInStmsAndRes stms res
-
--- | Return the set of variable names that are free in the given
--- expression.
-freeInExp :: (FreeAttr (ExpAttr lore),
-              FreeAttr (BodyAttr lore),
-              FreeIn (FParamAttr lore),
-              FreeIn (LParamAttr lore),
-              FreeIn (LetAttr lore),
-              FreeIn (Op lore)) =>
-             Exp lore -> Names
-freeInExp (DoLoop ctxmerge valmerge form loopbody) =
-  let (ctxparams, ctxinits) = unzip ctxmerge
-      (valparams, valinits) = unzip valmerge
-      bound_here = S.fromList $ M.keys $
-                   scopeOf form <>
-                   scopeOfFParams (ctxparams ++ valparams)
-  in (freeIn (ctxinits ++ valinits) <> freeIn form <>
-      freeIn (ctxparams ++ valparams) <> freeInBody loopbody)
-     `S.difference` bound_here
-freeInExp e = execWriter $ walkExpM freeWalker e
-
--- | Return the set of variable names that are free in the given
--- binding.
-freeInStm :: (FreeAttr (ExpAttr lore),
-              FreeAttr (BodyAttr lore),
-              FreeIn (FParamAttr lore),
-              FreeIn (LParamAttr lore),
-              FreeIn (LetAttr lore),
-              FreeIn (Op lore)) =>
-             Stm lore -> Names
-freeInStm (Let pat (StmAux cs attr) e) =
-  freeIn cs <> precomputed attr (freeIn attr <> freeInExp e <> freeIn pat)
-
--- | Return the set of variable names that are free in the given
--- lambda, including shape annotations in the parameters.
-freeInLambda :: (FreeAttr (ExpAttr lore),
-                 FreeAttr (BodyAttr lore),
-                 FreeIn (FParamAttr lore),
-                 FreeIn (LParamAttr lore),
-                 FreeIn (LetAttr lore),
-                 FreeIn (Op lore)) =>
-                Lambda lore -> Names
-freeInLambda (Lambda params body rettype) =
-  S.filter (`notElem` paramnames) $ inRet <> inParams <> inBody
-  where inRet = mconcat $ map freeIn rettype
-        inParams = mconcat $ map freeIn params
-        inBody = freeInBody body
-        paramnames = map paramName params
-
 -- | A class indicating that we can obtain free variable information
 -- from values of this type.
 class FreeIn a where
@@ -142,6 +78,54 @@
 instance FreeIn a => FreeIn [a] where
   freeIn = fold . fmap freeIn
 
+instance (FreeAttr (ExpAttr lore),
+          FreeAttr (BodyAttr lore),
+          FreeIn (FParamAttr lore),
+          FreeIn (LParamAttr lore),
+          FreeIn (LetAttr lore),
+          FreeIn (Op lore)) => FreeIn (Lambda lore) where
+  freeIn (Lambda params body rettype) =
+    S.filter (`notElem` paramnames) $ inRet <> inParams <> inBody
+    where inRet = mconcat $ map freeIn rettype
+          inParams = mconcat $ map freeIn params
+          inBody = freeIn body
+          paramnames = map paramName params
+
+instance (FreeAttr (ExpAttr lore),
+          FreeAttr (BodyAttr lore),
+          FreeIn (FParamAttr lore),
+          FreeIn (LParamAttr lore),
+          FreeIn (LetAttr lore),
+          FreeIn (Op lore)) => FreeIn (Body lore) where
+  freeIn (Body attr stms res) =
+    precomputed attr $ freeIn attr <> freeInStmsAndRes stms res
+
+instance (FreeAttr (ExpAttr lore),
+          FreeAttr (BodyAttr lore),
+          FreeIn (FParamAttr lore),
+          FreeIn (LParamAttr lore),
+          FreeIn (LetAttr lore),
+          FreeIn (Op lore)) => FreeIn (Exp lore) where
+  freeIn (DoLoop ctxmerge valmerge form loopbody) =
+    let (ctxparams, ctxinits) = unzip ctxmerge
+        (valparams, valinits) = unzip valmerge
+        bound_here = S.fromList $ M.keys $
+                     scopeOf form <>
+                     scopeOfFParams (ctxparams ++ valparams)
+    in (freeIn (ctxinits ++ valinits) <> freeIn form <>
+        freeIn (ctxparams ++ valparams) <> freeIn loopbody)
+       `S.difference` bound_here
+  freeIn e = execWriter $ walkExpM freeWalker e
+
+instance (FreeAttr (ExpAttr lore),
+          FreeAttr (BodyAttr lore),
+          FreeIn (FParamAttr lore),
+          FreeIn (LParamAttr lore),
+          FreeIn (LetAttr lore),
+          FreeIn (Op lore)) => FreeIn (Stm lore) where
+  freeIn (Let pat (StmAux cs attr) e) =
+    freeIn cs <> precomputed attr (freeIn attr <> freeIn e <> freeIn pat)
+
 instance FreeIn (Stm lore) => FreeIn (Stms lore) where
   freeIn = fold . fmap freeIn
 
@@ -173,7 +157,7 @@
 
 instance FreeIn shape => FreeIn (TypeBase shape u) where
   freeIn (Array _ shape _) = freeIn shape
-  freeIn (Mem size _)      = freeIn size
+  freeIn (Mem _)           = mempty
   freeIn (Prim _)          = mempty
 
 instance FreeIn attr => FreeIn (ParamT attr) where
diff --git a/src/Futhark/Representation/AST/Attributes/TypeOf.hs b/src/Futhark/Representation/AST/Attributes/TypeOf.hs
--- a/src/Futhark/Representation/AST/Attributes/TypeOf.hs
+++ b/src/Futhark/Representation/AST/Attributes/TypeOf.hs
@@ -24,8 +24,6 @@
        , primOpType
        , mapType
        , subExpShapeContext
-       , loopResultContext
-       , loopExtType
 
        -- * Return type
        , module Futhark.Representation.AST.RetType
@@ -46,7 +44,6 @@
 import Futhark.Representation.AST.Attributes.Types
 import Futhark.Representation.AST.Attributes.Patterns
 import Futhark.Representation.AST.Attributes.Constants
-import Futhark.Representation.AST.Attributes.Names
 import Futhark.Representation.AST.RetType
 import Futhark.Representation.AST.Attributes.Scope
 
@@ -165,15 +162,6 @@
                       [TypeBase ExtShape u] -> [SubExp] -> m [SubExp]
 subExpShapeContext rettype ses =
   extractShapeContext rettype <$> traverse (fmap arrayDims . subExpType) ses
-
--- | A loop returns not only its value merge parameters, but may also
--- have an existential context.  Thus, @loopResult ctxmergeparams
--- valmergeparams@ returns those paramters in @ctxmergeparams@ that
--- constitute the returned context.
-loopResultContext :: FreeIn attr => [Param attr] -> [Param attr] -> [Param attr]
-loopResultContext ctx val = filter usedInValue ctx
-  where usedInValue = (`S.member` used) . paramName
-        used = freeIn val <> freeIn ctx
 
 -- | Given the context and value merge parameters of a Futhark @loop@,
 -- produce the return type.
diff --git a/src/Futhark/Representation/AST/Attributes/Types.hs b/src/Futhark/Representation/AST/Attributes/Types.hs
--- a/src/Futhark/Representation/AST/Attributes/Types.hs
+++ b/src/Futhark/Representation/AST/Attributes/Types.hs
@@ -84,7 +84,7 @@
 rankShaped :: ArrayShape shape => TypeBase shape u -> TypeBase Rank u
 rankShaped (Array et sz u) = Array et (Rank $ shapeRank sz) u
 rankShaped (Prim et) = Prim et
-rankShaped (Mem size space) = Mem size space
+rankShaped (Mem space) = Mem space
 
 -- | Return the dimensionality of a type.  For non-arrays, this is
 -- zero.  For a one-dimensional array it is one, for a two-dimensional
@@ -107,8 +107,8 @@
   | shapeRank ds' == 0 = Prim t
   | otherwise          = Array t (f ds) u
   where ds' = f ds
-modifyArrayShape _ (Prim t)        = Prim t
-modifyArrayShape _ (Mem size space) = Mem size space
+modifyArrayShape _ (Prim t)    = Prim t
+modifyArrayShape _ (Mem space) = Mem space
 
 -- | Set the shape of an array.  If the given type is not an
 -- array, return the type unchanged.
@@ -152,8 +152,8 @@
   Prim bt
 staticShapes1 (Array bt (Shape shape) u) =
   Array bt (Shape $ map Free shape) u
-staticShapes1 (Mem size space) =
-  Mem size space
+staticShapes1 (Mem space) =
+  Mem space
 
 -- | @arrayOf t s u@ constructs an array type.  The convenience
 -- compared to using the 'Array' constructor directly is that @t@ can
@@ -317,7 +317,7 @@
   t1 == t2 &&
   shape1 `subShapeOf` shape2
 subtypeOf (Prim t1) (Prim t2) = t1 == t2
-subtypeOf (Mem _ space1) (Mem _ space2) = space1 == space2
+subtypeOf (Mem space1) (Mem space2) = space1 == space2
 subtypeOf _ _ = False
 
 -- | @xs \`subtypesOf\` ys@ is true if @xs@ is the same size as @ys@,
@@ -335,13 +335,13 @@
        -> TypeBase shape Uniqueness
 toDecl (Prim bt) _ = Prim bt
 toDecl (Array et shape _) u = Array et shape u
-toDecl (Mem size space) _ = Mem size space
+toDecl (Mem space) _ = Mem space
 
 fromDecl :: TypeBase shape Uniqueness
          -> TypeBase shape NoUniqueness
 fromDecl (Prim bt) = Prim bt
 fromDecl (Array et shape _) = Array et shape NoUniqueness
-fromDecl (Mem size space) = Mem size space
+fromDecl (Mem space) = Mem space
 
 -- | Given the existential return type of a function, and the shapes
 -- of the values returned by the function, return the existential
@@ -374,10 +374,8 @@
 -- | If all dimensions of the given 'RetType' are statically known,
 -- return the corresponding list of 'Type'.
 hasStaticShape :: ExtType -> Maybe Type
-hasStaticShape (Prim bt) =
-  Just $ Prim bt
-hasStaticShape (Mem size space) =
-  Just $ Mem size space
+hasStaticShape (Prim bt) = Just $ Prim bt
+hasStaticShape (Mem space) = Just $ Mem space
 hasStaticShape (Array bt (Shape shape) u) =
   Array bt <$> (Shape <$> mapM isFree shape) <*> pure u
   where isFree (Free s) = Just s
diff --git a/src/Futhark/Representation/AST/Pretty.hs b/src/Futhark/Representation/AST/Pretty.hs
--- a/src/Futhark/Representation/AST/Pretty.hs
+++ b/src/Futhark/Representation/AST/Pretty.hs
@@ -85,22 +85,22 @@
   ppr (Prim et) = ppr et
   ppr (Array et (Shape ds) u) =
     ppr u <> mconcat (map (brackets . ppr) ds) <> ppr et
-  ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s)
-  ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp
+  ppr (Mem DefaultSpace) = text "mem"
+  ppr (Mem (Space sp)) = text "mem" <> text "@" <> text sp
 
 instance Pretty u => Pretty (TypeBase ExtShape u) where
   ppr (Prim et) = ppr et
   ppr (Array et (Shape ds) u) =
     ppr u <> mconcat (map (brackets . ppr) ds) <> ppr et
-  ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s)
-  ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp
+  ppr (Mem DefaultSpace) = text "mem"
+  ppr (Mem (Space sp)) = text "mem" <> text "@" <> text sp
 
 instance Pretty u => Pretty (TypeBase Rank u) where
   ppr (Prim et) = ppr et
   ppr (Array et (Rank n) u) =
     ppr u <> mconcat (replicate n $ brackets mempty) <> ppr et
-  ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s)
-  ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp
+  ppr (Mem DefaultSpace) = text "mem"
+  ppr (Mem (Space sp)) = text "mem" <> text "@" <> text sp
 
 instance Pretty Ident where
   ppr ident = ppr (identType ident) <+> ppr (identName ident)
diff --git a/src/Futhark/Representation/AST/Syntax/Core.hs b/src/Futhark/Representation/AST/Syntax/Core.hs
--- a/src/Futhark/Representation/AST/Syntax/Core.hs
+++ b/src/Futhark/Representation/AST/Syntax/Core.hs
@@ -165,7 +165,7 @@
 -- comparing types for equality with '==', shapes must match.
 data TypeBase shape u = Prim PrimType
                       | Array PrimType shape u
-                      | Mem SubExp Space
+                      | Mem Space
                     deriving (Show, Eq, Ord)
 
 -- | A type with shape information, used for describing the type of
diff --git a/src/Futhark/Representation/AST/Traversals.hs b/src/Futhark/Representation/AST/Traversals.hs
--- a/src/Futhark/Representation/AST/Traversals.hs
+++ b/src/Futhark/Representation/AST/Traversals.hs
@@ -26,7 +26,6 @@
   -- * Mapping
     Mapper(..)
   , identityMapper
-  , mapBody
   , mapExpM
   , mapExp
   , mapOnType
@@ -37,8 +36,6 @@
   , Walker(..)
   , identityWalker
   , walkExpM
-  , walkExp
-  -- * Simple wrappers
   )
   where
 
@@ -59,7 +56,6 @@
     -- ^ Most bodies are enclosed in a scope, which is passed along
     -- for convenience.
   , mapOnVName :: VName -> m VName
-  , mapOnCertificates :: Certificates -> m Certificates
   , mapOnRetType :: RetType flore -> m (RetType tlore)
   , mapOnBranchType :: BranchType flore -> m (BranchType tlore)
   , mapOnFParam :: FParam flore -> m (FParam tlore)
@@ -73,7 +69,6 @@
                    mapOnSubExp = return
                  , mapOnBody = const return
                  , mapOnVName = return
-                 , mapOnCertificates = return
                  , mapOnRetType = return
                  , mapOnBranchType = return
                  , mapOnFParam = return
@@ -81,10 +76,6 @@
                  , mapOnOp = return
                  }
 
--- | Map across the bindings of a 'Body'.
-mapBody :: (Stm lore -> Stm lore) -> Body lore -> Body lore
-mapBody f (Body attr stms res) = Body attr (fmap f stms) res
-
 -- | Map a monadic action across the immediate children of an
 -- expression.  Importantly, the 'mapOnExp' action is not invoked for
 -- the expression itself, and the mapping does not descend recursively
@@ -171,7 +162,7 @@
   where mapOnExtSize (Ext x)   = return $ Ext x
         mapOnExtSize (Free se) = Free <$> mapOnSubExp tv se
 mapOnExtType _ (Prim bt) = return $ Prim bt
-mapOnExtType tv (Mem size space) = Mem <$> mapOnSubExp tv size <*> pure space
+mapOnExtType _ (Mem space) = pure $ Mem space
 
 mapOnLoopForm :: Monad m =>
                  Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore)
@@ -189,7 +180,7 @@
 mapOnType :: Monad m =>
              (SubExp -> m SubExp) -> Type -> m Type
 mapOnType _ (Prim bt) = return $ Prim bt
-mapOnType f (Mem size space) = Mem <$> f size <*> pure space
+mapOnType _ (Mem space) = pure $ Mem space
 mapOnType f (Array bt shape u) =
   Array bt <$> (Shape <$> mapM f (shapeDims shape)) <*> pure u
 
@@ -200,7 +191,6 @@
     walkOnSubExp :: SubExp -> m ()
   , walkOnBody :: Body lore -> m ()
   , walkOnVName :: VName -> m ()
-  , walkOnCertificates :: Certificates -> m ()
   , walkOnRetType :: RetType lore -> m ()
   , walkOnBranchType :: BranchType lore -> m ()
   , walkOnFParam :: FParam lore -> m ()
@@ -214,7 +204,6 @@
                    walkOnSubExp = const $ return ()
                  , walkOnBody = const $ return ()
                  , walkOnVName = const $ return ()
-                 , walkOnCertificates = const $ return ()
                  , walkOnRetType = const $ return ()
                  , walkOnBranchType = const $ return ()
                  , walkOnFParam = const $ return ()
@@ -227,7 +216,6 @@
                  mapOnSubExp = wrap walkOnSubExp
                , mapOnBody = const $ wrap walkOnBody
                , mapOnVName = wrap walkOnVName
-               , mapOnCertificates = wrap walkOnCertificates
                , mapOnRetType = wrap walkOnRetType
                , mapOnBranchType = wrap walkOnBranchType
                , mapOnFParam = wrap walkOnFParam
@@ -240,7 +228,3 @@
 walkExpM :: Monad m => Walker lore m -> Exp lore -> m ()
 walkExpM f = void . mapExpM m
   where m = walkMapper f
-
--- | As 'walkExp', but runs in the 'Identity' monad..
-walkExp :: Walker lore Identity -> Exp lore -> ()
-walkExp f = runIdentity . walkExpM f
diff --git a/src/Futhark/Representation/Aliases.hs b/src/Futhark/Representation/Aliases.hs
--- a/src/Futhark/Representation/Aliases.hs
+++ b/src/Futhark/Representation/Aliases.hs
@@ -246,7 +246,7 @@
 mkAliasedBody innerlore bnds res =
   Body (mkBodyAliases bnds res, innerlore) bnds res
 
-mkPatternAliases :: (Attributes lore, Aliased lore, Typed attr) =>
+mkPatternAliases :: (Aliased lore, Typed attr) =>
                     PatternT attr -> Exp lore
                  -> ([PatElemT (VarAliases, attr)],
                      [PatElemT (VarAliases, attr)])
@@ -265,14 +265,13 @@
           where names' =
                   case patElemType bindee of
                     Array {} -> names
-                    Mem _ _  -> names
+                    Mem _    -> names
                     _        -> mempty
 
-mkContextAliases :: (Attributes lore, Aliased lore) =>
-                    PatternT attr -> Exp lore
-                 -> [Names]
+mkContextAliases :: Aliased lore =>
+                    PatternT attr -> Exp lore -> [Names]
 mkContextAliases pat (DoLoop ctxmerge valmerge _ body) =
-  let ctx = loopResultContext (map fst ctxmerge) (map fst valmerge)
+  let ctx = map fst ctxmerge
       init_als = zip mergenames $ map (subExpAliases . snd) $ ctxmerge ++ valmerge
       expand als = als <> S.unions (mapMaybe (`lookup` init_als) (S.toList als))
       merge_als = zip mergenames $
diff --git a/src/Futhark/Representation/ExplicitMemory.hs b/src/Futhark/Representation/ExplicitMemory.hs
--- a/src/Futhark/Representation/ExplicitMemory.hs
+++ b/src/Futhark/Representation/ExplicitMemory.hs
@@ -76,7 +76,6 @@
        , sliceInfo
        , lookupMemInfo
        , subExpMemInfo
-       , lookupMemSize
        , lookupArraySummary
        , fullyLinear
        , ixFunMatchesInnerShape
@@ -124,7 +123,6 @@
 import Futhark.Representation.Aliases
   (Aliases, removeScopeAliases, removeExpAliases, removePatternAliases)
 import Futhark.Representation.AST.Attributes.Ranges
-import Futhark.Analysis.Usage
 import qualified Futhark.Analysis.SymbolTable as ST
 
 -- | A lore containing explicit memory information.
@@ -157,7 +155,7 @@
   freeIn (Inner k) = freeIn k
 
 instance TypedOp inner => TypedOp (MemOp inner) where
-  opType (Alloc size space) = pure [Mem size space]
+  opType (Alloc _ space) = pure [Mem space]
   opType (Inner k) = opType k
 
 instance AliasedOp inner => AliasedOp (MemOp inner) where
@@ -212,10 +210,6 @@
   cheapOp (Inner k) = cheapOp k
   cheapOp Alloc{} = True
 
-instance UsageInOp inner => UsageInOp (MemOp inner) where
-  usageInOp Alloc {} = mempty
-  usageInOp (Inner k) = usageInOp k
-
 instance CanBeWise inner => CanBeWise (MemOp inner) where
   type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner)
   removeOpWisdom (Alloc size space) = Alloc size space
@@ -252,7 +246,7 @@
 -- over uniqueness, dimension, and auxiliary array information.
 data MemInfo d u ret = MemPrim PrimType
                      -- ^ A primitive value.
-                     | MemMem d Space
+                     | MemMem Space
                      -- ^ A memory block.
                      | MemArray PrimType (ShapeBase d) u ret
                      -- ^ The array is stored in the named memory block,
@@ -267,19 +261,17 @@
 
 instance FixExt ret => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where
   declExtTypeOf (MemPrim pt) = Prim pt
-  declExtTypeOf (MemMem (Free size) space) = Mem size space
-  declExtTypeOf (MemMem Ext{} space) = Mem (intConst Int32 0) space -- XXX
+  declExtTypeOf (MemMem space) = Mem space
   declExtTypeOf (MemArray pt shape u _) = Array pt shape u
 
 instance FixExt ret => ExtTyped (MemInfo ExtSize NoUniqueness ret) where
   extTypeOf (MemPrim pt) = Prim pt
-  extTypeOf (MemMem (Free size) space) = Mem size space
-  extTypeOf (MemMem Ext{} space) = Mem (intConst Int32 0) space -- XXX
+  extTypeOf (MemMem space) = Mem space
   extTypeOf (MemArray pt shape u _) = Array pt shape u
 
 instance FixExt ret => FixExt (MemInfo ExtSize u ret) where
   fixExt _ _ (MemPrim pt) = MemPrim pt
-  fixExt i se (MemMem size space) = MemMem (fixExt i se size) space
+  fixExt _ _ (MemMem space) = MemMem space
   fixExt i se (MemArray pt shape u ret) =
     MemArray pt (fixExt i se shape) u (fixExt i se ret)
 
@@ -288,26 +280,26 @@
 
 instance Typed (MemInfo SubExp NoUniqueness ret) where
   typeOf (MemPrim pt) = Prim pt
-  typeOf (MemMem size space) = Mem size space
+  typeOf (MemMem space) = Mem space
   typeOf (MemArray bt shape u _) = Array bt shape u
 
 instance DeclTyped (MemInfo SubExp Uniqueness ret) where
   declTypeOf (MemPrim bt) = Prim bt
-  declTypeOf (MemMem size space) = Mem size space
+  declTypeOf (MemMem space) = Mem space
   declTypeOf (MemArray bt shape u _) = Array bt shape u
 
 instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where
   freeIn (MemArray _ shape _ ret) = freeIn shape <> freeIn ret
-  freeIn (MemMem size _) = freeIn size
-  freeIn (MemPrim _) = mempty
+  freeIn MemMem{} = mempty
+  freeIn MemPrim{} = mempty
 
 instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where
   substituteNames subst (MemArray bt shape u ret) =
     MemArray bt
     (substituteNames subst shape) u
     (substituteNames subst ret)
-  substituteNames substs (MemMem size space) =
-    MemMem (substituteNames substs size) space
+  substituteNames _ (MemMem space) =
+    MemMem space
   substituteNames _ (MemPrim bt) =
     MemPrim bt
 
@@ -331,18 +323,18 @@
          Engine.Simplifiable (MemInfo d u ret) where
   simplify (MemPrim bt) =
     return $ MemPrim bt
-  simplify (MemMem size space) =
-    MemMem <$> Engine.simplify size <*> pure space
+  simplify (MemMem space) =
+    pure $ MemMem space
   simplify (MemArray bt shape u ret) =
     MemArray bt <$> Engine.simplify shape <*> pure u <*> Engine.simplify ret
 
 instance (PP.Pretty (TypeBase (ShapeBase d) u),
           PP.Pretty d, PP.Pretty u, PP.Pretty ret) => PP.Pretty (MemInfo d u ret) where
   ppr (MemPrim bt) = PP.ppr bt
-  ppr (MemMem s DefaultSpace) =
-    PP.text "mem" <> PP.parens (PP.ppr s)
-  ppr (MemMem s (Space sp)) =
-    PP.text "mem" <> PP.parens (PP.ppr s) <> PP.text "@" <> PP.text sp
+  ppr (MemMem DefaultSpace) =
+    PP.text "mem"
+  ppr (MemMem (Space sp)) =
+    PP.text "mem" <> PP.text "@" <> PP.text sp
   ppr (MemArray bt shape u ret) =
     PP.ppr (Array bt shape u) <> PP.text "@" <> PP.ppr ret
 
@@ -386,9 +378,9 @@
 data MemReturn = ReturnsInBlock VName ExtIxFun
                  -- ^ The array is located in a memory block that is
                  -- already in scope.
-               | ReturnsNewBlock Space Int ExtSize ExtIxFun
-                 -- ^ The operation returns a new (existential) block,
-                 -- with an existential or known size.
+               | ReturnsNewBlock Space Int ExtIxFun
+                 -- ^ The operation returns a new (existential) memory
+                 -- block.
                deriving (Show)
 
 instance Eq MemReturn where
@@ -403,15 +395,15 @@
 instance Substitute MemReturn where
   substituteNames substs (ReturnsInBlock ident ixfun) =
     ReturnsInBlock (substituteNames substs ident) (substituteNames substs ixfun)
-  substituteNames substs (ReturnsNewBlock space i size ixfun) =
-    ReturnsNewBlock space i (substituteNames substs size) (substituteNames substs ixfun)
+  substituteNames substs (ReturnsNewBlock space i ixfun) =
+    ReturnsNewBlock space i (substituteNames substs ixfun)
 
 instance FixExt MemReturn where
-  fixExt i (Var v) (ReturnsNewBlock _ j _ ixfun)
+  fixExt i (Var v) (ReturnsNewBlock _ j ixfun)
     | j == i = ReturnsInBlock v $ fixExtIxFun i
                (primExpFromSubExp int32 (Var v)) ixfun
-  fixExt i se (ReturnsNewBlock space j size ixfun) =
-    ReturnsNewBlock space j' (fixExt i se size)
+  fixExt i se (ReturnsNewBlock space j ixfun) =
+    ReturnsNewBlock space j'
     (fixExtIxFun i (primExpFromSubExp int32 se) ixfun)
     where j' | i < j     = j-1
              | otherwise = j
@@ -435,9 +427,8 @@
 instance PP.Pretty MemReturn where
   ppr (ReturnsInBlock v ixfun) =
     PP.parens $ PP.text (pretty v) <> PP.text "->" <> PP.ppr ixfun
-  ppr (ReturnsNewBlock space i size ixfun) =
-    PP.text ("?" ++ show i) <> space' <> PP.parens (PP.ppr size)
-    <> PP.text "->" <> PP.ppr ixfun
+  ppr (ReturnsNewBlock space i ixfun) =
+    PP.text ("?" ++ show i) <> space' <> PP.text "->" <> PP.ppr ixfun
     where space' = case space of DefaultSpace -> mempty
                                  Space s -> PP.text $ "@" ++ s
 
@@ -446,8 +437,8 @@
   freeIn _                        = mempty
 
 instance Engine.Simplifiable MemReturn where
-  simplify (ReturnsNewBlock space i size ixfun) =
-    ReturnsNewBlock space i <$> Engine.simplify size <*> simplifyExtIxFun ixfun
+  simplify (ReturnsNewBlock space i ixfun) =
+    ReturnsNewBlock space i <$> simplifyExtIxFun ixfun
   simplify (ReturnsInBlock v ixfun) =
     ReturnsInBlock <$> Engine.simplify v <*> simplifyExtIxFun ixfun
 
@@ -486,16 +477,16 @@
   MemArray bt shape u $ Just ret
 maybeReturns (MemPrim bt) =
   MemPrim bt
-maybeReturns (MemMem size space) =
-  MemMem size space
+maybeReturns (MemMem space) =
+  MemMem space
 
 noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r
 noUniquenessReturns (MemArray bt shape _ r) =
   MemArray bt shape NoUniqueness r
 noUniquenessReturns (MemPrim bt) =
   MemPrim bt
-noUniquenessReturns (MemMem size space) =
-  MemMem size space
+noUniquenessReturns (MemMem space) =
+  MemMem space
 
 funReturnsToExpReturns :: FunReturns -> ExpReturns
 funReturnsToExpReturns = noUniquenessReturns . maybeReturns
@@ -619,8 +610,8 @@
 
       checkReturn (MemPrim x) (MemPrim y)
         | x == y = return ()
-      checkReturn (MemMem x _) (MemMem y _) =
-        checkDim x y
+      checkReturn (MemMem x) (MemMem y)
+        | x == y = return ()
       checkReturn (MemArray x_pt x_shape _ x_ret)
                   (MemArray y_pt y_shape _ y_ret)
         | x_pt == y_pt, shapeRank x_shape == shapeRank y_shape = do
@@ -648,7 +639,7 @@
                     "\nixfun of body result: ", pretty y_ixfun',
                     "\nixfun of return type: ", pretty x_ixfun',
                     "\nand context elements: ", pretty ctx_res]
-      checkMemReturn (ReturnsNewBlock x_space x_ext x_mem_size x_ixfun)
+      checkMemReturn (ReturnsNewBlock x_space x_ext x_ixfun)
                      (ArrayIn y_mem y_ixfun) = do
         (x_mem, x_mem_type)  <- fetchCtx x_ext
         let x_ixfun' = substConstsInExtIndFun x_ixfun
@@ -659,14 +650,13 @@
             "\nixfun of return type: ", pretty x_ixfun',
             "\nand context elements: ", pretty ctx_res]
         case x_mem_type of
-          MemMem y_mem_size y_space -> do
+          MemMem y_space -> do
             unless (x_mem == Var y_mem) $
               throwError $ unwords ["Expected memory", pretty x_ext, "=>", pretty x_mem,
                                     "but got", pretty y_mem]
             unless (x_space == y_space) $
               throwError $ unwords ["Expected memory", pretty y_mem, "in space", pretty x_space,
                                     "but actually in space", pretty y_space]
-            checkDim x_mem_size y_mem_size
           t ->
             throwError $ unwords ["Expected memory", pretty x_ext, "=>", pretty x_mem,
                                   "but but has type", pretty t]
@@ -705,8 +695,8 @@
                             "\ncannot match pattern type:\n  " ++ prettyTuple val_ts ++
                             "\nwith context elements: " ++ pretty ctx_ids
   where matches _ _ (MemPrim x) (MemPrim y) = x == y
-        matches _ _ (MemMem x_size x_space) (MemMem y_size y_space) =
-          x_size == y_size && x_space == y_space
+        matches _ _ (MemMem x_space) (MemMem y_space) =
+          x_space == y_space
         matches ctxids ctxexts (MemArray x_pt x_shape _ x_ret) (MemArray y_pt y_shape _ y_ret) =
           x_pt == y_pt && x_shape == y_shape &&
           case (x_ret, y_ret) of
@@ -715,16 +705,15 @@
                   y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun
               in  x_mem == y_mem && x_ixfun' == y_ixfun'
             (ReturnsInBlock _ x_ixfun,
-             Just (ReturnsNewBlock _ _ _ y_ixfun)) ->
+             Just (ReturnsNewBlock _ _ y_ixfun)) ->
               let x_ixfun' = IxFun.substituteInIxFun ctxids  x_ixfun
                   y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun
               in  x_ixfun' == y_ixfun'
-            (ReturnsNewBlock x_space x_i x_size x_ixfun,
-             Just (ReturnsNewBlock y_space y_i y_size y_ixfun)) ->
+            (ReturnsNewBlock x_space x_i x_ixfun,
+             Just (ReturnsNewBlock y_space y_i y_ixfun)) ->
               let x_ixfun' = IxFun.substituteInIxFun  ctxids x_ixfun
                   y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun
-              in  x_space == y_space && x_i == y_i &&
-                  x_size == y_size && x_ixfun' == y_ixfun'
+              in  x_space == y_space && x_i == y_i && x_ixfun' == y_ixfun'
             (_, Nothing) -> True
             _ -> False
         matches _ _ _ _ = False
@@ -767,19 +756,11 @@
     _ ->
       fail $ "Variable " ++ pretty name ++ " does not look like an array."
 
-lookupMemSize :: (HasScope lore m, Monad m) =>
-                 VName -> m SubExp
-lookupMemSize v = do
-  t <- lookupType v
-  case t of Mem size _ -> return size
-            _ -> fail $ "lookupMemSize: " ++ pretty v ++ " is not a memory block."
-
 checkMemInfo :: TC.Checkable lore =>
                  VName -> MemInfo SubExp u MemBind
              -> TC.TypeM lore ()
 checkMemInfo _ (MemPrim _) = return ()
-checkMemInfo _ (MemMem size _) =
-  TC.require [Prim int64] size
+checkMemInfo _ (MemMem _) = return ()
 checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do
   t <- lookupType v
   case t of
@@ -823,12 +804,12 @@
          (patElemName pe,
           case patElemAttr pe of
             MemPrim pt -> MemPrim pt
-            MemMem size space -> MemMem (ext size) space
+            MemMem space -> MemMem space
             MemArray pt shape u (ArrayIn mem ixfun) ->
               MemArray pt (Shape $ map ext $ shapeDims shape) u $
               case find ((==mem) . patElemName . snd) $ zip [0..] ctx  of
-                Just (i, PatElem _ (MemMem size space)) ->
-                  ReturnsNewBlock space i (ext size) $
+                Just (i, PatElem _ (MemMem space)) ->
+                  ReturnsNewBlock space i $
                   existentialiseIxFun (map patElemName ctx) ixfun
                 _ -> ReturnsInBlock mem $ existentialiseIxFun [] ixfun
          )
@@ -863,13 +844,13 @@
     evalState (mapM addAttr ts) 0
     where addAttr (Prim bt) =
             return $ MemPrim bt
-          addAttr (Mem size space) =
-            return $ MemMem (Free size) space
+          addAttr (Mem space) =
+            return $ MemMem space
           addAttr t@(Array bt shape u)
             | existential t = do
-              i <- get <* modify (+2)
+              i <- get <* modify (+1)
               return $ MemArray bt shape u $ Just $
-                ReturnsNewBlock DefaultSpace (i+1) (Ext i) $
+                ReturnsNewBlock DefaultSpace i $
                 IxFun.iota $ map convert $ shapeDims shape
             | otherwise =
               return $ MemArray bt shape u Nothing
@@ -897,8 +878,8 @@
     MemArray et shape _ (ArrayIn mem ixfun) ->
       return $ MemArray et (fmap Free shape) NoUniqueness $
                Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun
-    MemMem size space ->
-      return $ MemMem (Free size) space
+    MemMem space ->
+      return $ MemMem space
 
 -- | The return information of an expression.  This can be seen as the
 -- "return type with memory annotations" of the expression.
@@ -948,7 +929,7 @@
       return [MemArray et (fmap Free shape) u $
               Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun]
     MemPrim pt -> return [MemPrim pt]
-    MemMem d space -> return [MemMem (Free d) space]
+    MemMem space -> return [MemMem space]
 
 expReturns (BasicOp (Update v _ _)) =
   pure <$> varReturns v
@@ -956,18 +937,15 @@
 expReturns (BasicOp op) =
   extReturns . staticShapes <$> primOpType op
 
-expReturns (DoLoop ctx val _ _) =
-  zipWithM typeWithAttr
-  (loopExtType (map (paramIdent . fst) ctx) (map (paramIdent . fst) val)) $ map fst val
+expReturns e@(DoLoop ctx val _ _) = do
+  t <- expExtType e
+  zipWithM typeWithAttr t $ map fst val
     where typeWithAttr t p =
             case (t, paramAttr p) of
               (Array bt shape u, MemArray _ _ _ (ArrayIn mem ixfun))
                 | Just (i, mem_p) <- isMergeVar mem,
-                  Mem mem_size space <- paramType mem_p ->
-                    let ext_size
-                          | Just (j, _) <- isMergeVar =<< subExpVar mem_size = Ext j
-                          | otherwise                                        = Free mem_size
-                    in return $ MemArray bt shape u $ Just $ ReturnsNewBlock space i ext_size ixfun'
+                  Mem space <- paramType mem_p ->
+                    return $ MemArray bt shape u $ Just $ ReturnsNewBlock space i ixfun'
                 | otherwise ->
                   return (MemArray bt shape u $
                           Just $ ReturnsInBlock mem ixfun')
@@ -1008,14 +986,11 @@
   opReturns op = extReturns <$> opType op
 
 instance OpReturns ExplicitMemory where
-  opReturns (Alloc size space) =
-    return [MemMem (Free size) space]
+  opReturns (Alloc _ space) =
+    return [MemMem space]
   opReturns (Inner (HostOp k@(Kernel _ _ _ body))) =
     zipWithM correct (kernelBodyResult body) =<< (extReturns <$> opType k)
     where correct (WriteReturn _ arr _) _ = varReturns arr
-          correct (KernelInPlaceReturn arr) _ =
-            extendedScope (varReturns arr)
-            (castScope $ scopeOf $ kernelBodyStms body)
           correct _ ret = return ret
   opReturns (Inner (HostOp (SegGenRed _ ops _ _))) =
     concat <$> mapM (mapM varReturns . genReduceDest) ops
@@ -1023,8 +998,8 @@
     extReturns <$> opType k
 
 instance OpReturns InKernel where
-  opReturns (Alloc size space) =
-    return [MemMem (Free size) space]
+  opReturns (Alloc _ space) =
+    return [MemMem space]
 
   opReturns (Inner (GroupStream _ _ lam _ _)) =
     forM (groupStreamAccParams lam) $ \param ->
@@ -1034,8 +1009,8 @@
         MemArray et shape _ (ArrayIn mem ixfun) ->
           return $ MemArray et (Shape $ map Free $ shapeDims shape) NoUniqueness $
           Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun
-        MemMem size space ->
-          return $ MemMem (Free size) space
+        MemMem space ->
+          return $ MemMem space
 
   opReturns (Inner (GroupScan _ _ input)) =
     mapM varReturns arrs
@@ -1078,10 +1053,8 @@
 
         correctDims (MemPrim t) =
           MemPrim t
-        correctDims (MemMem (Free se) space) =
-          MemMem (Free $ substSubExp se) space
-        correctDims (MemMem (Ext d) space) =
-          MemMem (Ext d) space
+        correctDims (MemMem space) =
+          MemMem space
         correctDims (MemArray et shape u memsummary) =
           MemArray et (correctShape shape) u $
           correctSummary memsummary
@@ -1090,8 +1063,8 @@
         correctDim (Ext i)   = Ext i
         correctDim (Free se) = Free $ substSubExp se
 
-        correctSummary (ReturnsNewBlock space i size ixfun) =
-          ReturnsNewBlock space i size ixfun
+        correctSummary (ReturnsNewBlock space i ixfun) =
+          ReturnsNewBlock space i ixfun
         correctSummary (ReturnsInBlock mem ixfun) =
           -- FIXME: we should also do a replacement in ixfun here.
           ReturnsInBlock mem' ixfun
diff --git a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs
--- a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs
+++ b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs
@@ -248,25 +248,24 @@
 -- | Compute the flat memory index for a complete set `inds` of array indices
 -- and a certain element size `elem_size`.
 index :: (IntegralExp num, Eq num) =>
-          IxFun num -> Indices num -> num -> num
+          IxFun num -> Indices num -> num
 index = indexFromLMADs . ixfunLMADs
   where indexFromLMADs :: (IntegralExp num, Eq num) =>
-                          NonEmpty (LMAD num) -> Indices num -> num -> num
-        indexFromLMADs (lmad :| []) inds elm_size = indexLMAD lmad inds elm_size
-        indexFromLMADs (lmad1 :| lmad2 : lmads) inds elm_size =
-          let i_flat   = indexLMAD lmad1 inds 1
+                          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 elm_size
+          in indexFromLMADs (lmad2 :| lmads) new_inds
 
         -- | Compute the flat index of an LMAD.
         indexLMAD :: (IntegralExp num, Eq num) =>
-                     LMAD num -> Indices num -> num -> num
-        indexLMAD lmad@(LMAD off dims) inds elm_size =
+                     LMAD num -> Indices num -> num
+        indexLMAD lmad@(LMAD off dims) inds =
           let prod = sum $ zipWith flatOneDim
                              (map (\(LMADDim s r n _ _) -> (s, r, n)) dims)
                              (permuteInv (lmadPermutation lmad) inds)
-              ind = off + prod
-          in if elm_size == 1 then ind else ind * elm_size
+          in off + prod
 
 -- | iota.
 iota :: IntegralExp num => Shape num -> IxFun num
diff --git a/src/Futhark/Representation/ExplicitMemory/Simplify.hs b/src/Futhark/Representation/ExplicitMemory/Simplify.hs
--- a/src/Futhark/Representation/ExplicitMemory/Simplify.hs
+++ b/src/Futhark/Representation/ExplicitMemory/Simplify.hs
@@ -3,7 +3,6 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE ConstraintKinds #-}
-{-# LANGUAGE LambdaCase #-}
 module Futhark.Representation.ExplicitMemory.Simplify
        ( simplifyExplicitMemory
        , simplifyStms
@@ -12,7 +11,6 @@
 
 import Control.Monad
 import qualified Data.Set as S
-import Data.Maybe
 import Data.List
 
 import qualified Futhark.Representation.AST.Syntax as AST
@@ -62,17 +60,13 @@
 
 -- | Getting the roots of what to hoist, for now only variable
 -- names that represent array and memory-block sizes.
-getShapeNames :: ExplicitMemorish lore =>
+getShapeNames :: (ExplicitMemorish lore, Op lore ~ MemOp op) =>
                  Stm (Wise lore) -> S.Set VName
-getShapeNames bnd =
-  let tps = map patElemType $ patternElements $ stmPattern bnd
-      ats = map (snd . patElemAttr) $ patternElements $ stmPattern bnd
-      nms = mapMaybe (\case
-                         MemMem (Var nm) _ -> Just nm
-                         MemArray _ _ _ (ArrayIn nm _) -> Just nm
-                         _ -> Nothing
-                     ) ats
-  in  S.fromList $ nms ++ subExpVars (concatMap arrayDims tps)
+getShapeNames stm =
+  let ts = map patElemType $ patternElements $ stmPattern stm
+  in freeIn (concatMap arrayDims ts) <>
+     case stmExp stm of Op (Alloc size _) -> freeIn size
+                        _                 -> mempty
 
 isAlloc0 :: Op lore ~ MemOp op => AST.Stm lore -> Bool
 isAlloc0 (Let _ _ (Op Alloc{})) = True
@@ -102,8 +96,8 @@
                           RuleIf unExistentialiseMemory] []
 
 -- | If a branch is returning some existential memory, but the size of
--- the array is existential, then we can create a block of the proper
--- size and always return there.
+-- the array is not existential, then we can create a block of the
+-- proper size and always return there.
 unExistentialiseMemory :: TopDownRuleIf (Wise InKernel)
 unExistentialiseMemory _ pat _ (cond, tbranch, fbranch, ifattr)
   | fixable <- foldl hasConcretisableMemory mempty $ patternElements pat,
@@ -111,12 +105,12 @@
 
       -- Create non-existential memory blocks big enough to hold the
       -- arrays.
-      (arr_to_mem, oldmem_to_mem, oldsize_to_size) <-
-        fmap unzip3 $ forM fixable $ \(arr_pe, oldmem, oldsize, space) -> do
+      (arr_to_mem, oldmem_to_mem) <-
+        fmap unzip $ forM fixable $ \(arr_pe, oldmem, space) -> do
           size <- letSubExp "size" =<<
                   toExp (arraySizeInBytesExp $ patElemType arr_pe)
           mem <- letExp "mem" $ Op $ Alloc size space
-          return ((patElemName arr_pe, mem), (oldmem, mem), (oldsize, size))
+          return ((patElemName arr_pe, mem), (oldmem, mem))
 
       -- Update the branches to contain Copy expressions putting the
       -- arrays where they are expected.
@@ -134,8 +128,6 @@
                 return $ Var v_copy
             | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem =
                 return $ Var mem
-            | Just size <- lookup (Var (patElemName pat_elem)) oldsize_to_size =
-                return size
           updateResult _ se =
             return se
       tbranch' <- updateBody tbranch
@@ -150,7 +142,7 @@
 
         hasConcretisableMemory fixable pat_elem
           | (_, MemArray _ shape _ (ArrayIn mem _)) <- patElemAttr pat_elem,
-            Just (j, Mem old_size space) <-
+            Just (j, Mem space) <-
               fmap patElemType <$> find ((mem==) . patElemName . snd)
                                         (zip [(0::Int)..] $ patternElements pat),
             Just tse <- maybeNth j $ bodyResult tbranch,
@@ -158,7 +150,7 @@
             mem `onlyUsedIn` patElemName pat_elem,
             all knownSize (shapeDims shape),
             fse /= tse =
-              (pat_elem, mem, old_size, space) : fixable
+              (pat_elem, mem, space) : fixable
           | otherwise =
               fixable
 unExistentialiseMemory _ _ _ _ = cannotSimplify
@@ -174,11 +166,11 @@
     Just (_, MemArray _ _ _ (ArrayIn srcmem src_ixfun)) <-
       ST.entryLetBoundAttr =<< ST.lookup v1 vtable,
 
-    Just (Mem _ src_space) <- ST.lookupType srcmem vtable,
+    Just (Mem src_space) <- ST.lookupType srcmem vtable,
 
     (_, MemArray _ _ _ (ArrayIn destmem dest_ixfun)) <- patElemAttr pat_elem,
 
-    Just (Mem _ dest_space) <- ST.lookupType destmem vtable,
+    Just (Mem dest_space) <- ST.lookupType destmem vtable,
 
     src_space == dest_space, dest_ixfun == src_ixfun =
 
diff --git a/src/Futhark/Representation/Kernels/Kernel.hs b/src/Futhark/Representation/Kernels/Kernel.hs
--- a/src/Futhark/Representation/Kernels/Kernel.hs
+++ b/src/Futhark/Representation/Kernels/Kernel.hs
@@ -10,14 +10,16 @@
 module Futhark.Representation.Kernels.Kernel
        ( Kernel(..)
        , kernelType
+       , kernelSpace
        , KernelDebugHints(..)
        , GenReduceOp(..)
+       , SegRedOp(..)
+       , segRedResults
        , KernelBody(..)
        , KernelSpace(..)
        , spaceDimensions
        , SpaceStructure(..)
        , scopeOfKernelSpace
-       , WhichThreads(..)
        , KernelResult(..)
        , kernelResultSubExp
        , KernelPath
@@ -50,7 +52,6 @@
 
 import Futhark.Representation.AST
 import qualified Futhark.Analysis.Alias as Alias
-import qualified Futhark.Analysis.UsageTable as UT
 import qualified Futhark.Analysis.SymbolTable as ST
 import Futhark.Analysis.PrimExp.Convert
 import qualified Futhark.Util.Pretty as PP
@@ -67,7 +68,6 @@
   (Aliases, removeLambdaAliases, removeBodyAliases, removeStmAliases)
 import Futhark.Representation.Kernels.KernelExp (SplitOrdering(..))
 import Futhark.Representation.Kernels.Sizes
-import Futhark.Analysis.Usage
 import qualified Futhark.TypeCheck as TC
 import Futhark.Analysis.Metrics
 import Futhark.Tools (partitionChunkedKernelLambdaParameters)
@@ -99,20 +99,49 @@
               }
   deriving (Eq, Ord, Show)
 
+data SegRedOp lore =
+  SegRedOp { segRedComm :: Commutativity
+           , segRedLambda :: Lambda lore
+           , segRedNeutral :: [SubExp]
+           , segRedShape :: Shape
+             -- ^ In case this operator is semantically a vectorised
+             -- operator (corresponding to a perfect map nest in the
+             -- SOACS representation), these are the logical
+             -- "dimensions".  This is used to generate more efficient
+             -- code.
+           }
+  deriving (Eq, Ord, Show)
+
+-- | How many reduction results are produced by these 'SegRedOp's?
+segRedResults :: [SegRedOp lore] -> Int
+segRedResults = sum . map (length . segRedNeutral)
+
 data Kernel lore
   = Kernel KernelDebugHints KernelSpace [Type] (KernelBody lore)
-  | SegRed KernelSpace Commutativity (Lambda lore) [SubExp] [Type] (Body lore)
+  | SegMap KernelSpace [Type] (KernelBody lore)
+  | SegRed KernelSpace [SegRedOp lore] [Type] (KernelBody lore)
     -- ^ The KernelSpace must always have at least two dimensions,
     -- implying that the result of a SegRed is always an array.
-  | SegGenRed KernelSpace [GenReduceOp lore] [Type] (Body lore)
+  | SegScan KernelSpace (Lambda lore) [SubExp] [Type] (KernelBody lore)
+  | SegGenRed KernelSpace [GenReduceOp lore] [Type] (KernelBody lore)
     deriving (Eq, Show, Ord)
 
+kernelSpace :: Kernel lore -> KernelSpace
+kernelSpace (Kernel _ kspace _ _) = kspace
+kernelSpace (SegMap kspace _ _) = kspace
+kernelSpace (SegRed kspace _ _ _) = kspace
+kernelSpace (SegScan kspace _ _ _ _) = kspace
+kernelSpace (SegGenRed kspace _ _ _) = kspace
+
 data KernelSpace = KernelSpace { spaceGlobalId :: VName
                                , spaceLocalId :: VName
                                , spaceGroupId :: VName
                                , spaceNumThreads :: SubExp
                                , spaceNumGroups :: SubExp
                                , spaceGroupSize :: SubExp -- flat group size
+                               , spaceNumVirtGroups :: SubExp
+                                 -- How many groups should we pretend
+                                 -- exist?
                                , spaceStructure :: SpaceStructure
                                -- TODO: document what this spaceStructure is
                                -- used for
@@ -151,7 +180,11 @@
 deriving instance Annotations lore => Show (KernelBody lore)
 deriving instance Annotations lore => Eq (KernelBody lore)
 
-data KernelResult = ThreadsReturn WhichThreads SubExp
+data KernelResult = ThreadsReturn SubExp
+                    -- ^ Each thread in the kernel space (which must
+                    -- be non-empty) returns this.
+                  | GroupsReturn SubExp
+                    -- ^ Each group returns this.
                   | WriteReturn
                     [SubExp] -- Size of array.  Must match number of dims.
                     VName -- Which array
@@ -163,21 +196,14 @@
                     SubExp -- Per-thread (max) chunk size.
                     (Maybe SubExp) -- Optional precalculated offset.
                     VName -- Chunk by this thread.
-                  | KernelInPlaceReturn VName -- HACK!
                   deriving (Eq, Show, Ord)
 
 kernelResultSubExp :: KernelResult -> SubExp
-kernelResultSubExp (ThreadsReturn _ se) = se
+kernelResultSubExp (ThreadsReturn se) = se
+kernelResultSubExp (GroupsReturn se) = se
 kernelResultSubExp (WriteReturn _ arr _) = Var arr
 kernelResultSubExp (ConcatReturns _ _ _ _ v) = Var v
-kernelResultSubExp (KernelInPlaceReturn v) = Var v
 
-data WhichThreads = AllThreads
-                  | OneResultPerGroup
-                  | ThreadsPerGroup [(VName,SubExp)] -- All threads before this one.
-                  | ThreadsInSpace
-                  deriving (Eq, Show, Ord)
-
 -- | Like 'Mapper', but just for 'Kernel's.
 data KernelMapper flore tlore m = KernelMapper {
     mapOnKernelSubExp :: SubExp -> m SubExp
@@ -203,20 +229,35 @@
 -- and is done left-to-right.
 mapKernelM :: (Applicative m, Monad m) =>
               KernelMapper flore tlore m -> Kernel flore -> m (Kernel tlore)
-mapKernelM tv (SegRed space comm red_op nes ts lam) =
+mapKernelM tv (SegMap space ts body) =
+  SegMap
+  <$> mapOnKernelSpace tv space
+  <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts
+  <*> mapOnKernelKernelBody tv body
+mapKernelM tv (SegRed space reds ts body) =
   SegRed
   <$> mapOnKernelSpace tv space
-  <*> pure comm
-  <*> mapOnKernelLambda tv red_op
+  <*> mapM onSegOp reds
+  <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts
+  <*> mapOnKernelKernelBody tv body
+  where onSegOp (SegRedOp comm red_op nes shape) =
+          SegRedOp comm
+          <$> mapOnKernelLambda tv red_op
+          <*> mapM (mapOnKernelSubExp tv) nes
+          <*> (Shape <$> mapM (mapOnKernelSubExp tv) (shapeDims shape))
+mapKernelM tv (SegScan space scan_op nes ts body) =
+  SegScan
+  <$> mapOnKernelSpace tv space
+  <*> mapOnKernelLambda tv scan_op
   <*> mapM (mapOnKernelSubExp tv) nes
   <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts
-  <*> mapOnKernelBody tv lam
+  <*> mapOnKernelKernelBody tv body
 mapKernelM tv (SegGenRed space ops ts body) =
   SegGenRed
   <$> mapOnKernelSpace tv space
   <*> mapM onGenRedOp ops
   <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts
-  <*> mapOnKernelBody tv body
+  <*> mapOnKernelKernelBody tv body
   where onGenRedOp (GenReduceOp w arrs nes shape op) =
           GenReduceOp <$> mapOnKernelSubExp tv w
           <*> mapM (mapOnKernelVName tv) arrs
@@ -234,11 +275,12 @@
 
 mapOnKernelSpace :: Monad f =>
                     KernelMapper flore tlore f -> KernelSpace -> f KernelSpace
-mapOnKernelSpace tv (KernelSpace gtid ltid gid num_threads num_groups group_size structure) =
+mapOnKernelSpace tv (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) =
   KernelSpace gtid ltid gid -- all in binding position
   <$> mapOnKernelSubExp tv num_threads
   <*> mapOnKernelSubExp tv num_groups
   <*> mapOnKernelSubExp tv group_size
+  <*> mapOnKernelSubExp tv virt_groups
   <*> mapOnKernelStructure structure
   where mapOnKernelStructure (FlatThreadSpace dims) =
           FlatThreadSpace <$> (zip gtids <$> mapM (mapOnKernelSubExp tv) gdim_sizes)
@@ -255,15 +297,15 @@
 mapOnKernelType _tv (Prim pt) = pure $ Prim pt
 mapOnKernelType tv (Array pt shape u) = Array pt <$> f shape <*> pure u
   where f (Shape dims) = Shape <$> mapM (mapOnKernelSubExp tv) dims
-mapOnKernelType _tv (Mem se s) = pure $ Mem se s
+mapOnKernelType _tv (Mem s) = pure $ Mem s
 
 instance (Attributes lore, FreeIn (LParamAttr lore)) =>
          FreeIn (Kernel lore) where
   freeIn e = execWriter $ mapKernelM free e
     where walk f x = tell (f x) >> return x
           free = KernelMapper { mapOnKernelSubExp = walk freeIn
-                              , mapOnKernelLambda = walk freeInLambda
-                              , mapOnKernelBody = walk freeInBody
+                              , mapOnKernelLambda = walk freeIn
+                              , mapOnKernelBody = walk freeIn
                               , mapOnKernelVName = walk freeIn
                               , mapOnKernelLParam = walk freeIn
                               , mapOnKernelKernelBody = walk freeIn
@@ -309,22 +351,16 @@
   where m = walkKernelMapper f
 
 instance FreeIn KernelResult where
-  freeIn (ThreadsReturn which what) = freeIn which <> freeIn what
+  freeIn (GroupsReturn what) = freeIn what
+  freeIn (ThreadsReturn what) = freeIn what
   freeIn (WriteReturn rws arr res) = freeIn rws <> freeIn arr <> freeIn res
   freeIn (ConcatReturns o w per_thread_elems moffset v) =
     freeIn o <> freeIn w <> freeIn per_thread_elems <> freeIn moffset <> freeIn v
-  freeIn (KernelInPlaceReturn what) = freeIn what
 
-instance FreeIn WhichThreads where
-  freeIn AllThreads = mempty
-  freeIn OneResultPerGroup = mempty
-  freeIn (ThreadsPerGroup limit) = freeIn limit
-  freeIn ThreadsInSpace = mempty
-
 instance Attributes lore => FreeIn (KernelBody lore) where
   freeIn (KernelBody attr stms res) =
     (freeIn attr <> free_in_stms <> free_in_res) `S.difference` bound_in_stms
-    where free_in_stms = fold $ fmap freeInStm stms
+    where free_in_stms = fold $ fmap freeIn stms
           free_in_res = freeIn res
           bound_in_stms = fold $ fmap boundByStm stms
 
@@ -336,8 +372,10 @@
     (substituteNames subst res)
 
 instance Substitute KernelResult where
-  substituteNames subst (ThreadsReturn who se) =
-    ThreadsReturn (substituteNames subst who) (substituteNames subst se)
+  substituteNames subst (GroupsReturn se) =
+    GroupsReturn $ substituteNames subst se
+  substituteNames subst (ThreadsReturn se) =
+    ThreadsReturn $ substituteNames subst se
   substituteNames subst (WriteReturn rws arr res) =
     WriteReturn
     (substituteNames subst rws) (substituteNames subst arr)
@@ -349,24 +387,16 @@
     (substituteNames subst per_thread_elems)
     (substituteNames subst moffset)
     (substituteNames subst v)
-  substituteNames subst (KernelInPlaceReturn what) =
-    KernelInPlaceReturn (substituteNames subst what)
 
-instance Substitute WhichThreads where
-  substituteNames _ AllThreads = AllThreads
-  substituteNames _ OneResultPerGroup = OneResultPerGroup
-  substituteNames _ ThreadsInSpace = ThreadsInSpace
-  substituteNames subst (ThreadsPerGroup limit) =
-    ThreadsPerGroup $ substituteNames subst limit
-
 instance Substitute KernelSpace where
-  substituteNames subst (KernelSpace gtid ltid gid num_threads num_groups group_size structure) =
+  substituteNames subst (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) =
     KernelSpace (substituteNames subst gtid)
     (substituteNames subst ltid)
     (substituteNames subst gid)
     (substituteNames subst num_threads)
     (substituteNames subst num_groups)
     (substituteNames subst group_size)
+    (substituteNames subst virt_groups)
     (substituteNames subst structure)
 
 instance Substitute SpaceStructure where
@@ -400,11 +430,8 @@
 instance Rename KernelResult where
   rename = substituteRename
 
-instance Rename WhichThreads where
-  rename = substituteRename
-
 scopeOfKernelSpace :: KernelSpace -> Scope lore
-scopeOfKernelSpace (KernelSpace gtid ltid gid _ _ _ structure) =
+scopeOfKernelSpace (KernelSpace gtid ltid gid _ _ _ _ structure) =
   M.fromList $ zip ([gtid, ltid, gid] ++ structure') $ repeat $ IndexInfo Int32
   where structure' = case structure of
                        FlatThreadSpace dims -> map fst dims
@@ -416,34 +443,39 @@
   rename = mapKernelM renamer
     where renamer = KernelMapper rename rename rename rename rename rename
 
+
+kernelResultShape :: KernelSpace -> Type -> KernelResult -> Type
+kernelResultShape _ t (WriteReturn rws _ _) =
+  t `arrayOfShape` Shape rws
+kernelResultShape space t (GroupsReturn _) =
+  t `arrayOfRow` spaceNumGroups space
+kernelResultShape space t (ThreadsReturn _) =
+  foldr (flip arrayOfRow . snd) t $ spaceDimensions space
+kernelResultShape _ t (ConcatReturns _ w _ _ _) =
+  t `arrayOfRow` w
+
 kernelType :: Kernel lore -> [Type]
 kernelType (Kernel _ space ts body) =
-  zipWith resultShape ts $ kernelBodyResult body
-  where dims = map snd $ spaceDimensions space
-        num_groups = spaceNumGroups space
-        num_threads = spaceNumThreads space
-        resultShape t (WriteReturn rws _ _) =
-          t `arrayOfShape` Shape rws
-        resultShape t (ThreadsReturn AllThreads _) =
-          t `arrayOfRow` num_threads
-        resultShape t (ThreadsReturn OneResultPerGroup _) =
-          t `arrayOfRow` num_groups
-        resultShape t (ThreadsReturn (ThreadsPerGroup limit) _) =
-          t `arrayOfShape` Shape (map snd limit) `arrayOfRow` num_groups
-        resultShape t (ThreadsReturn ThreadsInSpace _) =
-          foldr (flip arrayOfRow) t dims
-        resultShape t (ConcatReturns _ w _ _ _) =
-          t `arrayOfRow` w
-        resultShape t KernelInPlaceReturn{} =
-          t
+  zipWith (kernelResultShape space) ts $ kernelBodyResult body
 
-kernelType (SegRed space _ _ nes ts _) =
-  map (`arrayOfShape` Shape outer_dims) red_ts ++
-  map (`arrayOfShape` Shape dims) map_ts
-  where (red_ts, map_ts) = splitAt (length nes) ts
-        dims = map snd $ spaceDimensions space
-        outer_dims = init dims
+kernelType (SegMap space ts body) =
+  zipWith (kernelResultShape space) ts $ kernelBodyResult body
 
+kernelType (SegRed space reds ts body) =
+  red_ts ++
+  zipWith (kernelResultShape space) map_ts
+  (drop (length red_ts) $ kernelBodyResult body)
+  where map_ts = drop (length red_ts) ts
+        segment_dims = init $ map snd $ spaceDimensions space
+        red_ts = do
+          op <- reds
+          let shape = Shape segment_dims <> segRedShape op
+          map (`arrayOfShape` shape) (lambdaReturnType $ segRedLambda op)
+
+kernelType (SegScan space _ _ ts _) =
+  map (`arrayOfShape` Shape dims) ts
+  where dims = map snd $ spaceDimensions space
+
 kernelType (SegGenRed space ops _ _) = do
   op <- ops
   let shape = Shape (segment_dims <> [genReduceWidth op]) <> genReduceShape op
@@ -468,10 +500,15 @@
     mconcat (map consumedByReturn (kernelBodyResult kbody))
     where consumedByReturn (WriteReturn _ a _) = S.singleton a
           consumedByReturn _                   = mempty
-  consumedInOp (SegGenRed _ ops _ body) =
+  consumedInOp (SegGenRed _ ops _ kbody) =
     S.fromList (concatMap genReduceDest ops) <>
-    consumedInBody body
-  consumedInOp _ = mempty
+    consumedInKernelBody kbody
+  consumedInOp (SegMap _ _ kbody) =
+    consumedInKernelBody kbody
+  consumedInOp (SegRed _ _ _ kbody) =
+    consumedInKernelBody kbody
+  consumedInOp (SegScan _ _ _ _ kbody) =
+    consumedInKernelBody kbody
 
 aliasAnalyseKernelBody :: (Attributes lore,
                            CanBeAliased (Op lore)) =>
@@ -479,15 +516,7 @@
                        -> KernelBody (Aliases lore)
 aliasAnalyseKernelBody (KernelBody attr stms res) =
   let Body attr' stms' _ = Alias.analyseBody $ Body attr stms []
-  in KernelBody attr' stms' $ map aliasAnalyseKernelResult res
-  where aliasAnalyseKernelResult (ThreadsReturn which what) =
-          ThreadsReturn which what
-        aliasAnalyseKernelResult (WriteReturn rws arr res') =
-          WriteReturn rws arr res'
-        aliasAnalyseKernelResult (ConcatReturns o w per_thread_elems moffset v) =
-          ConcatReturns o w per_thread_elems moffset v
-        aliasAnalyseKernelResult (KernelInPlaceReturn what) =
-          KernelInPlaceReturn what
+  in KernelBody attr' stms' res
 
 instance (Attributes lore,
           Attributes (Aliases lore),
@@ -530,17 +559,7 @@
           addKernelBodyRanges (KernelBody attr stms res) =
             Range.analyseStms stms $ \stms' -> do
             let attr' = (mkBodyRanges stms $ map kernelResultSubExp res, attr)
-            res' <- mapM addKernelResultRanges res
-            return $ KernelBody attr' stms' res'
-
-          addKernelResultRanges (ThreadsReturn which what) =
-            return $ ThreadsReturn which what
-          addKernelResultRanges (WriteReturn rws arr res) =
-            return $ WriteReturn rws arr res
-          addKernelResultRanges (ConcatReturns o w per_thread_elems moffset v) =
-            return $ ConcatReturns o w per_thread_elems moffset v
-          addKernelResultRanges (KernelInPlaceReturn what) =
-            return $ KernelInPlaceReturn what
+            return $ KernelBody attr' stms' res
 
 instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (Kernel lore) where
   type OpWithWisdom (Kernel lore) = Kernel (Wise lore)
@@ -559,19 +578,11 @@
 
 instance Attributes lore => ST.IndexOp (Kernel lore) where
   indexOp vtable k (Kernel _ space _ kbody) is = do
-    ThreadsReturn which se <- maybeNth k $ kernelBodyResult kbody
-
-    prim_table <- case (which, is) of
-      (AllThreads, [i]) ->
-        Just $ M.singleton (spaceGlobalId space) (i,mempty)
-      (ThreadsInSpace, _)
-        | (gtids, _) <- unzip $ spaceDimensions space,
-          length gtids == length is ->
-            Just $ M.fromList $ zip gtids $ zip is $ repeat mempty
-      _ ->
-        Nothing
-
-    let prim_table' = foldl expandPrimExpTable prim_table $ kernelBodyStms kbody
+    ThreadsReturn se <- maybeNth k $ kernelBodyResult kbody
+    let (gtids, _) = unzip $ spaceDimensions space
+    guard $ length gtids == length is
+    let prim_table = M.fromList $ zip gtids $ zip is $ repeat mempty
+        prim_table' = foldl expandPrimExpTable prim_table $ kernelBodyStms kbody
     case se of
       Var v -> M.lookup v prim_table'
       _ -> Nothing
@@ -591,15 +602,6 @@
 
   indexOp _ _ _ _ = Nothing
 
-instance Aliased lore => UsageInOp (Kernel lore) where
-  usageInOp (Kernel _ _ _ kbody) =
-    mconcat $ map UT.consumedUsage $ S.toList $ consumedInKernelBody kbody
-  usageInOp (SegRed _ _ _ _ _ body) =
-    mconcat $ map UT.consumedUsage $ S.toList $ consumedInBody body
-  usageInOp (SegGenRed _ ops _ body) =
-    mconcat $ map UT.consumedUsage $ S.toList (consumedInBody body) <>
-    concatMap genReduceDest ops
-
 consumedInKernelBody :: Aliased lore =>
                         KernelBody lore -> Names
 consumedInKernelBody (KernelBody attr stms _) =
@@ -607,53 +609,53 @@
 
 typeCheckKernel :: TC.Checkable lore => Kernel (Aliases lore) -> TC.TypeM lore ()
 
-typeCheckKernel (SegRed space _ red_op nes ts body) = do
+typeCheckKernel (SegMap space ts kbody) = do
   checkSpace space
   mapM_ TC.checkType ts
-
-  ne_ts <- mapM subExpType nes
+  TC.binding (scopeOfKernelSpace space) $ checkKernelBody ts kbody
 
-  let asArg t = (t, mempty)
-  TC.binding (scopeOfKernelSpace space) $ do
-    TC.checkLambda red_op $ map asArg $ ne_ts ++ ne_ts
-    unless (lambdaReturnType red_op == ne_ts &&
-            take (length nes) ts == ne_ts) $
-      TC.bad $ TC.TypeError
-      "SegRed: wrong type for reduction or neutral elements."
+typeCheckKernel (SegRed space reds ts body) =
+  checkScanRed space reds' ts body
+  where reds' = zip3
+                (map segRedLambda reds)
+                (map segRedNeutral reds)
+                (map segRedShape reds)
 
-    TC.checkLambdaBody ts body
+typeCheckKernel (SegScan space scan_op nes ts body) =
+  checkScanRed space [(scan_op, nes, mempty)] ts body
 
 typeCheckKernel (SegGenRed space ops ts body) = do
   checkSpace space
   mapM_ TC.checkType ts
 
   TC.binding (scopeOfKernelSpace space) $ do
-    forM_ ops $ \(GenReduceOp dest_w dests nes shape op) -> do
+    nes_ts <- forM ops $ \(GenReduceOp dest_w dests nes shape op) -> do
       TC.require [Prim int32] dest_w
       nes' <- mapM TC.checkArg nes
       mapM_ (TC.require [Prim int32]) $ shapeDims shape
 
       -- Operator type must match the type of neutral elements.
       let stripVecDims = stripArray $ shapeRank shape
-      TC.checkLambda op $ map (TC.noArgAliases .first stripVecDims) $ nes' ++ nes'
+      TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes'
       let nes_t = map TC.argType nes'
-      unless (nes_t == map (`arrayOfShape` shape) (lambdaReturnType op)) $
+      unless (nes_t == lambdaReturnType op) $
         TC.bad $ TC.TypeError $ "SegGenRed operator has return type " ++
         prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++
         prettyTuple nes_t
 
       -- Arrays must have proper type.
-      let dest_shape = Shape $ segment_dims <> [dest_w]
+      let dest_shape = Shape (segment_dims <> [dest_w]) <> shape
       forM_ (zip nes_t dests) $ \(t, dest) -> do
         TC.requireI [t `arrayOfShape` dest_shape] dest
         TC.consume =<< TC.lookupAliases dest
 
-    TC.checkLambdaBody ts body
+      return $ map (`arrayOfShape` shape) nes_t
 
+    checkKernelBody ts body
+
     -- Return type of bucket function must be an index for each
     -- operation followed by the values to write.
-    nes_ts <- concat <$> mapM (mapM subExpType . genReduceNeutral) ops
-    let bucket_ret_t = replicate (length ops) (Prim int32) ++ nes_ts
+    let bucket_ret_t = replicate (length ops) (Prim int32) ++ concat nes_ts
     unless (bucket_ret_t == ts) $
       TC.bad $ TC.TypeError $ "SegGenRed body has return type " ++
       prettyTuple ts ++ " but should have type " ++
@@ -668,17 +670,21 @@
 
   TC.binding (scopeOfKernelSpace space) $
     checkKernelBody kts kbody
-  where checkKernelBody ts (KernelBody (_, attr) stms res) = do
-          TC.checkBodyLore attr
-          TC.checkStms stms $ do
-            unless (length ts == length res) $
-              TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++
-              ", but body returns " ++ show (length res) ++ " values."
-            zipWithM_ checkKernelResult res ts
 
-        checkKernelResult (ThreadsReturn which what) t = do
-          checkWhich which
+checkKernelBody :: TC.Checkable lore =>
+                   [Type] -> KernelBody (Aliases lore) -> TC.TypeM lore ()
+checkKernelBody ts (KernelBody (_, attr) stms kres) = do
+  TC.checkBodyLore attr
+  TC.checkStms stms $ do
+    unless (length ts == length kres) $
+      TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++
+      ", but body returns " ++ show (length kres) ++ " values."
+    zipWithM_ checkKernelResult kres ts
+
+  where checkKernelResult (GroupsReturn what) t =
           TC.require [t] what
+        checkKernelResult (ThreadsReturn what) t =
+          TC.require [t] what
         checkKernelResult (WriteReturn rws arr res) t = do
           mapM_ (TC.require [Prim int32]) rws
           arr_t <- lookupType arr
@@ -700,19 +706,45 @@
           vt <- lookupType v
           unless (vt == t `arrayOfRow` arraySize 0 vt) $
             TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v
-        checkKernelResult (KernelInPlaceReturn what) t =
-          TC.requireI [t] what
 
-        checkWhich AllThreads = return ()
-        checkWhich OneResultPerGroup = return ()
-        checkWhich ThreadsInSpace = return ()
-        checkWhich (ThreadsPerGroup limit) = do
-          mapM_ (TC.requireI [Prim int32] . fst) limit
-          mapM_ (TC.require [Prim int32] . snd) limit
+checkScanRed :: TC.Checkable lore =>
+                KernelSpace
+             -> [(Lambda (Aliases lore), [SubExp], Shape)]
+             -> [Type]
+             -> KernelBody (Aliases lore)
+             -> TC.TypeM lore ()
+checkScanRed space ops ts kbody = do
+  checkSpace space
+  mapM_ TC.checkType ts
 
+  TC.binding (scopeOfKernelSpace space) $ do
+    ne_ts <- forM ops $ \(lam, nes, shape) -> do
+      mapM_ (TC.require [Prim int32]) $ shapeDims shape
+      nes' <- mapM TC.checkArg nes
+
+      -- Operator type must match the type of neutral elements.
+      let stripVecDims = stripArray $ shapeRank shape
+      TC.checkLambda lam $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes'
+      let nes_t = map TC.argType nes'
+
+      unless (lambdaReturnType lam == nes_t) $
+        TC.bad $ TC.TypeError "wrong type for operator or neutral elements."
+
+      return $ map (`arrayOfShape` shape) nes_t
+
+    let expecting = concat ne_ts
+        got = take (length expecting) ts
+    unless (expecting == got) $
+      TC.bad $ TC.TypeError $
+      "Wrong return for body (does not match neutral elements; expected " ++
+      pretty expecting ++ "; found " ++
+      pretty got ++ ")"
+
+    checkKernelBody ts kbody
+
 checkSpace :: TC.Checkable lore => KernelSpace -> TC.TypeM lore ()
-checkSpace (KernelSpace _ _ _ num_threads num_groups group_size structure) = do
-  mapM_ (TC.require [Prim int32]) [num_threads,num_groups,group_size]
+checkSpace (KernelSpace _ _ _ num_threads num_groups group_size virt_groups structure) = do
+  mapM_ (TC.require [Prim int32]) [num_threads,num_groups,group_size,virt_groups]
   case structure of
     FlatThreadSpace dims ->
       mapM_ (TC.require [Prim int32] . snd) dims
@@ -723,28 +755,50 @@
 instance OpMetrics (Op lore) => OpMetrics (Kernel lore) where
   opMetrics (Kernel _ _ _ kbody) =
     inside "Kernel" $ kernelBodyMetrics kbody
-    where kernelBodyMetrics :: KernelBody lore -> MetricsM ()
-          kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms
-  opMetrics (SegRed _ _ red_op _ _ body) =
-    inside "SegRed" $ lambdaMetrics red_op >> bodyMetrics body
+  opMetrics (SegMap _ _ body) =
+    inside "SegMap" $ kernelBodyMetrics body
+  opMetrics (SegRed _ reds _ body) =
+    inside "SegRed" $ do mapM_ (lambdaMetrics . segRedLambda) reds
+                         kernelBodyMetrics body
+  opMetrics (SegScan _ scan_op _ _ body) =
+    inside "SegScan" $ lambdaMetrics scan_op >> kernelBodyMetrics body
   opMetrics (SegGenRed _ ops _ body) =
     inside "SegGenRed" $ do mapM_ (lambdaMetrics . genReduceOp) ops
-                            bodyMetrics body
+                            kernelBodyMetrics body
 
+kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM ()
+kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms
+
 instance PrettyLore lore => PP.Pretty (Kernel lore) where
   ppr (Kernel desc space ts body) =
     text "kernel" <+> text (kernelName desc) <>
     PP.align (ppr space) <+>
     PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body)
 
-  ppr (SegRed space comm red_op nes ts body) =
-    text name <> PP.parens (ppr red_op <> PP.comma </>
-                             PP.braces (PP.commasep $ map ppr nes)) </>
+  ppr (SegMap space ts body) =
+    text "segmap" <>
     PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+>
     PP.nestedBlock "{" "}" (ppr body)
-    where name = case comm of Commutative    -> "segred_comm"
-                              Noncommutative -> "segred"
 
+  ppr (SegRed space reds ts body) =
+    text "segred" <>
+    PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp reds)) </>
+    PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+>
+    PP.nestedBlock "{" "}" (ppr body)
+    where ppOp (SegRedOp comm lam nes shape) =
+            PP.braces (PP.commasep $ map ppr nes) <> PP.comma </>
+            ppr shape <> PP.comma </>
+            comm' <> ppr lam
+            where comm' = case comm of Commutative -> text "commutative "
+                                       Noncommutative -> mempty
+
+
+  ppr (SegScan space scan_op nes ts body) =
+    text "segscan" <> PP.parens (ppr scan_op <> PP.comma </>
+                                 PP.braces (PP.commasep $ map ppr nes)) </>
+    PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+>
+    PP.nestedBlock "{" "}" (ppr body)
+
   ppr (SegGenRed space ops ts body) =
     text "seggenred" <>
     PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) </>
@@ -758,9 +812,10 @@
             ppr op
 
 instance Pretty KernelSpace where
-  ppr (KernelSpace f_gtid f_ltid gid num_threads num_groups group_size structure) =
+  ppr (KernelSpace f_gtid f_ltid gid num_threads num_groups group_size virt_groups structure) =
     parens (commasep [text "num groups:" <+> ppr num_groups,
                       text "group size:" <+> ppr group_size,
+                      text "virt_num_groups:" <+> ppr virt_groups,
                       text "num threads:" <+> ppr num_threads,
                       text "global TID ->" <+> ppr f_gtid,
                       text "local TID ->" <+> ppr f_ltid,
@@ -782,14 +837,10 @@
     text "return" <+> PP.braces (PP.commasep $ map ppr res)
 
 instance Pretty KernelResult where
-  ppr (ThreadsReturn AllThreads what) =
-    ppr what
-  ppr (ThreadsReturn OneResultPerGroup what) =
-    text "group" <+> "returns" <+> ppr what
-  ppr (ThreadsReturn (ThreadsPerGroup limit) what) =
-    text "thread <" <+> ppr limit <+> text "returns" <+> ppr what
-  ppr (ThreadsReturn ThreadsInSpace what) =
-    text "thread in space returns" <+> ppr what
+  ppr (GroupsReturn what) =
+    text "group returns" <+> ppr what
+  ppr (ThreadsReturn what) =
+    text "thread returns" <+> ppr what
   ppr (WriteReturn rws arr res) =
     ppr arr <+> text "with" <+> PP.apply (map ppRes res)
     where ppRes (is, e) =
@@ -803,8 +854,6 @@
                            SplitStrided stride -> text "Strided" <> parens (ppr stride)
           offset_text = case offset of Nothing -> ""
                                        Just se -> "," <+> "offset=" <> ppr se
-  ppr (KernelInPlaceReturn what) =
-    text "kernel returns" <+> ppr what
 
 --- Host operations
 
@@ -907,7 +956,7 @@
 
   ppr (CmpSizeLe name size_class x) =
     text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+>
-    text "<" <+> ppr x
+    text "<=" <+> ppr x
 
   ppr (HostOp op) = ppr op
 
@@ -916,12 +965,6 @@
   opMetrics GetSizeMax{} = seen "GetSizeMax"
   opMetrics CmpSizeLe{} = seen "CmpSizeLe"
   opMetrics (HostOp op) = opMetrics op
-
-instance UsageInOp inner => UsageInOp (HostOp lore inner) where
-  usageInOp GetSize{} = mempty
-  usageInOp GetSizeMax{} = mempty
-  usageInOp CmpSizeLe{} = mempty
-  usageInOp (HostOp op) = usageInOp op
 
 typeCheckHostOp :: TC.Checkable lore =>
                    (inner -> TC.TypeM lore ())
diff --git a/src/Futhark/Representation/Kernels/KernelExp.hs b/src/Futhark/Representation/Kernels/KernelExp.hs
--- a/src/Futhark/Representation/Kernels/KernelExp.hs
+++ b/src/Futhark/Representation/Kernels/KernelExp.hs
@@ -27,13 +27,11 @@
 
 import qualified Futhark.Analysis.Alias as Alias
 import qualified Futhark.Analysis.Range as Range
-import qualified Futhark.Analysis.UsageTable as UT
 import Futhark.Representation.Aliases
 import Futhark.Representation.Ranges
 import Futhark.Transform.Substitute
 import Futhark.Transform.Rename
 import Futhark.Optimise.Simplify.Lore
-import Futhark.Analysis.Usage
 import Futhark.Analysis.Metrics
 import qualified Futhark.Analysis.ScalExp as SE
 import qualified Futhark.Analysis.SymbolTable as ST
@@ -185,20 +183,20 @@
   freeIn (SplitSpace o w i elems_per_thread) =
     freeIn o <> freeIn [w, i, elems_per_thread]
   freeIn (Combine (CombineSpace scatter cspace) ts active body) =
-    freeIn scatter <> freeIn (map snd cspace) <> freeIn ts <> freeIn active <> freeInBody body
+    freeIn scatter <> freeIn (map snd cspace) <> freeIn ts <> freeIn active <> freeIn body
   freeIn (GroupReduce w lam input) =
-    freeIn w <> freeInLambda lam <> freeIn input
+    freeIn w <> freeIn lam <> freeIn input
   freeIn (GroupScan w lam input) =
-    freeIn w <> freeInLambda lam <> freeIn input
+    freeIn w <> freeIn lam <> freeIn input
   freeIn (GroupStream w maxchunk lam accs arrs) =
     freeIn w <> freeIn maxchunk <> freeIn lam <> freeIn accs <> freeIn arrs
   freeIn (GroupGenReduce w dests op bucket values locks) =
-    freeIn w <> freeIn dests <> freeInLambda op <> freeIn bucket <> freeIn values <> freeIn locks
+    freeIn w <> freeIn dests <> freeIn op <> freeIn bucket <> freeIn values <> freeIn locks
   freeIn (Barrier ses) = freeIn ses
 
 instance Attributes lore => FreeIn (GroupStreamLambda lore) where
   freeIn (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) =
-    freeInBody body `S.difference` bound_here
+    freeIn body `S.difference` bound_here
     where bound_here = S.fromList $
                        chunk_offset : chunk_size :
                        map paramName (acc_params ++ arr_params)
@@ -437,12 +435,6 @@
 
 instance ST.IndexOp (KernelExp lore) where
 
-instance Aliased lore => UsageInOp (KernelExp lore) where
-  usageInOp (Combine cspace _ _ body) =
-    mconcat $ map UT.consumedUsage $ S.toList (consumedInBody body) <>
-    [ arr | (_, _, arr) <- cspaceScatter cspace ]
-  usageInOp _ = mempty
-
 instance OpMetrics (Op lore) => OpMetrics (KernelExp lore) where
   opMetrics SplitSpace{} = seen "SplitSpace"
   opMetrics Combine{} = seen "Combine"
@@ -552,12 +544,9 @@
 checkScanOrReduce w lam input = do
   TC.require [Prim int32] w
   let (nes, arrs) = unzip input
-      asArg t = (t, mempty)
   neargs <- mapM TC.checkArg nes
   arrargs <- TC.checkSOACArrayArgs w arrs
-  TC.checkLambda lam $
-    map asArg [Prim int32, Prim int32] ++
-    map TC.noArgAliases (neargs ++ arrargs)
+  TC.checkLambda lam $ map TC.noArgAliases (neargs ++ arrargs)
 
 instance Scoped lore (GroupStreamLambda lore) where
   scopeOf (GroupStreamLambda chunk_size chunk_offset acc_params arr_params _) =
diff --git a/src/Futhark/Representation/Kernels/Simplify.hs b/src/Futhark/Representation/Kernels/Simplify.hs
--- a/src/Futhark/Representation/Kernels/Simplify.hs
+++ b/src/Futhark/Representation/Kernels/Simplify.hs
@@ -34,6 +34,7 @@
 import qualified Futhark.Analysis.SymbolTable as ST
 import qualified Futhark.Analysis.UsageTable as UT
 import Futhark.Analysis.Rephrase (castStm)
+import Futhark.Util (chunks)
 
 simpleKernels :: Simplify.SimpleOps Kernels
 simpleKernels = Simplify.bindableSimpleOps (simplifyKernelOp simpleInKernel inKernelEnv)
@@ -83,23 +84,56 @@
         scope_vtable = ST.fromScope scope
         bound_here = S.fromList $ M.keys scope
 
-simplifyKernelOp mk_ops env (HostOp (SegRed space comm red_op nes ts body)) = do
+simplifyKernelOp mk_ops env (HostOp (SegMap space ts body)) = do
   space' <- Engine.simplify space
-  nes' <- mapM Engine.simplify nes
   ts' <- mapM Engine.simplify ts
+
+  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body
+
+  return (HostOp $ SegMap space' ts' body',
+          body_hoisted)
+
+simplifyKernelOp mk_ops env (HostOp (SegRed space reds ts body)) = do
+  ts' <- mapM Engine.simplify ts
+  space' <- Engine.simplify space
   outer_vtable <- Engine.askVtable
 
-  (red_op', red_op_hoisted) <-
+  (reds', reds_hoisted) <- fmap unzip $ forM reds $ \(SegRedOp comm lam nes shape) -> do
+    (lam', hoisted) <- Engine.subSimpleM (mk_ops space) env outer_vtable $
+                       Engine.localVtable (<>scope_vtable) $
+                       Engine.simplifyLambda lam $
+                       replicate (length nes * 2) Nothing
+    shape' <- Engine.simplify shape
+    nes' <- mapM Engine.simplify nes
+    return (SegRedOp comm lam' nes' shape', hoisted)
+  red_op_hoisted' <- mapM processHoistedStm $ mconcat reds_hoisted
+
+  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body
+
+  return (HostOp $ SegRed space' reds' ts' body',
+          red_op_hoisted' <> body_hoisted)
+
+  where scope_vtable = ST.fromScope scope
+        scope = scopeOfKernelSpace space
+
+simplifyKernelOp mk_ops env (HostOp (SegScan space scan_op nes ts body)) = do
+  outer_vtable <- Engine.askVtable
+  space' <- Engine.simplify space
+
+  (scan_op', scan_op_hoisted) <-
     Engine.subSimpleM (mk_ops space) env outer_vtable $
     Engine.localVtable (<>scope_vtable) $
-    Engine.simplifyLambda red_op $ replicate (length nes * 2) Nothing
-  red_op_hoisted' <- mapM processHoistedStm red_op_hoisted
+    Engine.simplifyLambda scan_op $ replicate (length nes * 2) Nothing
+  scan_op_hoisted' <- mapM processHoistedStm scan_op_hoisted
 
-  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env ts body
+  nes' <- mapM Engine.simplify nes
+  ts' <- mapM Engine.simplify ts
 
-  return (HostOp $ SegRed space' comm red_op' nes' ts' body',
-          red_op_hoisted' <> body_hoisted)
+  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body
 
+  return (HostOp $ SegScan space' scan_op' nes' ts' body',
+          scan_op_hoisted' <> body_hoisted)
+
   where scope_vtable = ST.fromScope scope
         scope = scopeOfKernelSpace space
 
@@ -125,7 +159,7 @@
 
   red_op_hoisted' <- mapM processHoistedStm $ mconcat ops_hoisted
 
-  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env ts body
+  (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env body
 
   return (HostOp $ SegGenRed space' ops' ts' body',
           red_op_hoisted' <> body_hoisted)
@@ -148,9 +182,9 @@
                   RetType lore ~ RetType outerlore,
                   BranchType lore ~ BranchType outerlore) =>
                  KernelSpace -> Simplify.SimpleOps lore -> Engine.Env lore
-              -> [Type] -> Body lore
-              -> Engine.SimpleM outerlore (Body (Wise lore), Stms (Wise outerlore))
-hoistFromBody kspace ops env ts body = do
+              -> KernelBody lore
+              -> Engine.SimpleM outerlore (KernelBody (Wise lore), Stms (Wise outerlore))
+hoistFromBody kspace ops env kbody = do
   outer_vtable <- Engine.askVtable
 
   ((body_stms, body_res), body_hoisted) <-
@@ -161,11 +195,11 @@
                         `Engine.orIf` Engine.isOp
                         `Engine.orIf` par_blocker
                         `Engine.orIf` Engine.isConsumed) $
-        Engine.simplifyBody (replicate (length ts) Observe) body
+        simplifyKernelBodyM kbody
 
   body_hoisted' <- mapM processHoistedStm body_hoisted
 
-  return (mkWiseBody () body_stms body_res,
+  return (mkWiseKernelBody () body_stms body_res,
           body_hoisted')
 
   where scope_vtable = ST.fromScope scope
@@ -191,11 +225,7 @@
 mkWiseKernelBody attr bnds res =
   let Body attr' _ _ = mkWiseBody attr bnds res_vs
   in KernelBody attr' bnds res
-  where res_vs = map resValue res
-        resValue (ThreadsReturn _ se) = se
-        resValue (WriteReturn _ arr _) = Var arr
-        resValue (ConcatReturns _ _ _ _ v) = Var v
-        resValue (KernelInPlaceReturn v) = Var v
+  where res_vs = map kernelResultSubExp res
 
 inKernelEnv :: Engine.Env InKernel
 inKernelEnv = Engine.emptyEnv inKernelRules Simplify.noExtraHoistBlockers
@@ -321,11 +351,12 @@
   return (GroupStreamLambda block_size block_offset acc_params' arr_params' body', hoisted)
 
 instance Engine.Simplifiable KernelSpace where
-  simplify (KernelSpace gtid ltid gid num_threads num_groups group_size structure) =
+  simplify (KernelSpace gtid ltid gid num_threads num_groups group_size virt_groups structure) =
     KernelSpace gtid ltid gid
     <$> Engine.simplify num_threads
     <*> Engine.simplify num_groups
     <*> Engine.simplify group_size
+    <*> Engine.simplify virt_groups
     <*> Engine.simplify structure
 
 instance Engine.Simplifiable SpaceStructure where
@@ -341,8 +372,10 @@
     where (gtids, gdims, ltids, ldims) = unzip4 dims
 
 instance Engine.Simplifiable KernelResult where
-  simplify (ThreadsReturn threads what) =
-    ThreadsReturn <$> Engine.simplify threads <*> Engine.simplify what
+  simplify (GroupsReturn what) =
+    GroupsReturn <$> Engine.simplify what
+  simplify (ThreadsReturn what) =
+    ThreadsReturn <$> Engine.simplify what
   simplify (WriteReturn ws a res) =
     WriteReturn <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res
   simplify (ConcatReturns o w pte moffset what) =
@@ -352,16 +385,7 @@
     <*> Engine.simplify pte
     <*> Engine.simplify moffset
     <*> Engine.simplify what
-  simplify (KernelInPlaceReturn what) =
-    KernelInPlaceReturn <$> Engine.simplify what
 
-instance Engine.Simplifiable WhichThreads where
-  simplify AllThreads = pure AllThreads
-  simplify OneResultPerGroup = pure OneResultPerGroup
-  simplify ThreadsInSpace = pure ThreadsInSpace
-  simplify (ThreadsPerGroup limit) =
-    ThreadsPerGroup <$> mapM Engine.simplify limit
-
 instance BinderOps (Wise Kernels) where
   mkExpAttrB = bindableMkExpAttrB
   mkBodyB = bindableMkBodyB
@@ -374,9 +398,10 @@
 
 kernelRules :: RuleBook (Wise Kernels)
 kernelRules = standardRules <>
-              ruleBook [RuleOp removeInvariantKernelResults]
-                       [RuleOp distributeKernelResults,
-                        RuleBasicOp removeUnnecessaryCopy]
+              ruleBook [ RuleOp removeInvariantKernelResults
+                       , RuleOp mergeSegRedOps]
+                       [ RuleOp distributeKernelResults
+                       , RuleBasicOp removeUnnecessaryCopy]
 
 fuseStreamIota :: TopDownRuleOp (Wise InKernel)
 fuseStreamIota vtable pat _ (GroupStream w max_chunk lam accs arrs)
@@ -428,22 +453,14 @@
   where isInvariant Constant{} = True
         isInvariant (Var v) = isJust $ ST.lookup v vtable
 
-        num_threads = spaceNumThreads space
         space_dims = map snd $ spaceDimensions space
 
-        checkForInvarianceResult (_, pe, ThreadsReturn threads se)
-          | isInvariant se =
-              case threads of
-                AllThreads -> do
-                  letBindNames_ [patElemName pe] $ BasicOp $
-                    Replicate (Shape [num_threads]) se
-                  return False
-                ThreadsInSpace -> do
-                  let rep a d = BasicOp . Replicate (Shape [d]) <$> letSubExp "rep" a
-                  letBindNames_ [patElemName pe] =<<
-                    foldM rep (BasicOp (SubExp se)) (reverse space_dims)
-                  return False
-                _ -> return True
+        checkForInvarianceResult (_, pe, ThreadsReturn se)
+          | isInvariant se = do
+              let rep a d = BasicOp . Replicate (Shape [d]) <$> letSubExp "rep" a
+              letBindNames_ [patElemName pe] =<<
+                foldM rep (BasicOp (SubExp se)) (reverse space_dims)
+              return False
         checkForInvarianceResult _ =
           return True
 removeInvariantKernelResults _ _ _ _ = cannotSimplify
@@ -465,7 +482,7 @@
   addStm $ Let (Pattern [] kpes') attr $ Op $ HostOp $
     Kernel desc kspace kts' $ mkWiseKernelBody () (stmsFromList $ reverse kstms_rev) kres'
   where
-    free_in_kstms = fold $ fmap freeInStm kstms
+    free_in_kstms = fold $ fmap freeIn kstms
 
     distribute (kpes', kts', kres', kstms_rev) bnd
       | Let (Pattern [] [pe]) _ (BasicOp (Index arr slice)) <- bnd,
@@ -501,9 +518,57 @@
           | (kpes'', kts'', kres'') <- unzip3 kpes_and_kres ->
               Just (kpe, kpes'', kts'', kres'')
         _ -> Nothing
-      where matches (_, _, kre) = kre == ThreadsReturn ThreadsInSpace (Var $ patElemName pe)
+      where matches (_, _, kre) = kre == ThreadsReturn (Var $ patElemName pe)
 distributeKernelResults _ _ _ _ = cannotSimplify
 
+-- If a SegRed contains two reduction operations that have the same
+-- vector shape, merge them together.  This saves on communication
+-- overhead, but can in principle lead to more local memory usage.
+mergeSegRedOps :: TopDownRuleOp (Wise Kernels)
+mergeSegRedOps _ (Pattern [] pes) _ (HostOp (SegRed space ops ts kbody))
+  | length ops > 1,
+    op_groupings <- groupBy sameShape $ zip ops $ chunks (map (length . segRedNeutral) ops) $
+                    zip3 red_pes red_ts red_res,
+    any ((>1) . length) op_groupings = do
+      let (ops', aux) = unzip $ mapMaybe combineOps op_groupings
+          (red_pes', red_ts', red_res') = unzip3 $ concat aux
+          pes' = red_pes' ++ map_pes
+          ts' = red_ts' ++ map_ts
+          kbody' = kbody { kernelBodyResult = red_res' ++ map_res }
+      letBind_ (Pattern [] pes') $ Op $ HostOp $ SegRed space ops' ts' kbody'
+  where (red_pes, map_pes) = splitAt (segRedResults ops) pes
+        (red_ts, map_ts) = splitAt (segRedResults ops) ts
+        (red_res, map_res) = splitAt (segRedResults ops) $ kernelBodyResult kbody
+
+        sameShape (op1, _) (op2, _) = segRedShape op1 == segRedShape op2
+
+        combineOps :: [(SegRedOp (Wise InKernel), [a])]
+                   -> Maybe (SegRedOp (Wise InKernel), [a])
+        combineOps [] = Nothing
+        combineOps (x:xs) = Just $ foldl' combine x xs
+
+        combine (op1, op1_aux) (op2, op2_aux) =
+          let lam1 = segRedLambda op1
+              lam2 = segRedLambda op2
+              (op1_xparams, op1_yparams) =
+                splitAt (length (segRedNeutral op1)) $ lambdaParams lam1
+              (op2_xparams, op2_yparams) =
+                splitAt (length (segRedNeutral op2)) $ lambdaParams lam2
+              lam = Lambda { lambdaParams = op1_xparams ++ op2_xparams ++
+                                            op1_yparams ++ op2_yparams
+                           , lambdaReturnType = lambdaReturnType lam1 ++ lambdaReturnType lam2
+                           , lambdaBody =
+                               mkBody (bodyStms (lambdaBody lam1) <> bodyStms (lambdaBody lam2)) $
+                               bodyResult (lambdaBody lam1) <> bodyResult (lambdaBody lam2)
+                           }
+          in (SegRedOp { segRedComm = segRedComm op1 <> segRedComm op2
+                       , segRedLambda = lam
+                       , segRedNeutral = segRedNeutral op1 ++ segRedNeutral op2
+                       , segRedShape = segRedShape op1 -- Same as shape of op2 due to the grouping.
+                       },
+               op1_aux ++ op2_aux)
+mergeSegRedOps _ _ _ _ = cannotSimplify
+
 simplifyKnownIterationStream :: TopDownRuleOp (Wise InKernel)
 -- Remove GroupStreams over single-element arrays.  Not much to stream
 -- here, and no information to exploit.
@@ -537,7 +602,7 @@
       letBind_ pat $ Op $ GroupStream w maxchunk lam' accs arrs'
   where GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam
 
-        isUsed = (`S.member` freeInBody body)
+        isUsed = (`S.member` freeIn body)
 removeUnusedStreamInputs _ _ _ _ = cannotSimplify
 
 inKernelRules :: RuleBook (Wise InKernel)
diff --git a/src/Futhark/Representation/Kernels/Sizes.hs b/src/Futhark/Representation/Kernels/Sizes.hs
--- a/src/Futhark/Representation/Kernels/Sizes.hs
+++ b/src/Futhark/Representation/Kernels/Sizes.hs
@@ -16,6 +16,9 @@
                | SizeGroup
                | SizeNumGroups
                | SizeTile
+               | SizeLocalMemory
+               -- ^ Likely not useful on its own, but querying the
+               -- maximum can be handy.
                deriving (Eq, Ord, Show)
 
 instance Pretty SizeClass where
@@ -25,3 +28,4 @@
   ppr SizeGroup = text "group_size"
   ppr SizeNumGroups = text "num_groups"
   ppr SizeTile = text "tile_size"
+  ppr SizeLocalMemory = text "local_memory"
diff --git a/src/Futhark/Representation/Primitive.hs b/src/Futhark/Representation/Primitive.hs
--- a/src/Futhark/Representation/Primitive.hs
+++ b/src/Futhark/Representation/Primitive.hs
@@ -88,7 +88,7 @@
 import           Prelude
 
 import           Futhark.Util.Pretty
-import           Futhark.Util (roundFloat, roundDouble)
+import           Futhark.Util (roundFloat, roundDouble, lgamma, lgammaf, tgamma, tgammaf)
 
 -- | An integer type, ordered by size.  Note that signedness is not a
 -- property of the type, but a property of the operations performed on
@@ -832,8 +832,9 @@
   , f32 "asin32" asin, f64 "asin64" asin
   , f32 "acos32" acos, f64 "acos64" acos
   , f32 "atan32" atan, f64 "atan64" atan
-
   , f32 "round32" roundFloat, f64 "round64" roundDouble
+  , f32 "gamma32" tgammaf, f64 "gamma64" tgamma
+  , f32 "lgamma32" lgammaf, f64 "lgamma64" lgamma
 
   , ("atan2_32",
      ([FloatType Float32, FloatType Float32], FloatType Float32,
diff --git a/src/Futhark/Representation/SOACS/SOAC.hs b/src/Futhark/Representation/SOACS/SOAC.hs
--- a/src/Futhark/Representation/SOACS/SOAC.hs
+++ b/src/Futhark/Representation/SOACS/SOAC.hs
@@ -10,7 +10,9 @@
        , ScremaForm(..)
        , GenReduceOp(..)
        , Scan
-       , Reduce
+       , Reduce(..)
+       , redResults
+       , singleReduce
 
          -- * Utility
        , getStreamOrder
@@ -63,7 +65,6 @@
 import Futhark.Representation.Ranges (Ranges, removeLambdaRanges)
 import Futhark.Representation.AST.Attributes.Ranges
 import Futhark.Representation.Aliases (Aliases, removeLambdaAliases)
-import Futhark.Analysis.Usage
 import qualified Futhark.Analysis.SymbolTable as ST
 import Futhark.Analysis.PrimExp.Convert
 import qualified Futhark.TypeCheck as TC
@@ -73,8 +74,8 @@
 import Futhark.Util (maybeNth, chunks, splitAt3)
 
 data SOAC lore =
-    Stream SubExp (StreamForm lore) (LambdaT lore) [VName]
-  | Scatter SubExp (LambdaT lore) [VName] [(SubExp, Int, VName)]
+    Stream SubExp (StreamForm lore) (Lambda lore) [VName]
+  | Scatter SubExp (Lambda lore) [VName] [(SubExp, Int, VName)]
     -- Scatter <cs> <length> <lambda> <original index and value arrays>
     --
     -- <input/output arrays along with their sizes and number of
@@ -90,11 +91,11 @@
     --     [index_0, index_1, ..., index_n, value_0, value_1, ..., value_n]
     --
     -- This must be consistent along all Scatter-related optimisations.
-  | GenReduce SubExp [GenReduceOp lore] (LambdaT lore) [VName]
+  | GenReduce SubExp [GenReduceOp lore] (Lambda lore) [VName]
     -- GenReduce <length> <dest-arrays-and-ops> <bucket fun> <input arrays>
     --
     -- The first SubExp is the length of the input arrays. The first
-    -- list describes the operations to perform.  The 'LambdaT' is the
+    -- list describes the operations to perform.  The 'Lambda' is the
     -- bucket function.  Finally comes the input images.
   | Screma SubExp (ScremaForm lore) [VName]
     -- ^ A combination of scan, reduction, and map.  The first
@@ -108,12 +109,12 @@
 data GenReduceOp lore = GenReduceOp { genReduceWidth :: SubExp
                                     , genReduceDest :: [VName]
                                     , genReduceNeutral :: [SubExp]
-                                    , genReduceOp :: LambdaT lore
+                                    , genReduceOp :: Lambda lore
                                     }
                       deriving (Eq, Ord, Show)
 
 data StreamForm lore  =
-    Parallel StreamOrd Commutativity (LambdaT lore) [SubExp]
+    Parallel StreamOrd Commutativity (Lambda lore) [SubExp]
   | Sequential [SubExp]
   deriving (Eq, Ord, Show)
 
@@ -121,18 +122,40 @@
 -- except the input arrays).
 data ScremaForm lore = ScremaForm
                          (Scan lore)
-                         (Reduce lore)
-                         (LambdaT lore)
+                         [Reduce lore]
+                         (Lambda lore)
   deriving (Eq, Ord, Show)
 
-type Scan lore = (LambdaT lore, [SubExp])
-type Reduce lore = (Commutativity, LambdaT lore, [SubExp])
+type Scan lore = (Lambda lore, [SubExp])
+data Reduce lore = Reduce { redComm :: Commutativity
+                          , redLambda :: Lambda lore
+                          , redNeutral :: [SubExp]
+                          }
+                   deriving (Eq, Ord, Show)
 
+-- | How many reduction results are produced by these 'Reduce's?
+redResults :: [Reduce lore] -> Int
+redResults = sum . map (length . redNeutral)
+
+-- | Combine multiple reduction operators to a single operator.
+singleReduce :: Bindable lore => [Reduce lore] -> Reduce lore
+singleReduce reds =
+  let red_nes = concatMap redNeutral reds
+      red_lam =
+        let xParams red = take (length (redNeutral red)) (lambdaParams (redLambda red))
+            yParams red = drop (length (redNeutral red)) (lambdaParams (redLambda red))
+        in Lambda { lambdaParams = concatMap xParams reds ++ concatMap yParams reds
+                  , lambdaReturnType = concatMap (lambdaReturnType . redLambda) reds
+                  , lambdaBody = mkBody (mconcat (map (bodyStms . lambdaBody . redLambda) reds))
+                                 (concatMap (bodyResult . lambdaBody . redLambda) reds)
+                  }
+  in Reduce (mconcat (map redComm reds)) red_lam red_nes
+
 scremaType :: SubExp -> ScremaForm lore -> [Type]
-scremaType w (ScremaForm (scan_lam, _scan_nes) (_, red_lam, _red_nes) map_lam) =
+scremaType w (ScremaForm (scan_lam, _scan_nes) reds map_lam) =
   map (`arrayOfRow` w) scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps
   where scan_tps = lambdaReturnType scan_lam
-        red_tps  = lambdaReturnType red_lam
+        red_tps  = concatMap (lambdaReturnType . redLambda) reds
         map_tps  = drop (length scan_tps + length red_tps) $ lambdaReturnType map_lam
 
 -- | Construct a lambda that takes parameters of the given types and
@@ -182,36 +205,36 @@
         (red_x_params, red_y_params) = splitAt m $ lambdaParams red_fun
 
 -- | A lambda with no parameters that returns no values.
-nilFn :: Bindable lore => LambdaT lore
+nilFn :: Bindable lore => Lambda lore
 nilFn = Lambda mempty (mkBody mempty mempty) mempty
 
-isNilFn :: LambdaT lore -> Bool
+isNilFn :: Lambda lore -> Bool
 isNilFn (Lambda ps body ts) =
   null ps && null ts &&
   null (bodyStms body) && null (bodyResult body)
 
-scanomapSOAC :: Bindable lore =>
-                Lambda lore -> [SubExp] -> Lambda lore -> ScremaForm lore
-scanomapSOAC lam nes = ScremaForm (lam, nes) (mempty, nilFn, mempty)
+scanomapSOAC :: Lambda lore -> [SubExp] -> Lambda lore -> ScremaForm lore
+scanomapSOAC lam nes = ScremaForm (lam, nes) []
 
 redomapSOAC :: Bindable lore =>
-               Commutativity -> Lambda lore -> [SubExp] -> Lambda lore -> ScremaForm lore
-redomapSOAC comm lam nes = ScremaForm (nilFn, mempty) (comm, lam, nes)
+               [Reduce lore] -> Lambda lore -> ScremaForm lore
+redomapSOAC = ScremaForm (nilFn, mempty)
 
 scanSOAC :: (Bindable lore, MonadFreshNames m) =>
             Lambda lore -> [SubExp] -> m (ScremaForm lore)
 scanSOAC lam nes = scanomapSOAC lam nes <$> mkIdentityLambda (lambdaReturnType lam)
 
 reduceSOAC :: (Bindable lore, MonadFreshNames m) =>
-              Commutativity -> Lambda lore -> [SubExp] -> m (ScremaForm lore)
-reduceSOAC comm lam nes = redomapSOAC comm lam nes <$> mkIdentityLambda (lambdaReturnType lam)
+              [Reduce lore] -> m (ScremaForm lore)
+reduceSOAC reds = redomapSOAC reds <$> mkIdentityLambda ts
+  where ts = concatMap (lambdaReturnType . redLambda) reds
 
 mapSOAC :: Bindable lore => Lambda lore -> ScremaForm lore
-mapSOAC = ScremaForm (nilFn, mempty) (mempty, nilFn, mempty)
+mapSOAC = ScremaForm (nilFn, mempty) []
 
 isScanomapSOAC :: ScremaForm lore -> Maybe (Lambda lore, [SubExp], Lambda lore)
-isScanomapSOAC (ScremaForm (scan_lam, scan_nes) (_, _, red_nes) map_lam) = do
-  guard $ null red_nes
+isScanomapSOAC (ScremaForm (scan_lam, scan_nes) reds map_lam) = do
+  guard $ null reds
   guard $ not $ null scan_nes
   return (scan_lam, scan_nes, map_lam)
 
@@ -220,21 +243,21 @@
                      guard $ isIdentityLambda map_lam
                      return (scan_lam, scan_nes)
 
-isRedomapSOAC :: ScremaForm lore -> Maybe (Commutativity, Lambda lore, [SubExp], Lambda lore)
-isRedomapSOAC (ScremaForm (_, scan_nes) (comm, red_lam, red_nes) map_lam) = do
+isRedomapSOAC :: ScremaForm lore -> Maybe ([Reduce lore], Lambda lore)
+isRedomapSOAC (ScremaForm (_, scan_nes) reds map_lam) = do
   guard $ null scan_nes
-  guard $ not $ null red_nes
-  return (comm, red_lam, red_nes, map_lam)
+  guard $ not $ null reds
+  return (reds, map_lam)
 
-isReduceSOAC :: ScremaForm lore -> Maybe (Commutativity, Lambda lore, [SubExp])
-isReduceSOAC form = do (comm, red_lam, red_nes, map_lam) <- isRedomapSOAC form
+isReduceSOAC :: ScremaForm lore -> Maybe [Reduce lore]
+isReduceSOAC form = do (reds, map_lam) <- isRedomapSOAC form
                        guard $ isIdentityLambda map_lam
-                       return (comm, red_lam, red_nes)
+                       return reds
 
 isMapSOAC :: ScremaForm lore -> Maybe (Lambda lore)
-isMapSOAC (ScremaForm (_, scan_nes) (_, _, red_nes) map_lam) = do
+isMapSOAC (ScremaForm (_, scan_nes) reds map_lam) = do
   guard $ null scan_nes
-  guard $ null red_nes
+  guard $ null reds
   return map_lam
 
 -- | Like 'Mapper', but just for 'SOAC's.
@@ -283,11 +306,13 @@
               <*> mapOnSOACLambda tv op) ops
   <*> mapOnSOACLambda tv bucket_fun
   <*> mapM (mapOnSOACVName tv) imgs
-mapSOACM tv (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) =
+mapSOACM tv (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) =
   Screma <$> mapOnSOACSubExp tv w <*>
   (ScremaForm <$>
    ((,) <$> mapOnSOACLambda tv scan_lam <*> mapM (mapOnSOACSubExp tv) scan_nes) <*>
-   ((,,) comm <$> mapOnSOACLambda tv red_lam <*> mapM (mapOnSOACSubExp tv) red_nes) <*>
+   forM reds (\(Reduce comm red_lam red_nes) ->
+                Reduce comm <$> mapOnSOACLambda tv red_lam <*>
+               mapM (mapOnSOACSubExp tv) red_nes) <*>
    mapOnSOACLambda tv map_lam)
   <*> mapM (mapOnSOACVName tv) arrs
 mapSOACM tv (CmpThreshold what s) = CmpThreshold <$> mapOnSOACSubExp tv what <*> pure s
@@ -296,7 +321,7 @@
   freeIn = execWriter . mapSOACM free
     where walk f x = tell (f x) >> return x
           free = SOACMapper { mapOnSOACSubExp = walk freeIn
-                            , mapOnSOACLambda = walk freeInLambda
+                            , mapOnSOACLambda = walk freeIn
                             , mapOnSOACVName = walk freeIn
                             }
 
@@ -363,7 +388,7 @@
     S.fromList $ concatMap genReduceDest ops
   consumedInOp CmpThreshold{} = mempty
 
-mapGenReduceOp :: (LambdaT flore -> LambdaT tlore)
+mapGenReduceOp :: (Lambda flore -> Lambda tlore)
                -> GenReduceOp flore -> GenReduceOp tlore
 mapGenReduceOp f (GenReduceOp w dests nes lam) =
   GenReduceOp w dests nes $ f lam
@@ -384,12 +409,13 @@
   addOpAliases (GenReduce len ops bucket_fun imgs) =
     GenReduce len (map (mapGenReduceOp Alias.analyseLambda) ops)
     (Alias.analyseLambda bucket_fun) imgs
-  addOpAliases (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) =
+  addOpAliases (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) =
     Screma w (ScremaForm
                 (Alias.analyseLambda scan_lam, scan_nes)
-                (comm, Alias.analyseLambda red_lam, red_nes)
+                (map onRed reds)
                 (Alias.analyseLambda map_lam))
                arrs
+    where onRed red = red { redLambda = Alias.analyseLambda $ redLambda red }
   addOpAliases (CmpThreshold what s) = CmpThreshold what s
 
   removeOpAliases = runIdentity . mapSOACM remove
@@ -402,8 +428,7 @@
 
 substNamesInType :: M.Map VName SubExp -> Type -> Type
 substNamesInType _ tp@(Prim _) = tp
-substNamesInType subs (Mem se space) =
-  Mem (substNamesInSubExp subs se) space
+substNamesInType _ (Mem space) = Mem space
 substNamesInType subs (Array btp shp u) =
   let shp' = Shape $ map (substNamesInSubExp subs) (shapeDims shp)
   in  Array btp shp' u
@@ -436,12 +461,14 @@
   addOpRanges (GenReduce len ops bucket_fun imgs) =
     GenReduce len (map (mapGenReduceOp $ Range.runRangeM . Range.analyseLambda) ops)
     (Range.runRangeM $ Range.analyseLambda bucket_fun) imgs
-  addOpRanges (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) =
+  addOpRanges (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) =
     Screma w (ScremaForm
                 (Range.runRangeM $ Range.analyseLambda scan_lam, scan_nes)
-                (comm, Range.runRangeM $ Range.analyseLambda red_lam, red_nes)
+                (map onRed reds)
                 (Range.runRangeM $ Range.analyseLambda map_lam))
                arrs
+    where onRed red = red { redLambda = Range.runRangeM $ Range.analyseLambda $
+                                        redLambda red }
   addOpRanges (CmpThreshold what s) = CmpThreshold what s
 
 instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (SOAC lore) where
@@ -458,8 +485,8 @@
     case se of
       Var v -> M.lookup v arr_indexes'
       _ -> Nothing
-      where lambdaAndSubExp (Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) map_lam) arrs) =
-              nthMapOut (length scan_nes + length red_nes) map_lam arrs
+      where lambdaAndSubExp (Screma _ (ScremaForm (_, scan_nes) reds map_lam) arrs) =
+              nthMapOut (length scan_nes + redResults reds) map_lam arrs
             lambdaAndSubExp _ =
               Nothing
 
@@ -487,10 +514,6 @@
               | otherwise = lift Nothing
   indexOp _ _ _ _ = Nothing
 
-instance Aliased lore => UsageInOp (SOAC lore) where
-  usageInOp (Screma _ (ScremaForm _ _ f) arrs) = usageInLambda f arrs
-  usageInOp _ = mempty
-
 typeCheckSOAC :: TC.Checkable lore => SOAC (Aliases lore) -> TC.TypeM lore ()
 typeCheckSOAC (CmpThreshold what _) = TC.require [Prim int32] what
 typeCheckSOAC (Stream size form lam arrexps) = do
@@ -611,16 +634,24 @@
     prettyTuple (lambdaReturnType bucket_fun) ++ " but should have type " ++
     prettyTuple bucket_ret_t
 
-typeCheckSOAC (Screma w (ScremaForm (scan_lam, scan_nes) (_, red_lam, red_nes) map_lam) arrs) = do
+typeCheckSOAC (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) = do
   TC.require [Prim int32] w
   arrs' <- TC.checkSOACArrayArgs w arrs
   scan_nes' <- mapM TC.checkArg scan_nes
-  red_nes' <- mapM TC.checkArg red_nes
   TC.checkLambda map_lam $ map TC.noArgAliases arrs'
   TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes'
-  TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes'
+
+  red_nes' <- fmap concat $ forM reds $ \(Reduce _ red_lam red_nes) -> do
+    red_nes' <- mapM TC.checkArg red_nes
+    let red_t = map TC.argType red_nes'
+    TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes'
+    unless (red_t == lambdaReturnType red_lam) $
+      TC.bad $ TC.TypeError $ "Reduce function returns type " ++
+      prettyTuple (lambdaReturnType red_lam) ++ " but neutral element has type " ++
+      prettyTuple red_t
+    return red_nes'
+
   let scan_t = map TC.argType scan_nes'
-      red_t = map TC.argType red_nes'
       map_lam_ts = lambdaReturnType map_lam
 
   unless (scan_t == lambdaReturnType scan_lam) $
@@ -628,12 +659,7 @@
     prettyTuple (lambdaReturnType scan_lam) ++ " but neutral element has type " ++
     prettyTuple scan_t
 
-  unless (red_t == lambdaReturnType red_lam) $
-    TC.bad $ TC.TypeError $ "Reduce function returns type " ++
-    prettyTuple (lambdaReturnType red_lam) ++ " but neutral element has type " ++
-    prettyTuple red_t
-
-  unless (take (length scan_nes + length red_nes) map_lam_ts ==
+  unless (take (length scan_nes + length red_nes') map_lam_ts ==
           map TC.argType (scan_nes'++ red_nes')) $
     TC.bad $ TC.TypeError $ "Map function return type " ++ prettyTuple map_lam_ts ++
     " wrong for given scan and reduction functions."
@@ -654,9 +680,10 @@
     inside "Scatter" $ lambdaMetrics lam
   opMetrics (GenReduce _len ops bucket_fun _imgs) =
     inside "GenReduce" $ mapM_ (lambdaMetrics . genReduceOp) ops >> lambdaMetrics bucket_fun
-  opMetrics (Screma _ (ScremaForm (scan_lam, _) (_, red_lam, _) map_lam) _) =
-    inside "Screma" $
-    lambdaMetrics scan_lam >> lambdaMetrics red_lam >> lambdaMetrics map_lam
+  opMetrics (Screma _ (ScremaForm (scan_lam, _) reds map_lam) _) =
+    inside "Screma" $ do lambdaMetrics scan_lam
+                         mapM_ (lambdaMetrics . redLambda) reds
+                         lambdaMetrics map_lam
   opMetrics CmpThreshold{} = seen "CmpThreshold"
 
 instance PrettyLore lore => PP.Pretty (SOAC lore) where
@@ -677,42 +704,50 @@
     ppSOAC "scatter" len [lam] (Just (map Var ivs)) (map (\(_,n,a) -> (n,a)) as)
   ppr (GenReduce len ops bucket_fun imgs) =
     ppGenReduce len ops bucket_fun imgs
-  ppr (Screma w (ScremaForm (scan_lam, scan_nes) (_, red_lam, red_nes) map_lam) arrs)
-    | isNilFn scan_lam, null scan_nes,
-      isNilFn red_lam, null red_nes =
-        text "map" <> parens (ppr w <> comma </>
-                              ppr map_lam <> comma </>
-                              commasep (map ppr arrs))
+  ppr (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs)
+    | isNilFn scan_lam, null scan_nes, null reds =
+        text "map" <>
+        parens (ppr w <> comma </>
+                ppr map_lam <> comma </>
+                commasep (map ppr arrs))
 
     | isNilFn scan_lam, null scan_nes =
-        text "redomap" <> parens (ppr w <> comma </>
-                                  ppr red_lam <> comma </>
-                                  commasep (map ppr red_nes) <> comma </>
-                                  ppr map_lam <> comma </>
-                                  commasep (map ppr arrs))
+        text "redomap" <>
+        parens (ppr w <> comma </>
+                PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma </>
+                ppr map_lam <> comma </>
+                commasep (map ppr arrs))
 
-    | isNilFn red_lam, null red_nes =
-        text "scanomap" <> parens (ppr w <> comma </>
-                                   ppr scan_lam <> comma </>
-                                   commasep (map ppr scan_nes) <> comma </>
-                                   ppr map_lam <> comma </>
-                                   commasep (map ppr arrs))
+    | null reds =
+        text "scanomap" <>
+        parens (ppr w <> comma </>
+                ppr scan_lam <> comma </>
+                commasep (map ppr scan_nes) <> comma </>
+                ppr map_lam <> comma </>
+                commasep (map ppr arrs))
 
   ppr (Screma w form arrs) = ppScrema w form arrs
   ppr (CmpThreshold what s) = text "cmpThreshold(" <> ppr what <> comma PP.<+> text (show s) <> text ")"
 
 ppScrema :: (PrettyLore lore, Pretty inp) =>
               SubExp -> ScremaForm lore -> [inp] -> Doc
-ppScrema w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs =
-  text s <> parens (ppr w <> comma </>
-                      ppr scan_lam <> comma </>
-                      PP.braces (commasep $ map ppr scan_nes) </>
-                      ppr red_lam <> comma </>
-                      PP.braces (commasep $ map ppr red_nes) </>
-                      ppr map_lam <> comma </>
-                      commasep (map ppr arrs))
-    where s = case comm of Noncommutative -> "screma"
-                           Commutative -> "scremaComm"
+ppScrema w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs =
+  text "screma" <>
+  parens (ppr w <> comma </>
+          ppr scan_lam <> comma </>
+          PP.braces (commasep $ map ppr scan_nes) </>
+          PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma </>
+          ppr map_lam <> comma </>
+          commasep (map ppr arrs))
+
+ppComm :: Commutativity -> Doc
+ppComm Noncommutative = mempty
+ppComm Commutative = text "commutative "
+
+instance PrettyLore lore => Pretty (Reduce lore) where
+  ppr (Reduce comm red_lam red_nes) =
+    ppComm comm <> ppr red_lam <> comma </>
+    PP.braces (commasep $ map ppr red_nes)
 
 ppGenReduce :: (PrettyLore lore, Pretty inp) =>
                SubExp -> [GenReduceOp lore] -> Lambda lore -> [inp] -> Doc
diff --git a/src/Futhark/Representation/SOACS/Simplify.hs b/src/Futhark/Representation/SOACS/Simplify.hs
--- a/src/Futhark/Representation/SOACS/Simplify.hs
+++ b/src/Futhark/Representation/SOACS/Simplify.hs
@@ -10,10 +10,15 @@
        , simplifyStms
 
        , simpleSOACS
+
+       , soacRules
        )
 where
 
 import Control.Monad
+import Control.Monad.Identity
+import Control.Monad.State
+import Control.Monad.Writer
 import Data.Foldable
 import Data.Either
 import Data.List
@@ -103,19 +108,19 @@
   (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun $ map Just imgs
   return (GenReduce w' ops' bfun' imgs', mconcat hoisted <> bfun_hoisted)
 
-simplifySOAC (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) = do
+simplifySOAC (Screma w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) = do
   (scan_lam', scan_lam_hoisted) <-
     Engine.simplifyLambda scan_lam $ replicate (length scan_nes) Nothing
-  (red_lam', red_lam_hoisted) <-
-    Engine.simplifyLambda red_lam $ replicate (length red_nes) Nothing
+  (reds', reds_hoisted) <- fmap unzip $ forM reds $ \(Reduce comm lam nes) -> do
+    (lam', hoisted) <- Engine.simplifyLambda lam $ replicate (length nes) Nothing
+    nes' <- Engine.simplify nes
+    return (Reduce comm lam' nes', hoisted)
   (map_lam', map_lam_hoisted) <- Engine.simplifyLambda map_lam $ map Just arrs
   (,) <$> (Screma <$> Engine.simplify w <*>
-           (ScremaForm <$>
-             ((,) scan_lam' <$> Engine.simplify scan_nes) <*>
-             ((,,) comm red_lam' <$> Engine.simplify red_nes) <*>
-             pure map_lam') <*>
+           (ScremaForm <$> ((,) scan_lam' <$> Engine.simplify scan_nes) <*>
+             pure reds' <*> pure map_lam') <*>
             Engine.simplify arrs) <*>
-    pure (scan_lam_hoisted <> red_lam_hoisted <> map_lam_hoisted)
+    pure (scan_lam_hoisted <> mconcat reds_hoisted <> map_lam_hoisted)
 
 instance BinderOps (Wise SOACS) where
   mkExpAttrB = bindableMkExpAttrB
@@ -148,12 +153,15 @@
 soacRules = standardRules <> ruleBook topDownRules bottomUpRules
 
 topDownRules :: [TopDownRule (Wise SOACS)]
-topDownRules = [RuleOp removeReplicateMapping,
+topDownRules = [RuleOp hoistCertificates,
+                RuleOp removeReplicateMapping,
                 RuleOp removeReplicateWrite,
                 RuleOp removeUnusedSOACInput,
                 RuleOp simplifyClosedFormReduce,
                 RuleOp simplifyKnownIterationSOAC,
-                RuleOp fuseConcatScatter
+                RuleOp fuseConcatScatter,
+                RuleOp simplifyMapIota,
+                RuleOp moveTransformToInput
                ]
 
 bottomUpRules :: [BottomUpRule (Wise SOACS)]
@@ -167,11 +175,34 @@
                  RuleOp mapOpToOp
                 ]
 
+-- Any certificates attached to a trivial Stm in the body might as
+-- well be applied to the SOAC itself.
+hoistCertificates :: TopDownRuleOp (Wise SOACS)
+hoistCertificates vtable pat aux soac
+  | (soac', hoisted) <- runState (mapSOACM mapper soac) mempty,
+    hoisted /= mempty =
+      certifying (hoisted <> stmAuxCerts aux) $ letBind_ pat $ Op soac'
+  where mapper = identitySOACMapper { mapOnSOACLambda = onLambda }
+        onLambda lam = do
+          stms' <- mapM onStm $ bodyStms $ lambdaBody lam
+          return lam { lambdaBody =
+                       mkBody stms' $ bodyResult $ lambdaBody lam }
+        onStm (Let se_pat se_aux (BasicOp (SubExp se))) = do
+          let (invariant, variant) =
+                partition (`ST.elem` vtable) $
+                unCertificates $ stmAuxCerts se_aux
+              se_aux' = se_aux { stmAuxCerts = Certificates variant }
+          modify (Certificates invariant<>)
+          return $ Let se_pat se_aux' $ BasicOp $ SubExp se
+        onStm stm = return stm
+hoistCertificates _ _ _ _ =
+  cannotSimplify
+
 liftIdentityMapping :: BottomUpRuleOp (Wise SOACS)
 liftIdentityMapping (_, usages) pat _ (Screma w form arrs)
   | Just fun <- isMapSOAC form = do
   let inputMap = M.fromList $ zip (map paramName $ lambdaParams fun) arrs
-      free = freeInBody $ lambdaBody fun
+      free = freeIn $ lambdaBody fun
       rettype = lambdaReturnType fun
       ses = bodyResult $ lambdaBody fun
 
@@ -294,7 +325,7 @@
           map_lam' = map_lam { lambdaParams = used_params }
       letBind_ pat $ Op $ Screma w (ScremaForm scan reduce map_lam') used_arrs
   where params_and_arrs = zip (lambdaParams map_lam) arrs
-        used_in_body = freeInBody $ lambdaBody map_lam
+        used_in_body = freeIn $ lambdaBody map_lam
         usedInput (param, _) = paramName param `S.member` used_in_body
 removeUnusedSOACInput _ _ _ _ = cannotSimplify
 
@@ -400,14 +431,15 @@
 -- actually used for computing one of the live ones.
 removeDeadReduction :: BottomUpRuleOp (Wise SOACS)
 removeDeadReduction (_, used) pat (StmAux cs _) (Screma w form arrs)
-  | Just (comm, redlam, nes, maplam) <- isRedomapSOAC form,
+  | Just ([Reduce comm redlam nes], maplam) <- isRedomapSOAC form,
     not $ all (`UT.used` used) $ patternNames pat, -- Quick/cheap check
 
+    let (red_pes, map_pes) = splitAt (length nes) $ patternElements pat,
     let redlam_deps = dataDependencies $ lambdaBody redlam,
     let redlam_res = bodyResult $ lambdaBody redlam,
     let redlam_params = lambdaParams redlam,
     let used_after = map snd $ filter ((`UT.used` used) . patElemName . fst) $
-                     zip (patternElements pat) redlam_params,
+                     zip red_pes redlam_params,
     let necessary = findNecessaryForReturned (`elem` used_after)
                     (zip redlam_params $ redlam_res <> redlam_res) redlam_deps,
     let alive_mask = map ((`S.member` necessary) . paramName) redlam_params,
@@ -416,15 +448,15 @@
 
   let fixDeadToNeutral lives ne = if lives then Nothing else Just ne
       dead_fix = zipWith fixDeadToNeutral alive_mask nes
-      (used_pes, _, used_nes) =
+      (used_red_pes, _, used_nes) =
         unzip3 $ filter (\(_,x,_) -> paramName x `S.member` necessary) $
-        zip3 (patternElements pat) redlam_params nes
+        zip3 red_pes redlam_params nes
 
-  let maplam' = removeLambdaResults alive_mask maplam
-  redlam' <- removeLambdaResults alive_mask <$> fixLambdaParams redlam (dead_fix++dead_fix)
+  let maplam' = removeLambdaResults (take (length nes) alive_mask) maplam
+  redlam' <- removeLambdaResults (take (length nes) alive_mask) <$> fixLambdaParams redlam (dead_fix++dead_fix)
 
-  certifying cs $ letBind_ (Pattern [] used_pes) $
-    Op $ Screma w (redomapSOAC comm redlam' used_nes maplam') arrs
+  certifying cs $ letBind_ (Pattern [] $ used_red_pes ++ map_pes) $
+    Op $ Screma w (redomapSOAC [Reduce comm redlam' used_nes] maplam') arrs
 
 removeDeadReduction _ _ _ _ = cannotSimplify
 
@@ -477,23 +509,29 @@
             y_ws<- mapM sizeOf ys
             guard $ all (x_w==) y_ws
             return (x_w, x:ys, cs)
+          Just (BasicOp (Reshape reshape arr), cs) -> do
+            guard $ isJust $ shapeCoercion reshape
+            (a, b, cs') <- isConcat arr
+            return (a, b, cs <> cs')
           _ -> Nothing
 
 fuseConcatScatter _ _ _ _ = cannotSimplify
 
 simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS)
+simplifyClosedFormReduce _ pat _ (Screma (Constant w) form _)
+  | Just nes <- concatMap redNeutral . fst <$> isRedomapSOAC form,
+    zeroIsh w =
+      forM_ (zip (patternNames pat) nes) $ \(v, ne) ->
+      letBindNames_ [v] $ BasicOp $ SubExp ne
 simplifyClosedFormReduce vtable pat _ (Screma _ form arrs)
-  | Just (_, red_fun, nes) <- isReduceSOAC form =
+  | Just [Reduce _ red_fun nes] <- isReduceSOAC form =
       foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs
 simplifyClosedFormReduce _ _ _ _ = cannotSimplify
 
 -- For now we just remove singleton SOACs.
-simplifyKnownIterationSOAC :: (BinderOps lore, Op lore ~ SOAC lore) =>
-                              TopDownRuleOp lore
+simplifyKnownIterationSOAC :: TopDownRuleOp (Wise SOACS)
 simplifyKnownIterationSOAC _ pat _ (Screma (Constant k)
-                                    (ScremaForm (scan_lam, scan_nes)
-                                                  (_, red_lam, red_nes)
-                                                  map_lam)
+                                    (ScremaForm (scan_lam, scan_nes) reds map_lam)
                                     arrs)
   | oneIsh k = do
       zipWithM_ bindMapParam (lambdaParams map_lam) arrs
@@ -506,7 +544,8 @@
       zipWithM_ bindResult red_pes red_res
       zipWithM_ bindArrayResult map_pes map_res
 
-        where (scan_pes, red_pes, map_pes) = splitAt3 (length scan_nes) (length red_nes) $
+        where (Reduce _ red_lam red_nes) = singleReduce reds
+              (scan_pes, red_pes, map_pes) = splitAt3 (length scan_nes) (length red_nes) $
                                              patternElements pat
               bindMapParam p a = do
                 a_t <- lookupType a
@@ -518,3 +557,178 @@
               bindResult pe se =
                 letBindNames_ [patElemName pe] $ BasicOp $ SubExp se
 simplifyKnownIterationSOAC _ _ _ _ = cannotSimplify
+
+data ArrayOp = ArrayIndexing Certificates VName (Slice SubExp)
+             | ArrayRearrange Certificates VName [Int]
+             | ArrayVar Certificates VName -- ^ Never constructed.
+  deriving (Eq, Ord, Show)
+
+arrayOpArr :: ArrayOp -> VName
+arrayOpArr (ArrayIndexing _ arr _) = arr
+arrayOpArr (ArrayRearrange _ arr _) = arr
+arrayOpArr (ArrayVar _ arr) = arr
+
+arrayOpCerts :: ArrayOp -> Certificates
+arrayOpCerts (ArrayIndexing cs _ _) = cs
+arrayOpCerts (ArrayRearrange cs _ _) = cs
+arrayOpCerts (ArrayVar cs _) = cs
+
+isArrayOp :: Certificates -> AST.Exp (Wise SOACS) -> Maybe ArrayOp
+isArrayOp cs (BasicOp (Index arr slice)) =
+  Just $ ArrayIndexing cs arr slice
+isArrayOp cs (BasicOp (Rearrange perm arr)) =
+  Just $ ArrayRearrange cs arr perm
+isArrayOp _ _ =
+  Nothing
+
+fromArrayOp :: ArrayOp -> (Certificates, AST.Exp (Wise SOACS))
+fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice)
+fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr)
+fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr)
+
+arrayOps :: AST.Body (Wise SOACS) -> S.Set ArrayOp
+arrayOps = mconcat . map onStm . stmsToList . bodyStms
+  where onStm (Let _ aux e) =
+          case isArrayOp (stmAuxCerts aux) e of
+            Just op -> S.singleton op
+            Nothing -> execWriter $ walkExpM walker e
+        onOp = execWriter . mapSOACM identitySOACMapper { mapOnSOACLambda = onLambda }
+        onLambda lam = do tell $ arrayOps $ lambdaBody lam
+                          return lam
+        walker = identityWalker { walkOnBody = tell . arrayOps
+                                , walkOnOp = tell . onOp }
+
+replaceArrayOps :: M.Map ArrayOp ArrayOp
+                -> AST.Body (Wise SOACS) -> AST.Body (Wise SOACS)
+replaceArrayOps substs (Body _ stms res) =
+  mkBody (fmap onStm stms) res
+  where onStm (Let pat aux e) =
+          let (cs', e') = onExp (stmAuxCerts aux) e
+          in certify cs' $ mkLet (patternContextIdents pat) (patternValueIdents pat) e'
+        onExp cs e
+          | Just op <- isArrayOp cs e,
+            Just op' <- M.lookup op substs =
+              fromArrayOp op'
+        onExp cs e = (cs, mapExp mapper e)
+        mapper = identityMapper { mapOnBody = const $ return . replaceArrayOps substs
+                                , mapOnOp = return . onOp }
+        onOp = runIdentity . mapSOACM identitySOACMapper { mapOnSOACLambda = return . onLambda }
+        onLambda lam = lam { lambdaBody = replaceArrayOps substs $ lambdaBody lam }
+
+-- Turn
+--
+--    map (\i -> ... xs[i] ...) (iota n)
+--
+-- into
+--
+--    map (\i x -> ... x ...) (iota n) xs
+--
+-- This is not because we want to encourage the map-iota pattern, but
+-- it may be present in generated code.  This is an unfortunately
+-- expensive simplification rule, since it requires multiple passes
+-- over the entire lambda body.  It only handles the very simplest
+-- case - if you find yourself planning to extend it to handle more
+-- complex situations (rotate or whatnot), consider turning it into a
+-- separate compiler pass instead.
+simplifyMapIota :: TopDownRuleOp (Wise SOACS)
+simplifyMapIota vtable pat _ (Screma w (ScremaForm scan reduce map_lam) arrs)
+  | Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs),
+    indexings <- filter (indexesWith (paramName p)) $ S.toList $
+                 arrayOps $ lambdaBody map_lam,
+    not $ null indexings = do
+      -- For each indexing with iota, add the corresponding array to
+      -- the Screma, and construct a new lambda parameter.
+      (more_arrs, more_params, replacements) <-
+        unzip3 . catMaybes <$> mapM mapOverArr indexings
+      let substs = M.fromList $ zip indexings replacements
+          map_lam' = map_lam { lambdaParams = lambdaParams map_lam <> more_params
+                             , lambdaBody = replaceArrayOps substs $
+                                            lambdaBody map_lam
+                             }
+      letBind_ pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs)
+  where isIota (_, arr) = case ST.lookupBasicOp arr vtable of
+                            Just (Iota _ (Constant o) (Constant s) _, _) ->
+                              zeroIsh o && oneIsh s
+                            _ -> False
+
+        indexesWith v (ArrayIndexing cs arr (DimFix (Var i) : _))
+          | arr `ST.elem` vtable,
+            all (`ST.elem` vtable) $ unCertificates cs =
+              i == v
+        indexesWith _ _ = False
+
+        mapOverArr (ArrayIndexing cs arr slice) = do
+          arr_elem <- newVName $ baseString arr ++ "_elem"
+          arr_t <- lookupType arr
+          arr' <- if arraySize 0 arr_t == w
+                  then return arr
+                  else certifying cs $ letExp (baseString arr ++ "_prefix") $
+                       BasicOp $ Index arr $
+                       fullSlice arr_t [DimSlice (intConst Int32 0) w (intConst Int32 1)]
+          return $ Just (arr',
+                         Param arr_elem (rowType arr_t),
+                         ArrayIndexing cs arr_elem (drop 1 slice))
+
+        mapOverArr _ = return Nothing
+
+simplifyMapIota  _ _ _ _ = cannotSimplify
+
+-- If a Screma's map function contains a transformation
+-- (e.g. transpose) on a parameter, create a new parameter
+-- corresponding to that transformation performed on the rows of the
+-- full array.
+moveTransformToInput :: TopDownRuleOp (Wise SOACS)
+moveTransformToInput vtable pat _ (Screma w (ScremaForm scan reduce map_lam) arrs)
+  | ops <- filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam,
+    not $ null ops = do
+      (more_arrs, more_params, replacements) <-
+        unzip3 . catMaybes <$> mapM mapOverArr ops
+
+      when (null more_arrs) cannotSimplify
+
+      let substs = M.fromList $ zip ops replacements
+          map_lam' = map_lam { lambdaParams = lambdaParams map_lam <> more_params
+                             , lambdaBody = replaceArrayOps substs $
+                                            lambdaBody map_lam
+                             }
+
+      letBind_ pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs)
+
+  where map_param_names = map paramName (lambdaParams map_lam)
+
+        -- It's not just about whether the array is a parameter;
+        -- everything else must be map-invariant.
+        arrayIsMapParam (ArrayIndexing cs arr slice) =
+          arr `elem` map_param_names &&
+          all (`ST.elem` vtable) (S.toList $ freeIn cs <> freeIn slice) &&
+          not (null slice) && not (null $ sliceDims slice)
+        arrayIsMapParam (ArrayRearrange cs arr perm) =
+          arr `elem` map_param_names &&
+          all (`ST.elem` vtable) (S.toList $ freeIn cs) &&
+          not (null perm)
+        arrayIsMapParam ArrayVar{} =
+          False
+
+        mapOverArr op
+         | Just (_, arr) <- find ((==arrayOpArr op) . fst) (zip map_param_names arrs) = do
+             arr_t <- lookupType arr
+             let whole_dim = DimSlice (intConst Int32 0) (arraySize 0 arr_t) (intConst Int32 1)
+             arr_transformed <- certifying (arrayOpCerts op) $
+                                letExp (baseString arr ++ "_transformed") $
+                                case op of
+                                  ArrayIndexing _ _ slice ->
+                                    BasicOp $ Index arr $ whole_dim : slice
+                                  ArrayRearrange _ _ perm ->
+                                    BasicOp $ Rearrange (0 : map (+1) perm) arr
+                                  ArrayVar{} ->
+                                    BasicOp $ SubExp $ Var arr
+             arr_transformed_t <- lookupType arr_transformed
+             arr_transformed_row <- newVName $ baseString arr ++ "_transformed_row"
+             return $ Just (arr_transformed,
+                            Param arr_transformed_row (rowType arr_transformed_t),
+                            ArrayVar mempty arr_transformed_row)
+
+        mapOverArr _ = return Nothing
+
+moveTransformToInput _ _ _ _ =
+  cannotSimplify
diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs
--- a/src/Futhark/Tools.hs
+++ b/src/Futhark/Tools.hs
@@ -62,7 +62,7 @@
   let map_bnd = mkLet [] map_pat $ Op $ Screma w (mapSOAC map_lam) arrs
       (nes, red_arrs) = unzip red_args
   red_bnd <- Let red_pat (defAux ()) . Op <$>
-             (Screma w <$> reduceSOAC comm redlam nes <*> pure red_arrs)
+             (Screma w <$> reduceSOAC [Reduce comm redlam nes] <*> pure red_arrs)
   return (map_bnd, red_bnd)
 redomapToMapAndReduce _ _ =
   error "redomapToMapAndReduce does not handle a non-empty 'patternContextElements'"
@@ -112,18 +112,17 @@
                     Bindable (Lore m)) =>
                    Pattern (Lore m) -> SubExp -> ScremaForm (Lore m) -> [VName]
                 -> m ()
-dissectScrema pat w (ScremaForm (scan_lam, scan_nes)
-                                    (comm, red_lam, red_nes)
-                                    map_lam) arrs = do
-  let (scan_res, red_res, map_res) = splitAt3 (length scan_nes) (length red_nes) $
-                                     patternNames pat
+dissectScrema pat w (ScremaForm (scan_lam, scan_nes) reds map_lam) arrs = do
+  let num_reds = redResults reds
+      (scan_res, red_res, map_res) =
+        splitAt3 (length scan_nes) num_reds $ patternNames pat
   -- First we perform the Map, then we perform the Reduce, and finally
   -- the Scan.
   to_scan <- replicateM (length scan_nes) $ newVName "to_scan"
-  to_red <- replicateM (length red_nes) $ newVName "to_red"
+  to_red <- replicateM num_reds $ newVName "to_red"
   letBindNames_ (to_scan <> to_red <> map_res) $ Op $ Screma w (mapSOAC map_lam) arrs
 
-  reduce <- reduceSOAC comm red_lam red_nes
+  reduce <- reduceSOAC reds
   letBindNames_ red_res $ Op $ Screma w reduce to_red
 
   scan <- scanSOAC scan_lam scan_nes
diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs
--- a/src/Futhark/Transform/FirstOrderTransform.hs
+++ b/src/Futhark/Transform/FirstOrderTransform.hs
@@ -91,10 +91,10 @@
 transformSOAC pat CmpThreshold{} =
   letBind_ pat $ BasicOp $ SubExp $ constant False -- close enough
 
-transformSOAC pat (Screma w form@(ScremaForm (scan_lam, scan_nes)
-                                                 (_, red_lam, red_nes)
-                                                 map_lam) arrs) = do
-  let (scan_arr_ts, _red_ts, map_arr_ts) =
+transformSOAC pat (Screma w form@(ScremaForm (scan_lam, scan_nes) reds map_lam) arrs) = do
+  -- Start by combining all the reduction parts into a single operator
+  let (Reduce _ red_lam red_nes) = singleReduce reds
+      (scan_arr_ts, _red_ts, map_arr_ts) =
         splitAt3 (length scan_nes) (length red_nes) $ scremaType w form
   scan_arrs <- resultArray scan_arr_ts
   map_arrs <- resultArray map_arr_ts
diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs
--- a/src/Futhark/Transform/Rename.hs
+++ b/src/Futhark/Transform/Rename.hs
@@ -268,7 +268,6 @@
                       mapOnBody = const rename
                     , mapOnSubExp = rename
                     , mapOnVName = rename
-                    , mapOnCertificates = rename
                     , mapOnRetType = rename
                     , mapOnBranchType = rename
                     , mapOnFParam = rename
@@ -282,7 +281,7 @@
     size' <- rename size
     return $ Array et size' u
   rename (Prim et) = return $ Prim et
-  rename (Mem e space) = Mem <$> rename e <*> pure space
+  rename (Mem space) = pure $ Mem space
 
 instance Renameable lore => Rename (Lambda lore) where
   rename (Lambda params body ret) =
diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs
--- a/src/Futhark/Transform/Substitute.hs
+++ b/src/Futhark/Transform/Substitute.hs
@@ -113,7 +113,6 @@
                    mapOnVName = return . substituteNames substs
                  , mapOnSubExp = return . substituteNames substs
                  , mapOnBody = const $ return . substituteNames substs
-                 , mapOnCertificates = return . substituteNames substs
                  , mapOnRetType = return . substituteNames substs
                  , mapOnBranchType = return . substituteNames substs
                  , mapOnFParam = return . substituteNames substs
@@ -142,8 +141,8 @@
   substituteNames _ (Prim et) = Prim et
   substituteNames substs (Array et sz u) =
     Array et (substituteNames substs sz) u
-  substituteNames substs (Mem sz space) =
-    Mem (substituteNames substs sz) space
+  substituteNames _ (Mem space) =
+    Mem space
 
 instance Substitutable lore => Substitute (Lambda lore) where
   substituteNames substs (Lambda params body rettype) =
diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs
--- a/src/Futhark/Util.hs
+++ b/src/Futhark/Util.hs
@@ -26,6 +26,7 @@
         directoryContents,
         roundFloat,
         roundDouble,
+        lgamma, lgammaf, tgamma, tgammaf,
         fromPOSIX,
         toPOSIX,
         trim,
@@ -174,6 +175,27 @@
 -- | Round a double-precision floating point number correctly.
 roundDouble :: Double -> Double
 roundDouble = c_nearbyint
+
+foreign import ccall "lgamma" c_lgamma :: Double -> Double
+foreign import ccall "lgammaf" c_lgammaf :: Float -> Float
+foreign import ccall "tgamma" c_tgamma :: Double -> Double
+foreign import ccall "tgammaf" c_tgammaf :: Float -> Float
+
+-- | The system-level @lgamma()@ function.
+lgamma :: Double -> Double
+lgamma = c_lgamma
+
+-- | The system-level @lgammaf()@ function.
+lgammaf :: Float -> Float
+lgammaf = c_lgammaf
+
+-- | The system-level @tgamma()@ function.
+tgamma :: Double -> Double
+tgamma = c_tgamma
+
+-- | The system-level @tgammaf()@ function.
+tgammaf :: Float -> Float
+tgammaf = c_tgammaf
 
 -- | Turn a POSIX filepath into a filepath for the native system.
 toPOSIX :: Native.FilePath -> Posix.FilePath
diff --git a/src/Language/Futhark/Attributes.hs b/src/Language/Futhark/Attributes.hs
--- a/src/Language/Futhark/Attributes.hs
+++ b/src/Language/Futhark/Attributes.hs
@@ -31,7 +31,6 @@
   , patternStructType
   , patternPatternType
   , patternParam
-  , patternNoShapeAnnotations
   , patternOrderZero
   , patternDimNames
 
@@ -463,7 +462,7 @@
 typeOf (Ascript _ _ (Info t) _) = t
 typeOf (Apply _ _ _ (Info t) _) = t
 typeOf (Negate e _) = typeOf e
-typeOf (LetPat _ _ _ _ (Info t) _) = t
+typeOf (LetPat _ _ _ (Info t) _) = t
 typeOf (LetFun _ _ body _) = typeOf body
 typeOf (LetWith _ _ _ _ _ (Info t) _) = t
 typeOf (Index _ _ (Info t) _) = t
@@ -471,12 +470,11 @@
 typeOf (RecordUpdate _ _ _ (Info t) _) = t
 typeOf (Unsafe e _) = typeOf e
 typeOf (Assert _ e _ _) = typeOf e
-typeOf (DoLoop _ pat _ _ _ _) = patternType pat
-typeOf (Lambda tparams params _ _ (Info (als, t)) _) =
+typeOf (DoLoop pat _ _ _ _) = patternType pat
+typeOf (Lambda params _ _ (Info (als, t)) _) =
   unscopeType bound_here $
   foldr (uncurry (Arrow ()) . patternParam) t params `setAliases` als
-  where bound_here = S.fromList (map typeParamName tparams) <>
-                     S.map identName (mconcat $ map patternIdents params)
+  where bound_here = S.map identName (mconcat $ map patternIdents params)
 typeOf (OpSection _ (Info t) _) =
   t
 typeOf (OpSectionLeft _ _ _ (_, Info pt2) (Info ret) _)  =
@@ -603,28 +601,11 @@
   patternParam p
 patternParam (PatternAscription (Id v _ _) td _) =
   (Just v, unInfo $ expandedType td)
+patternParam (Id v (Info t) _) =
+  (Just v, toStruct t)
 patternParam p =
   (Nothing, patternStructType p)
 
--- | Remove all shape annotations from a pattern, leaving them unnamed
--- instead.
-patternNoShapeAnnotations :: PatternBase Info VName -> PatternBase Info VName
-patternNoShapeAnnotations (PatternAscription p (TypeDecl te (Info t)) loc) =
-  PatternAscription (patternNoShapeAnnotations p)
-  (TypeDecl te $ Info $ anyDimShapeAnnotations t) loc
-patternNoShapeAnnotations (PatternParens p loc) =
-  PatternParens (patternNoShapeAnnotations p) loc
-patternNoShapeAnnotations (Id v (Info t) loc) =
-  Id v (Info $ anyDimShapeAnnotations t) loc
-patternNoShapeAnnotations (TuplePattern ps loc) =
-  TuplePattern (map patternNoShapeAnnotations ps) loc
-patternNoShapeAnnotations (RecordPattern ps loc) =
-  RecordPattern (map (fmap patternNoShapeAnnotations) ps) loc
-patternNoShapeAnnotations (Wildcard (Info t) loc) =
-  Wildcard (Info (anyDimShapeAnnotations t)) loc
-patternNoShapeAnnotations (PatternLit e (Info t) loc) =
-  PatternLit e (Info (anyDimShapeAnnotations t)) loc
-
 -- | Names of primitive types to types.  This is only valid if no
 -- shadowing is going on, but useful for tools.
 namesToPrimTypes :: M.Map Name PrimType
@@ -735,17 +716,16 @@
                tupleRecord [uarr_a, Array () Unique (ArrayPrimElem (Signed Int32)) (rank 1)]),
 
               ("stream_map",
-               IntrinsicPolyFun [tp_a, tp_b] [arr_a `arr` arr_b, arr_a] uarr_b),
+               IntrinsicPolyFun [tp_a, tp_b] [Prim (Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] uarr_b),
 
               ("stream_map_per",
-               IntrinsicPolyFun [tp_a, tp_b] [arr_a `arr` arr_b, arr_a] uarr_b),
+               IntrinsicPolyFun [tp_a, tp_b] [Prim (Signed Int32) `arr` (arr_a `arr` arr_b), arr_a] uarr_b),
 
               ("stream_red",
-               IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), arr_a `arr` t_b, arr_a] t_b),
+               IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), Prim (Signed Int32) `arr` (arr_a `arr` t_b), arr_a] t_b),
 
               ("stream_red_per",
-               IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), arr_a `arr` t_b, arr_a] t_b),
-
+               IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), Prim (Signed Int32) `arr` (arr_a `arr` t_b), arr_a] t_b),
 
               ("trace", IntrinsicPolyFun [tp_a] [t_a] t_a),
               ("break", IntrinsicPolyFun [tp_a] [t_a] t_a)]
diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs
--- a/src/Language/Futhark/Interpreter.hs
+++ b/src/Language/Futhark/Interpreter.hs
@@ -539,6 +539,40 @@
         matchPtoA _ _ = mempty
 evalType _ (Enum cs) = Enum cs
 
+evalFunction :: Env -> [TypeParam] -> [Pattern] -> Exp
+             -> (Aliasing, StructType) -> SrcLoc -> EvalM Value
+
+-- We treat zero-parameter lambdas as simply an expression to
+-- evaluate immediately.  Note that this is *not* the same as a lambda
+-- that takes an empty tuple '()' as argument!  Zero-parameter lambdas
+-- can never occur in a well-formed Futhark program, but they are
+-- convenient in the interpreter.
+evalFunction env tparams [] body (_, t) loc = do
+  -- All remaining size parameters that have not yet been assigned a
+  -- value (because they were inner dimensions of empty arrays) are
+  -- now assigned a zero.
+  let unbound_dims = bindToZero $ map typeParamName $ filter isDimParam tparams
+  v <- eval (env <> unbound_dims) body
+  case (t, v) of
+    (Arrow _ _ _ rt, ValueFun f) ->
+      return $ ValueFun $ \arg -> do r <- f arg
+                                     match (evalType env rt) r
+    _ -> match t v
+  where match vt v =
+          case matchValueToType env vt v of
+            Right _ -> return v
+            Left err ->
+              bad loc env $ "Value `" <> pretty v <>
+              "` cannot match type `" <> pretty vt <> "`: " ++ err
+
+        isDimParam TypeParamDim{} = True
+        isDimParam _ = False
+
+evalFunction env tparams (p:ps) body (als, ret) loc =
+  return $ ValueFun $ \v -> do
+    env' <- matchPattern env p v
+    evalFunction env' tparams ps body (als, ret) loc
+
 eval :: Env -> Exp -> EvalM Value
 
 eval _ (Literal v _) = return $ ValuePrim v
@@ -592,13 +626,13 @@
     Left err -> bad loc env $ "Value `" <> pretty v <> "` cannot match shape of type `" <>
                 pretty (declaredType td) <> "` (`" <> pretty t <> "`): " ++ err
 
-eval env (LetPat _ p e body _ _) = do
+eval env (LetPat p e body _ _) = do
   v <- eval env e
   env' <- matchPattern env p v
   eval env' body
 
 eval env (LetFun f (tparams, pats, _, Info ret, fbody) body loc) = do
-  v <- eval env $ Lambda tparams pats fbody Nothing (Info (mempty, ret)) loc
+  v <- evalFunction env tparams pats fbody (mempty, ret) loc
   let ftype = T.BoundV [] $ foldr (uncurry (Arrow ()) . patternParam) ret pats
   eval (valEnv (M.singleton f (Just ftype, v)) <> env) body
 
@@ -690,31 +724,8 @@
 -- that takes an empty tuple '()' as argument!  Zero-parameter lambdas
 -- can never occur in a well-formed Futhark program, but they are
 -- convenient in the interpreter.
-eval env (Lambda tparams [] body _ (Info (_, t)) loc) = do
-  -- All remaining size parameters that have not yet been assigned a
-  -- value (because they were inner dimensions of empty arrays) are
-  -- now assigned a zero.
-  let unbound_dims = bindToZero $ map typeParamName $ filter isDimParam tparams
-  v <- eval (env <> unbound_dims) body
-  case (t, v) of
-    (Arrow _ _ _ rt, ValueFun f) ->
-      return $ ValueFun $ \arg -> do r <- f arg
-                                     match (evalType env rt) r
-    _ -> match t v
-  where match vt v =
-          case matchValueToType env vt v of
-            Right _ -> return v
-            Left err ->
-              bad loc env $ "Value `" <> pretty v <>
-              "` cannot match type `" <> pretty vt <> "`: " ++ err
-
-        isDimParam TypeParamDim{} = True
-        isDimParam _ = False
-
-eval env (Lambda tparams (p:ps) body mrd (Info (als, ret)) loc) =
-  return $ ValueFun $ \v -> do
-    env' <- matchPattern env p v
-    eval env' $ Lambda tparams ps body mrd (Info (als, ret)) loc
+eval env (Lambda ps body _ (Info (als, ret)) loc) =
+  evalFunction env [] ps body (als, ret) loc
 
 eval env (OpSection qv _  _) = evalTermVar env qv
 
@@ -735,7 +746,7 @@
           | Just v' <- M.lookup f fs = return v'
         walk _ _ = fail "Value does not have expected field."
 
-eval env (DoLoop _ pat init_e form body _) = do
+eval env (DoLoop pat init_e form body _) = do
   init_v <- eval env init_e
   case form of For iv bound -> do
                  bound' <- asSigned <$> eval env bound
@@ -869,8 +880,8 @@
 
 evalDec env (ValDec (ValBind _ v _ (Info t) tps ps def _ loc)) = do
   let t' = evalType env t
-      ftype = T.BoundV tps $ foldr (uncurry (Arrow ()) . patternParam) t' ps
-  val <- eval env $ Lambda tps ps def Nothing (Info (mempty, t')) loc
+      ftype = T.BoundV [] $ foldr (uncurry (Arrow ()) . patternParam) t' ps
+  val <- evalFunction env tps ps def (mempty, t') loc
   return $ valEnv (M.singleton v (Just ftype, val)) <> env
 
 evalDec env (OpenDec me _) = do
@@ -1093,10 +1104,10 @@
       toArray . reverse . fst =<< foldM next ([], ne) (fromArray xs)
 
     def s | "stream_map" `isPrefixOf` s =
-              Just $ fun2t $ apply noLoc mempty
+              Just $ fun2t stream
 
     def s | "stream_red" `isPrefixOf` s =
-              Just $ fun3t $ \_ f xs -> apply noLoc mempty f xs
+              Just $ fun3t $ \_ f arg -> stream f arg
 
     def "scatter" = Just $ fun3t $ \arr is vs ->
       case arr of
@@ -1186,6 +1197,12 @@
     tdef s = do
       t <- nameFromString s `M.lookup` namesToPrimTypes
       return $ T.TypeAbbr Unlifted [] $ Prim t
+
+    stream f arg@(ValueArray xs) =
+      let n = ValuePrim $ SignedValue $ Int32Value $ arrayLength xs
+      in apply2 noLoc mempty f n arg
+    stream _ arg = error $ "Cannot stream: " ++ pretty arg
+
 
 interpretExp :: Ctx -> Exp -> F ExtOp Value
 interpretExp ctx e = runEvalM (ctxImports ctx) $ eval (ctxEnv ctx) e
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
@@ -325,9 +325,6 @@
             : TypeParam TypeParams { $1 : $2 }
             |                      { [] }
 
-TypeParams1 :: { (TypeParamBase Name, [TypeParamBase Name]) }
-            : TypeParam TypeParams { ($1, $2) }
-
 UnOp :: { (QualName Name, SrcLoc) }
       : qunop { let L loc (QUALUNOP qs v) = $1 in (QualName qs v, loc) }
       | unop  { let L loc (UNOP v) = $1 in (qualName v, loc) }
@@ -519,11 +516,11 @@
      : if Exp then Exp else Exp %prec ifprec
                       { If $2 $4 $6 NoInfo (srcspan $1 $>) }
 
-     | loop TypeParams Pattern LoopForm do Exp %prec ifprec
-         {% fmap (\t -> DoLoop $2 $3 t $4 $6 (srcspan $1 $>)) (patternExp $3) }
+     | loop Pattern LoopForm do Exp %prec ifprec
+         {% fmap (\t -> DoLoop $2 t $3 $5 (srcspan $1 $>)) (patternExp $2) }
 
-     | loop TypeParams Pattern '=' Exp LoopForm do Exp %prec ifprec
-         { DoLoop $2 $3 $5 $6 $8 (srcspan $1 $>) }
+     | loop Pattern '=' Exp LoopForm do Exp %prec ifprec
+         { DoLoop $2 $4 $5 $7 (srcspan $1 $>) }
 
      | LetExp %prec letprec { $1 }
 
@@ -580,8 +577,8 @@
      | Exp2 with FieldAccesses_ '=' Exp2
        { RecordUpdate $1 (map fst $3) $5 NoInfo (srcspan $1 $>) }
 
-     | '\\' TypeParams FunParams1 maybeAscription(TypeExpTerm) '->' Exp
-       { Lambda $2 (fst $3 : snd $3) $6 $4 NoInfo (srcspan $1 $>) }
+     | '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp
+       { Lambda (fst $2 : snd $2) $5 $3 NoInfo (srcspan $1 $>) }
 
      | Apply { $1 }
 
@@ -696,9 +693,7 @@
 
 LetExp :: { UncheckedExp }
      : let Pattern '=' Exp LetBody
-                      { LetPat [] $2 $4 $5 NoInfo (srcspan $1 $>) }
-     | let TypeParams1 Pattern '=' Exp LetBody
-                      { LetPat (fst $2 : snd $2) $3 $5 $6 NoInfo (srcspan $1 $>) }
+                      { LetPat $2 $4 $5 NoInfo (srcspan $1 $>) }
 
      | let id TypeParams FunParams1 maybeAscription(TypeExpDecl) '=' Exp LetBody
        { let L _ (ID name) = $2
diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs
--- a/src/Language/Futhark/Pretty.hs
+++ b/src/Language/Futhark/Pretty.hs
@@ -196,6 +196,10 @@
   ppr (DimSlice i Nothing Nothing) =
     maybe mempty ppr i <> text ":"
 
+letBody :: (Eq vn, IsName vn, Annot f) => ExpBase f vn -> Doc
+letBody body@LetPat{} = ppr body
+letBody body          = text "in" <+> align (ppr body)
+
 instance (Eq vn, IsName vn, Annot f) => Pretty (ExpBase f vn) where
   ppr = pprPrec (-1)
   pprPrec _ (Var name _ _) = ppr name
@@ -230,14 +234,13 @@
   pprPrec p (Apply f arg _ _ _) =
     parensIf (p >= 10) $ ppr f <+> pprPrec 10 arg
   pprPrec _ (Negate e _) = text "-" <> ppr e
-  pprPrec p (LetPat tparams pat e body _ _) =
+  pprPrec p (LetPat pat e body _ _) =
     parensIf (p /= -1) $ align $
-    text "let" <+> align (spread $ map ppr tparams ++ [ppr pat]) <+>
+    text "let" <+> align (ppr pat) <+>
     (if linebreak
      then equals </> indent 2 (ppr e)
      else equals <+> align (ppr e)) </>
-    (case body of LetPat{} -> ppr body
-                  _        -> text "in" <+> align (ppr body))
+    letBody body
     where linebreak = case e of
                         DoLoop{}    -> True
                         LetPat{}    -> True
@@ -247,21 +250,21 @@
                         _           -> hasArrayLit e
   pprPrec _ (LetFun fname (tparams, params, retdecl, rettype, e) body _) =
     text "let" <+> pprName fname <+> spread (map ppr tparams ++ map ppr params) <>
-    retdecl' <+> equals </> indent 2 (ppr e) <+> text "in" </>
-    ppr body
+    retdecl' <+> equals </> indent 2 (ppr e) </>
+    letBody body
     where retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of
                        Just rettype' -> text ":" <+> rettype'
                        Nothing       -> mempty
   pprPrec _ (LetWith dest src idxs ve body _ _)
     | dest == src =
       text "let" <+> ppr dest <> list (map ppr idxs) <+>
-      equals <+> align (ppr ve) <+>
-      text "in" </> ppr body
+      equals <+> align (ppr ve) </>
+      letBody body
     | otherwise =
       text "let" <+> ppr dest <+> equals <+> ppr src <+>
       text "with" <+> brackets (commasep (map ppr idxs)) <+>
-      text "<-" <+> align (ppr ve) <+>
-      text "in" </> ppr body
+      text "<-" <+> align (ppr ve) </>
+      letBody body
   pprPrec _ (Update src idxs ve _) =
     ppr src <+> text "with" <+>
     brackets (commasep (map ppr idxs)) <+>
@@ -274,10 +277,9 @@
     pprPrec 9 e <> brackets (commasep (map ppr idxs))
   pprPrec _ (Unsafe e _) = text "unsafe" <+> pprPrec (-1) e
   pprPrec _ (Assert e1 e2 _ _) = text "assert" <+> pprPrec 10 e1 <+> pprPrec 10 e2
-  pprPrec p (Lambda tparams params body rettype _ _) =
+  pprPrec p (Lambda params body rettype _ _) =
     parensIf (p /= -1) $
-    text "\\" <> spread (map ppr tparams ++ map ppr params) <>
-    ppAscription rettype <+>
+    text "\\" <> spread (map ppr params) <> ppAscription rettype <+>
     text "->" </> indent 2 (ppr body)
   pprPrec _ (OpSection binop _ _) =
     parens $ ppr binop
@@ -290,8 +292,8 @@
     where p name = text "." <> ppr name
   pprPrec _ (IndexSection idxs _ _) =
     parens $ text "." <> brackets (commasep (map ppr idxs))
-  pprPrec _ (DoLoop tparams pat initexp form loopbody _) =
-    text "loop" <+> spread (map ppr tparams ++ [ppr pat]) <+>
+  pprPrec _ (DoLoop pat initexp form loopbody _) =
+    text "loop" <+> ppr pat <+>
     equals <+> ppr initexp <+> ppr form <+> text "do" </>
     indent 2 (ppr loopbody)
   pprPrec _ (VConstr0 n _ _) = text "#" <> ppr n
diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs
--- a/src/Language/Futhark/Syntax.hs
+++ b/src/Language/Futhark/Syntax.hs
@@ -614,7 +614,7 @@
             | Ascript (ExpBase f vn) (TypeDeclBase f vn) (f PatternType) SrcLoc
             -- ^ Type ascription: @e : t@.
 
-            | LetPat [TypeParamBase vn] (PatternBase f vn) (ExpBase f vn) (ExpBase f vn) (f PatternType) SrcLoc
+            | LetPat (PatternBase f vn) (ExpBase f vn) (ExpBase f vn) (f PatternType) SrcLoc
 
             | LetFun vn ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn), f StructType, ExpBase f vn)
               (ExpBase f vn) SrcLoc
@@ -626,7 +626,7 @@
             | Negate (ExpBase f vn) SrcLoc
               -- ^ Numeric negation (ugly special case; Haskell did it first).
 
-            | Lambda [TypeParamBase vn] [PatternBase f vn] (ExpBase f vn)
+            | Lambda [PatternBase f vn] (ExpBase f vn)
               (Maybe (TypeExp vn)) (f (Aliasing, StructType)) SrcLoc
 
             | OpSection (QualName vn) (f PatternType) SrcLoc
@@ -643,7 +643,6 @@
               -- ^ Array indexing as a section: @(.[i,j])@.
 
             | DoLoop
-              [TypeParamBase vn]
               (PatternBase f vn) -- Merge variable pattern
               (ExpBase f vn) -- Initial values of merge variables.
               (LoopFormBase f vn) -- Do or while loop.
@@ -703,23 +702,23 @@
   locOf (Ascript _ _ _ loc)            = locOf loc
   locOf (Negate _ pos)                 = locOf pos
   locOf (Apply _ _ _ _ pos)            = locOf pos
-  locOf (LetPat _ _ _ _ _ loc)         = locOf loc
+  locOf (LetPat _ _ _ _ loc)           = locOf loc
   locOf (LetFun _ _ _ loc)             = locOf loc
   locOf (LetWith _ _ _ _ _ _ loc)      = locOf loc
   locOf (Index _ _ _ loc)              = locOf loc
   locOf (Update _ _ _ pos)             = locOf pos
   locOf (RecordUpdate _ _ _ _ pos)     = locOf pos
-  locOf (Lambda _ _ _ _ _ loc)         = locOf loc
+  locOf (Lambda _ _ _ _ loc)           = locOf loc
   locOf (OpSection _ _ loc)            = locOf loc
   locOf (OpSectionLeft _ _ _ _ _ loc)  = locOf loc
   locOf (OpSectionRight _ _ _ _ _ loc) = locOf loc
   locOf (ProjectSection _ _ loc)       = locOf loc
   locOf (IndexSection _ _ loc)         = locOf loc
-  locOf (DoLoop _ _ _ _ _ pos)         = locOf pos
+  locOf (DoLoop _ _ _ _ pos)           = locOf pos
   locOf (Unsafe _ loc)                 = locOf loc
   locOf (Assert _ _ _ loc)             = locOf loc
   locOf (VConstr0 _ _ loc)             = locOf loc
-  locOf (Match _ _ _ loc)                = locOf loc
+  locOf (Match _ _ _ loc)              = locOf loc
 
 -- | An entry in a record literal.
 data FieldBase f vn = RecordFieldExplicit Name (ExpBase f vn) SrcLoc
diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs
--- a/src/Language/Futhark/Traversals.hs
+++ b/src/Language/Futhark/Traversals.hs
@@ -86,9 +86,8 @@
     Apply <$> mapOnExp tv f <*> mapOnExp tv arg <*>
     pure d <*> (Info <$> mapOnPatternType tv t) <*>
     pure loc
-  astMap tv (LetPat tparams pat e body t loc) =
-    LetPat <$> mapM (astMap tv) tparams <*>
-    astMap tv pat <*> mapOnExp tv e <*>
+  astMap tv (LetPat pat e body t loc) =
+    LetPat <$> astMap tv pat <*> mapOnExp tv e <*>
     mapOnExp tv body <*> traverse (mapOnPatternType tv) t <*> pure loc
   astMap tv (LetFun name (fparams, params, ret, t, e) body loc) =
     LetFun <$> mapOnName tv name <*>
@@ -119,8 +118,8 @@
     Unsafe <$> mapOnExp tv e <*> pure loc
   astMap tv (Assert e1 e2 desc loc) =
     Assert <$> mapOnExp tv e1 <*> mapOnExp tv e2 <*> pure desc <*> pure loc
-  astMap tv (Lambda tparams params body ret t loc) =
-    Lambda <$> mapM (astMap tv) tparams <*> mapM (astMap tv) params <*>
+  astMap tv (Lambda params body ret t loc) =
+    Lambda <$> mapM (astMap tv) params <*>
     astMap tv body <*> traverse (astMap tv) ret <*>
     traverse (traverse $ mapOnStructType tv) t <*> pure loc
   astMap tv (OpSection name t loc) =
@@ -143,8 +142,8 @@
   astMap tv (IndexSection idxs t loc) =
     IndexSection <$> mapM (astMap tv) idxs <*>
     traverse (mapOnPatternType tv) t <*> pure loc
-  astMap tv (DoLoop tparams mergepat mergeexp form loopbody loc) =
-    DoLoop <$> mapM (astMap tv) tparams <*> astMap tv mergepat <*>
+  astMap tv (DoLoop mergepat mergeexp form loopbody loc) =
+    DoLoop <$> astMap tv mergepat <*>
     mapOnExp tv mergeexp <*> astMap tv form <*>
     mapOnExp tv loopbody <*> pure loc
   astMap tv (VConstr0 name t loc) =
diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs
--- a/src/Language/Futhark/TypeChecker.hs
+++ b/src/Language/Futhark/TypeChecker.hs
@@ -452,17 +452,21 @@
 
 checkTypeBind :: TypeBindBase NoInfo Name
               -> TypeM (Env, TypeBindBase Info VName)
-checkTypeBind (TypeBind name ps td doc loc) =
-  checkTypeParams ps $ \ps' -> do
-    (td', l) <- bindingTypeParams ps' $ checkTypeDecl ps' td
+checkTypeBind (TypeBind name tps (TypeDecl t NoInfo) doc loc) =
+  checkTypeParams tps $ \tps' -> do
+    (td', l) <- bindingTypeParams tps' $ do
+      checkForDuplicateNamesInType t
+      (t', st, l) <- checkTypeExp t
+      checkShapeParamUses typeExpUses tps' [t']
+      return (TypeDecl t' $ Info st, l)
     bindSpaced [(Type, name)] $ do
       name' <- checkName Type name loc
       return (mempty { envTypeTable =
-                         M.singleton name' $ TypeAbbr l ps' $ unInfo $ expandedType td',
+                         M.singleton name' $ TypeAbbr l tps' $ unInfo $ expandedType td',
                        envNameMap =
                          M.singleton (Type, name) $ qualName name'
                      },
-              TypeBind name' ps' td' doc loc)
+              TypeBind name' tps' td' doc loc)
 
 checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind)
 checkValBind (ValBind entry fname maybe_tdecl NoInfo tparams params body doc loc) = do
diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs
--- a/src/Language/Futhark/TypeChecker/Monad.hs
+++ b/src/Language/Futhark/TypeChecker/Monad.hs
@@ -210,6 +210,15 @@
   lookupImport :: SrcLoc -> FilePath -> m (FilePath, Env)
   lookupVar :: SrcLoc -> QualName Name -> m (QualName VName, PatternType)
 
+  checkNamedDim :: SrcLoc -> QualName Name -> m (QualName VName)
+  checkNamedDim loc v = do
+    (v', t) <- lookupVar loc v
+    case t of
+      Prim (Signed Int32) -> return v'
+      _                   -> throwError $ TypeError loc $
+                             "Dimension declaration " ++ pretty v ++
+                             " should be of type `i32`."
+
 checkName :: MonadTypeChecker m => Namespace -> Name -> SrcLoc -> m VName
 checkName space name loc = qualLeaf <$> checkQualName space (qualName name) loc
 
diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs
--- a/src/Language/Futhark/TypeChecker/Terms.hs
+++ b/src/Language/Futhark/TypeChecker/Terms.hs
@@ -330,6 +330,11 @@
               (map (maybe (toStruct argtype) Prim) pts,
                maybe (toStruct argtype) Prim rt)
 
+  checkNamedDim loc v = do
+    (v', t) <- lookupVar loc v
+    unify loc (toStructural t) (Prim $ Signed Int32)
+    return v'
+
 checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName)
 checkQualNameWithEnv space qn@(QualName [q] _) loc
   | nameToString q == "intrinsics" = do
@@ -678,19 +683,16 @@
 
     descend [] orig_ps
 
-bindingPattern :: [UncheckedTypeParam]
-               -> PatternBase NoInfo Name -> InferredType
-               -> ([TypeParam] -> Pattern -> TermTypeM a) -> TermTypeM a
-bindingPattern tps p t m = do
+bindingPattern :: PatternBase NoInfo Name -> InferredType
+               -> (Pattern -> TermTypeM a) -> TermTypeM a
+bindingPattern p t m = do
   checkForDuplicateNames [p]
-  checkTypeParams tps $ \tps' -> bindingTypeParams tps' $
-    checkPattern p t $ \p' -> binding (S.toList $ patternIdents p') $ do
-      -- Perform an observation of every declared dimension.  This
-      -- prevents unused-name warnings for otherwise unused dimensions.
-      mapM_ observe $ patternDims p'
-      checkShapeParamUses patternUses tps' [p']
+  checkPattern p t $ \p' -> binding (S.toList $ patternIdents p') $ do
+    -- Perform an observation of every declared dimension.  This
+    -- prevents unused-name warnings for otherwise unused dimensions.
+    mapM_ observe $ patternDims p'
 
-      m tps' p'
+    m p'
 
 -- | Return the shapes used in a given pattern in postive and negative
 -- position, respectively.
@@ -704,14 +706,6 @@
 patternUses (PatternAscription p (TypeDecl declte _) _) =
   patternUses p <> typeExpUses declte
 
-noTypeParamsPermitted :: [UncheckedTypeParam] -> TermTypeM ()
-noTypeParamsPermitted ps =
-  case mapMaybe typeParamLoc ps of
-    loc:_ -> typeError loc "Type parameters are not permitted here."
-    []    -> return ()
-  where typeParamLoc (TypeParamDim _ _) = Nothing
-        typeParamLoc tparam             = Just $ srclocOf tparam
-
 patternDims :: Pattern -> [Ident]
 patternDims (PatternParens p _) = patternDims p
 patternDims (TuplePattern pats _) = concatMap patternDims pats
@@ -930,8 +924,7 @@
   (t1, rt) <- checkApply loc t arg
   return $ Apply e1' e2' (Info $ diet t1) (Info rt) loc
 
-checkExp (LetPat tparams pat e body NoInfo loc) = do
-  noTypeParamsPermitted tparams
+checkExp (LetPat pat e body NoInfo loc) =
   sequentially (checkExp e) $ \e' e_occs -> do
     -- Not technically an ascription, but we want the pattern to have
     -- exactly the type of 'e'.
@@ -941,10 +934,10 @@
         let msg = "of value computed with consumption at " ++ locStr (location c)
         in zeroOrderType loc msg t
       _ -> return ()
-    bindingPattern tparams pat (Ascribed $ anyDimShapeAnnotations t) $ \tparams' pat' -> do
+    bindingPattern pat (Ascribed $ anyDimShapeAnnotations t) $ \pat' -> do
       body' <- checkExp body
       body_t <- unscopeType (S.map identName $ patternIdents pat') <$> expType body'
-      return $ LetPat tparams' pat' e' body' (Info body_t) loc
+      return $ LetPat pat' e' body' (Info body_t) loc
 
 checkExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body loc) =
   sequentially (checkFunDef' (name, maybe_retdecl, tparams, params, e, loc)) $
@@ -1042,9 +1035,9 @@
   e2' <- checkExp e2
   return $ Assert e1' e2' (Info (pretty e1)) loc
 
-checkExp (Lambda tparams params body rettype_te NoInfo loc) =
+checkExp (Lambda params body rettype_te NoInfo loc) =
   removeSeminullOccurences $
-  bindingPatternGroup tparams params $ \tparams' params' -> do
+  bindingPatternGroup [] params $ \_ params' -> do
     rettype_checked <- traverse checkTypeExp rettype_te
     let declared_rettype =
           case rettype_checked of Just (_, st, _) -> Just st
@@ -1061,7 +1054,7 @@
 
     closure' <- lexicalClosure params' closure
 
-    return $ Lambda tparams' params' body' rettype' (Info (closure', rettype_st)) loc
+    return $ Lambda params' body' rettype' (Info (closure', rettype_st)) loc
 
 checkExp (OpSection op _ loc) = do
   (op', ftype) <- lookupVar loc op
@@ -1104,11 +1097,9 @@
   where isFix DimFix{} = True
         isFix _        = False
 
-checkExp (DoLoop tparams mergepat mergeexp form loopbody loc) =
+checkExp (DoLoop mergepat mergeexp form loopbody loc) =
   sequentially (checkExp mergeexp) $ \mergeexp' _ -> do
 
-  noTypeParamsPermitted tparams
-
   zeroOrderType (srclocOf mergeexp) "used as loop variable" (typeOf mergeexp')
 
   merge_t <- do
@@ -1122,17 +1113,16 @@
   --
   -- Play a little with occurences to ensure it does not look like
   -- none of the merge variables are being used.
-  ((tparams', mergepat', form', loopbody'), bodyflow) <-
+  ((mergepat', form', loopbody'), bodyflow) <-
     case form of
       For i uboundexp -> do
         uboundexp' <- require anySignedType =<< checkExp uboundexp
         bound_t <- expType uboundexp'
         bindingIdent i bound_t $ \i' ->
-          noUnique $ bindingPattern tparams mergepat merge_t $
-          \tparams' mergepat' -> onlySelfAliasing $ tapOccurences $ do
+          noUnique $ bindingPattern mergepat merge_t $
+          \mergepat' -> onlySelfAliasing $ tapOccurences $ do
             loopbody' <- checkExp loopbody
-            return (tparams',
-                    mergepat',
+            return (mergepat',
                     For i' uboundexp',
                     loopbody')
 
@@ -1142,12 +1132,11 @@
         t <- expType e'
         case t of
           _ | Just t' <- peelArray 1 t ->
-                bindingPattern [] xpat (Ascribed t') $ \_ xpat' ->
-                noUnique $ bindingPattern tparams mergepat merge_t $
-                \tparams' mergepat' -> onlySelfAliasing $ tapOccurences $ do
+                bindingPattern xpat (Ascribed t') $ \xpat' ->
+                noUnique $ bindingPattern mergepat merge_t $
+                \mergepat' -> onlySelfAliasing $ tapOccurences $ do
                   loopbody' <- checkExp loopbody
-                  return (tparams',
-                          mergepat',
+                  return (mergepat',
                           ForIn xpat' e',
                           loopbody')
             | otherwise ->
@@ -1155,12 +1144,11 @@
                 "Iteratee of a for-in loop must be an array, but expression has type " ++ pretty t
 
       While cond ->
-        noUnique $ bindingPattern tparams mergepat merge_t $ \tparams' mergepat' ->
+        noUnique $ bindingPattern mergepat merge_t $ \mergepat' ->
         onlySelfAliasing $ tapOccurences $
         sequentially (unifies (Prim Bool) =<< checkExp cond) $ \cond' _ -> do
           loopbody' <- checkExp loopbody
-          return (tparams',
-                  mergepat',
+          return (mergepat',
                   While cond',
                   loopbody')
 
@@ -1179,7 +1167,7 @@
       consumeMerge _ _ =
         return ()
   consumeMerge mergepat'' =<< expType mergeexp'
-  return $ DoLoop tparams' mergepat'' mergeexp' form' loopbody' loc
+  return $ DoLoop mergepat'' mergeexp' form' loopbody' loc
 
   where
     convergePattern pat body_cons body_t body_loc = do
@@ -1288,7 +1276,7 @@
 checkCase :: PatternType -> CaseBase NoInfo Name
           -> TermTypeM (CaseBase Info VName, PatternType)
 checkCase mt (CasePat p caseExp loc) =
-  bindingPattern [] p (Ascribed mt) $ \_ p' -> do
+  bindingPattern p (Ascribed mt) $ \p' -> do
     caseExp' <- checkExp caseExp
     caseType <- expType caseExp'
     return (CasePat p' caseExp' loc, caseType)
diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs
--- a/src/Language/Futhark/TypeChecker/Types.hs
+++ b/src/Language/Futhark/TypeChecker/Types.hs
@@ -8,6 +8,7 @@
   , subuniqueOf
 
   , checkForDuplicateNames
+  , checkForDuplicateNamesInType
   , checkTypeParams
 
   , typeExpUses
@@ -245,16 +246,6 @@
   unless (length names <= 256) $
     throwError $ TypeError loc "Enums must have 256 or fewer constructors."
   return (TEEnum names loc, Enum names,  Unlifted)
-
-checkNamedDim :: MonadTypeChecker m =>
-                 SrcLoc -> QualName Name -> m (QualName VName)
-checkNamedDim loc v = do
-  (v', t) <- lookupVar loc v
-  case t of
-    Prim (Signed Int32) -> return v'
-    _                   -> throwError $ TypeError loc $
-                           "Dimension declaration " ++ pretty v ++
-                           " should be of type `i32`."
 
 -- | Check for duplication of names inside a pattern group.  Produces
 -- a description of all names used in the pattern group.
diff --git a/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs b/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs
--- a/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs
+++ b/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs
@@ -118,42 +118,42 @@
 
 
 index :: (IntegralExp num, Eq num) =>
-         IxFun num -> Indices num -> num -> num
-index (Direct dims) is element_size =
-  sum (zipWith (*) is slicesizes) * element_size
+         IxFun num -> Indices num -> num
+index (Direct dims) is =
+  sum $ zipWith (*) is slicesizes
   where slicesizes = drop 1 $ sliceSizes dims
-index (Permute fun perm) is_new element_size =
-  index fun is_old element_size
+index (Permute fun perm) is_new =
+  index fun is_old
   where is_old = rearrangeShape (rearrangeInverse perm) is_new
-index (Rotate fun offsets) is element_size =
-  index fun (zipWith mod (zipWith (+) is offsets) dims) element_size
+index (Rotate fun offsets) is =
+  index fun $ zipWith mod (zipWith (+) is offsets) dims
   where dims = shape fun
-index (Index fun js) is element_size =
-  index fun (adjust js is) element_size
+index (Index fun js) is =
+  index fun (adjust js is)
   where adjust (DimFix j:js') is' = j : adjust js' is'
         adjust (DimSlice j _ s:js') (i:is') = j + i * s : adjust js' is'
         adjust _ _ = []
-index (Reshape fun newshape) is element_size =
+index (Reshape fun newshape) is =
   let new_indices = reshapeIndex (shape fun) (newDims newshape) is
-  in index fun new_indices element_size
-index (Repeat fun outer_shapes _) is element_size =
+  in index fun new_indices
+index (Repeat fun outer_shapes _) is =
   -- Discard those indices that are just repeats.  It is intentional
   -- that we cut off those indices that correspond to the innermost
   -- repeated dimensions.
-  index fun is' element_size
+  index fun is'
   where flags dims = replicate (length dims) True ++ [False]
         is' = map snd $ filter (not . fst) $ zip (concatMap flags outer_shapes) is
-index (OffsetIndex fun i) is element_size =
+index (OffsetIndex fun i) is =
   case shape fun of
     d : ds ->
-      index (Index fun (DimSlice i (d-i) 1 : map (unitSlice 0) ds)) is element_size
+      index (Index fun (DimSlice i (d-i) 1 : map (unitSlice 0) ds)) is
     [] -> error "index: OffsetIndex: underlying index function has rank zero"
-index (StrideIndex fun s) is element_size =
+index (StrideIndex fun s) is =
   case shape fun of
     d : ds ->
-      index (Index fun (DimSlice 0 d s : map (unitSlice 0) ds)) is element_size
+      index (Index fun (DimSlice 0 d s : map (unitSlice 0) ds)) is
     [] -> error "index: StrideIndex: underlying index function has rank zero"
-index (Rebase new_base fun) is element_size =
+index (Rebase new_base fun) is =
   let fun' = case fun of
                Direct old_shape ->
                  if old_shape == shape new_base
@@ -175,4 +175,4 @@
                  offsetIndex (rebase new_base ixfun) s
                r@Rebase{} ->
                  r
-  in index fun' is element_size
+  in index fun' is
diff --git a/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs
--- a/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs
+++ b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs
@@ -50,8 +50,8 @@
   let lmadShape = IxFunLMAD.shape ixfunLMAD
       algShape = IxFunAlg.shape ixfunAlg
       points = allPoints lmadShape
-      resLMAD = map (\is -> IxFunLMAD.index ixfunLMAD is 4) points
-      resAlg = map (\is -> IxFunAlg.index ixfunAlg is 4) points
+      resLMAD = map (IxFunLMAD.index ixfunLMAD) points
+      resAlg = map (IxFunAlg.index ixfunAlg) points
       errorMessage = "lmad ixfun:  " ++ PR.pretty ixfunLMAD ++ "\n" ++
                      "alg ixfun:   " ++ PR.pretty ixfunAlg ++ "\n" ++
                      "lmad shape:  " ++ show lmadShape ++ "\n" ++