diff --git a/docs/installation.rst b/docs/installation.rst
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -219,13 +219,13 @@
 Also `Git for Windows`_ is required for its Linux command line tools.
 If you have not marked the option to add them to path, there are
 instructions below how to do so. The GUI alternative to ``git``,
-`Github Desktop`_ is optional and does not come with the required
+`GitHub Desktop`_ is optional and does not come with the required
 tools.
 
 .. _`CUDA 7.5`: https://developer.nvidia.com/cuda-downloads
 .. _`Anaconda`: https://www.continuum.io/downloads#_windows
 .. _`Git for Windows`: https://git-scm.com/download/win
-.. _`Github Desktop`: https://desktop.github.com/
+.. _`GitHub Desktop`: https://desktop.github.com/
 
 Setting up Futhark and OpenCL
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -1,7 +1,7 @@
 cabal-version: 2.4
 
 name:           futhark
-version:        0.15.4
+version:        0.15.5
 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,
diff --git a/rts/c/atomics.h b/rts/c/atomics.h
new file mode 100644
--- /dev/null
+++ b/rts/c/atomics.h
@@ -0,0 +1,197 @@
+// Start of atomics.h
+
+inline int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicAdd((int32_t*)p, x);
+#else
+  return atomic_add(p, x);
+#endif
+}
+
+inline int32_t atomic_add_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicAdd((int32_t*)p, x);
+#else
+  return atomic_add(p, x);
+#endif
+}
+
+inline float atomic_fadd_f32_global(volatile __global float *p, float x) {
+#ifdef FUTHARK_CUDA
+  return atomicAdd((float*)p, x);
+#else
+  union { int32_t i; float f; } old;
+  union { int32_t i; float f; } assumed;
+  old.f = *p;
+  do {
+    assumed.f = old.f;
+    old.f = old.f + x;
+    old.i = atomic_cmpxchg((volatile __global int32_t*)p, assumed.i, old.i);
+  } while (assumed.i != old.i);
+  return old.f;
+#endif
+}
+
+inline float atomic_fadd_f32_local(volatile __local float *p, float x) {
+#ifdef FUTHARK_CUDA
+  return atomicAdd((float*)p, x);
+#else
+  union { int32_t i; float f; } old;
+  union { int32_t i; float f; } assumed;
+  old.f = *p;
+  do {
+    assumed.f = old.f;
+    old.f = old.f + x;
+    old.i = atomic_cmpxchg((volatile __local int32_t*)p, assumed.i, old.i);
+  } while (assumed.i != old.i);
+  return old.f;
+#endif
+}
+
+inline int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMax((int32_t*)p, x);
+#else
+  return atomic_max(p, x);
+#endif
+}
+
+inline int32_t atomic_smax_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMax((int32_t*)p, x);
+#else
+  return atomic_max(p, x);
+#endif
+}
+
+inline int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMin((int32_t*)p, x);
+#else
+  return atomic_min(p, x);
+#endif
+}
+
+inline int32_t atomic_smin_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMin((int32_t*)p, x);
+#else
+  return atomic_min(p, x);
+#endif
+}
+
+inline uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMax((uint32_t*)p, x);
+#else
+  return atomic_max(p, x);
+#endif
+}
+
+inline uint32_t atomic_umax_i32_local(volatile __local uint32_t *p, uint32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMax((uint32_t*)p, x);
+#else
+  return atomic_max(p, x);
+#endif
+}
+
+inline uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMin((uint32_t*)p, x);
+#else
+  return atomic_min(p, x);
+#endif
+}
+
+inline uint32_t atomic_umin_i32_local(volatile __local uint32_t *p, uint32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicMin((uint32_t*)p, x);
+#else
+  return atomic_min(p, x);
+#endif
+}
+
+inline int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicAnd((int32_t*)p, x);
+#else
+  return atomic_and(p, x);
+#endif
+}
+
+inline int32_t atomic_and_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicAnd((int32_t*)p, x);
+#else
+  return atomic_and(p, x);
+#endif
+}
+
+inline int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicOr((int32_t*)p, x);
+#else
+  return atomic_or(p, x);
+#endif
+}
+
+inline int32_t atomic_or_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicOr((int32_t*)p, x);
+#else
+  return atomic_or(p, x);
+#endif
+}
+
+inline int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicXor((int32_t*)p, x);
+#else
+  return atomic_xor(p, x);
+#endif
+}
+
+inline int32_t atomic_xor_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicXor((int32_t*)p, x);
+#else
+  return atomic_xor(p, x);
+#endif
+}
+
+inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicExch((int32_t*)p, x);
+#else
+  return atomic_xor(p, x);
+#endif
+}
+
+inline int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x) {
+#ifdef FUTHARK_CUDA
+  return atomicExch((int32_t*)p, x);
+#else
+  return atomic_xor(p, x);
+#endif
+}
+
+inline int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p,
+                                         int32_t cmp, int32_t val) {
+#ifdef FUTHARK_CUDA
+  return atomicCAS((int32_t*)p, cmp, val);
+#else
+  return atomic_cmpxchg(p, cmp, val);
+#endif
+}
+
+inline int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p,
+                                         int32_t cmp, int32_t val) {
+#ifdef FUTHARK_CUDA
+  return atomicCAS((int32_t*)p, cmp, val);
+#else
+  return atomic_cmpxchg(p, cmp, val);
+#endif
+}
+
+// End of atomics.h
diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs
--- a/src/Futhark/Actions.hs
+++ b/src/Futhark/Actions.hs
@@ -54,5 +54,5 @@
 kernelImpCodeGenAction =
   Action { actionName = "Compile imperative kernels"
          , actionDescription = "Translate program into imperative IL with kernels and write it on standard output."
-         , actionProcedure = liftIO . putStrLn . pretty <=< ImpGenKernels.compileProg
+         , actionProcedure = liftIO . putStrLn . pretty <=< ImpGenKernels.compileProgOpenCL
          }
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
@@ -1173,7 +1173,7 @@
 
   let run_it = [C.citems|
                   int r;
-                  /* Run the program once. */
+                  // Run the program once.
                   $stms:pack_input
                   if ($id:sync_ctx(ctx) != 0) {
                     futhark_panic(1, "%s", $id:error_ctx(ctx));
@@ -1210,7 +1210,7 @@
     // We do not want to profile all the initialisation.
     $id:pause_profiling(ctx);
 
-    /* Declare and read input. */
+    // Declare and read input.
     set_binary_mode(stdin);
     $items:input_items
 
@@ -1220,13 +1220,13 @@
 
     $items:output_decls
 
-    /* Warmup run */
+    // Warmup run
     if (perform_warmup) {
       $items:run_it
       $stms:free_outputs
     }
     time_runs = 1;
-    /* Proper run. */
+    // Proper run.
     for (int run = 0; run < num_runs; run++) {
       // Only profile last run.
       profile_run = run == num_runs -1;
@@ -1236,10 +1236,10 @@
       }
     }
 
-    /* Free the parsed input. */
+    // Free the parsed input.
     $stms:free_parsed
 
-    /* Print the final result. */
+    // Print the final result.
     if (binary_output) {
       set_binary_mode(stdout);
     }
@@ -1339,25 +1339,25 @@
       option_parser = generateOptionParser "parse_options" $ benchmarkOptions++options
 
   let headerdefs = [C.cunit|
-$esc:("/*\n * Headers\n*/\n")
+$esc:("// Headers\n")
 $esc:("#include <stdint.h>")
 $esc:("#include <stddef.h>")
 $esc:("#include <stdbool.h>")
 $esc:(header_extra)
 
-$esc:("\n/*\n * Initialisation\n*/\n")
+$esc:("\n// Initialisation\n")
 $edecls:(initDecls endstate)
 
-$esc:("\n/*\n * Arrays\n*/\n")
+$esc:("\n// Arrays\n")
 $edecls:(arrayDecls endstate)
 
-$esc:("\n/*\n * Opaque values\n*/\n")
+$esc:("\n// Opaque values\n")
 $edecls:(opaqueDecls endstate)
 
-$esc:("\n/*\n * Entry points\n*/\n")
+$esc:("\n// Entry points\n")
 $edecls:(entryDecls endstate)
 
-$esc:("\n/*\n * Miscellaneous\n*/\n")
+$esc:("\n// Miscellaneous\n")
 $edecls:(miscDecls endstate)
                            |]
 
@@ -1367,9 +1367,9 @@
 $esc:("#include <stdbool.h>")
 $esc:("#include <math.h>")
 $esc:("#include <stdint.h>")
-/* If NDEBUG is set, the assert() macro will do nothing. Since Futhark
-   (unfortunately) makes use of assert() for error detection (and even some
-   side effects), we want to avoid that. */
+// If NDEBUG is set, the assert() macro will do nothing. Since Futhark
+// (unfortunately) makes use of assert() for error detection (and even some
+// side effects), we want to avoid that.
 $esc:("#undef NDEBUG")
 $esc:("#include <assert.h>")
 
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
@@ -1,4 +1,5 @@
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE OverloadedStrings #-}
 -- | Variation of "Futhark.CodeGen.ImpCode" that contains the notion
 -- of a kernel invocation.
 module Futhark.CodeGen.ImpCode.Kernels
@@ -17,8 +18,6 @@
   , KernelUse (..)
   , module Futhark.CodeGen.ImpCode
   , module Futhark.Representation.Kernels.Sizes
-  -- * Utility functions
-  , atomicBinOp
   )
   where
 
@@ -77,18 +76,6 @@
                | ConstUse VName KernelConstExp
                  deriving (Eq, Show)
 
--- | Get an atomic operator corresponding to a binary operator.
-atomicBinOp :: BinOp -> Maybe (VName -> VName -> Count Elements Imp.Exp -> Exp -> AtomicOp)
-atomicBinOp = flip lookup [ (Add Int32, AtomicAdd)
-                          , (SMax Int32, AtomicSMax)
-                          , (SMin Int32, AtomicSMin)
-                          , (UMax Int32, AtomicUMax)
-                          , (UMin Int32, AtomicUMin)
-                          , (And Int32, AtomicAnd)
-                          , (Or Int32, AtomicOr)
-                          , (Xor Int32, AtomicXor)
-                          ]
-
 instance Pretty KernelConst where
   ppr (SizeConst key) = text "get_size" <> parens (ppr key)
 
@@ -162,92 +149,97 @@
 
 -- Atomic operations return the value stored before the update.
 -- This value is stored in the first VName.
-data AtomicOp = AtomicAdd VName VName (Count Elements Imp.Exp) Exp
-              | AtomicSMax VName VName (Count Elements Imp.Exp) Exp
-              | AtomicSMin VName VName (Count Elements Imp.Exp) Exp
-              | AtomicUMax VName VName (Count Elements Imp.Exp) Exp
-              | AtomicUMin VName VName (Count Elements Imp.Exp) Exp
-              | AtomicAnd VName VName (Count Elements Imp.Exp) Exp
-              | AtomicOr VName VName (Count Elements Imp.Exp) Exp
-              | AtomicXor VName VName (Count Elements Imp.Exp) Exp
-              | AtomicCmpXchg VName VName (Count Elements Imp.Exp) Exp Exp
-              | AtomicXchg VName VName (Count Elements Imp.Exp) Exp
+data AtomicOp = AtomicAdd IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicFAdd FloatType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicSMax IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicSMin IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicUMax IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicUMin IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicAnd IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicOr IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicXor IntType VName VName (Count Elements Imp.Exp) Exp
+              | AtomicCmpXchg PrimType VName VName (Count Elements Imp.Exp) Exp Exp
+              | AtomicXchg PrimType VName VName (Count Elements Imp.Exp) Exp
               deriving (Show)
 
 instance FreeIn AtomicOp where
-  freeIn' (AtomicAdd _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicSMax _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicSMin _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicUMax _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicUMin _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicAnd _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicOr _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicXor _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
-  freeIn' (AtomicCmpXchg _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y
-  freeIn' (AtomicXchg _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicFAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicSMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicSMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicUMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicUMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
+  freeIn' (AtomicCmpXchg _ _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y
+  freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x
 
 instance Pretty KernelOp where
   ppr (GetGroupId dest i) =
-    ppr dest <+> text "<-" <+>
-    text "get_group_id" <> parens (ppr i)
+    ppr dest <+>  "<-" <+>
+     "get_group_id" <> parens (ppr i)
   ppr (GetLocalId dest i) =
-    ppr dest <+> text "<-" <+>
-    text "get_local_id" <> parens (ppr i)
+    ppr dest <+>  "<-" <+>
+     "get_local_id" <> parens (ppr i)
   ppr (GetLocalSize dest i) =
-    ppr dest <+> text "<-" <+>
-    text "get_local_size" <> parens (ppr i)
+    ppr dest <+>  "<-" <+>
+     "get_local_size" <> parens (ppr i)
   ppr (GetGlobalSize dest i) =
-    ppr dest <+> text "<-" <+>
-    text "get_global_size" <> parens (ppr i)
+    ppr dest <+>  "<-" <+>
+     "get_global_size" <> parens (ppr i)
   ppr (GetGlobalId dest i) =
-    ppr dest <+> text "<-" <+>
-    text "get_global_id" <> parens (ppr i)
+    ppr dest <+>  "<-" <+>
+     "get_global_id" <> parens (ppr i)
   ppr (GetLockstepWidth dest) =
-    ppr dest <+> text "<-" <+>
-    text "get_lockstep_width()"
+    ppr dest <+>  "<-" <+>
+     "get_lockstep_width()"
   ppr (Barrier FenceLocal) =
-    text "local_barrier()"
+     "local_barrier()"
   ppr (Barrier FenceGlobal) =
-    text "global_barrier()"
+     "global_barrier()"
   ppr (MemFence FenceLocal) =
-    text "mem_fence_local()"
+     "mem_fence_local()"
   ppr (MemFence FenceGlobal) =
-    text "mem_fence_global()"
+     "mem_fence_global()"
   ppr (LocalAlloc name size) =
-    ppr name <+> equals <+> text "local_alloc" <> parens (ppr size)
+    ppr name <+> equals <+>  "local_alloc" <> parens (ppr size)
   ppr (ErrorSync FenceLocal) =
-    text "error_sync_local()"
+     "error_sync_local()"
   ppr (ErrorSync FenceGlobal) =
-    text "error_sync_global()"
-  ppr (Atomic _ (AtomicAdd old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_add" <>
+     "error_sync_global()"
+  ppr (Atomic _ (AtomicAdd t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_add_" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicSMax old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_smax" <>
+  ppr (Atomic _ (AtomicFAdd t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_fadd_" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicSMin old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_smin" <>
+  ppr (Atomic _ (AtomicSMax t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_smax" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicUMax old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_umax" <>
+  ppr (Atomic _ (AtomicSMin t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_smin" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicUMin old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_umin" <>
+  ppr (Atomic _ (AtomicUMax t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_umax" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicAnd old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_and" <>
+  ppr (Atomic _ (AtomicUMin t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_umin" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicOr old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_or" <>
+  ppr (Atomic _ (AtomicAnd t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_and" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicXor old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_xor" <>
+  ppr (Atomic _ (AtomicOr t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_or" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
-  ppr (Atomic _ (AtomicCmpXchg old arr ind x y)) =
-    ppr old <+> text "<-" <+> text "atomic_cmp_xchg" <>
+  ppr (Atomic _ (AtomicXor t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_xor" <> ppr t <>
+    parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
+  ppr (Atomic _ (AtomicCmpXchg t old arr ind x y)) =
+    ppr old <+>  "<-" <+>  "atomic_cmp_xchg" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y])
-  ppr (Atomic _ (AtomicXchg old arr ind x)) =
-    ppr old <+> text "<-" <+> text "atomic_xchg" <>
+  ppr (Atomic _ (AtomicXchg t old arr ind x)) =
+    ppr old <+>  "<-" <+>  "atomic_xchg" <> ppr t <>
     parens (commasep [ppr arr <> brackets (ppr ind), ppr x])
 
 instance FreeIn KernelOp where
@@ -255,4 +247,4 @@
   freeIn' _ = mempty
 
 brace :: Doc -> Doc
-brace body = text " {" </> indent 2 body </> text "}"
+brace body =  " {" </> indent 2 body </>  "}"
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
@@ -19,7 +19,8 @@
 
     -- * Monadic Compiler Interface
   , ImpM
-  , Env (envDefaultSpace, envFunction)
+  , localDefaultSpace, askFunction
+  , askEnv, localEnv
   , VTable
   , getVTable
   , localVTable
@@ -89,7 +90,7 @@
 import qualified Data.Map.Strict as M
 import qualified Data.Set as S
 import Data.Maybe
-import Data.List (find, foldl', sortOn)
+import Data.List (find, sortOn)
 
 import qualified Futhark.CodeGen.ImpCode as Imp
 import Futhark.CodeGen.ImpCode
@@ -103,33 +104,33 @@
 import Futhark.Util
 
 -- | How to compile an 'Op'.
-type OpCompiler lore op = Pattern lore -> Op lore -> ImpM lore op ()
+type OpCompiler lore r op = Pattern lore -> Op lore -> ImpM lore r op ()
 
 -- | How to compile some 'Stms'.
-type StmsCompiler lore op = Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()
+type StmsCompiler lore r op = Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op ()
 
 -- | How to compile an 'Exp'.
-type ExpCompiler lore op = Pattern lore -> Exp lore -> ImpM lore op ()
+type ExpCompiler lore r op = Pattern lore -> Exp lore -> ImpM lore r op ()
 
-type CopyCompiler lore op = PrimType
+type CopyCompiler lore r op = PrimType
                            -> MemLocation
                            -> MemLocation
-                           -> ImpM lore op ()
+                           -> ImpM lore r op ()
 
 -- | An alternate way of compiling an allocation.
-type AllocCompiler lore op = VName -> Count Bytes Imp.Exp -> ImpM lore op ()
+type AllocCompiler lore r op = VName -> Count Bytes Imp.Exp -> ImpM lore r 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)
+data Operations lore r op = Operations { opsExpCompiler :: ExpCompiler lore r op
+                                     , opsOpCompiler :: OpCompiler lore r op
+                                     , opsStmsCompiler :: StmsCompiler lore r op
+                                     , opsCopyCompiler :: CopyCompiler lore r op
+                                     , opsAllocCompilers :: M.Map Space (AllocCompiler lore r op)
                                      }
 
 -- | An operations set for which the expression compiler always
 -- returns 'CompileExp'.
 defaultOperations :: (ExplicitMemorish lore, FreeIn op) =>
-                     OpCompiler lore op -> Operations lore op
+                     OpCompiler lore r op -> Operations lore r op
 defaultOperations opc = Operations { opsExpCompiler = defCompileExp
                                    , opsOpCompiler = opc
                                    , opsStmsCompiler = defCompileStms
@@ -185,20 +186,22 @@
                         -- takes care of this array.
                       deriving (Show)
 
-data Env lore op = Env {
-    envExpCompiler :: ExpCompiler lore op
-  , envStmsCompiler :: StmsCompiler lore op
-  , envOpCompiler :: OpCompiler lore op
-  , envCopyCompiler :: CopyCompiler lore op
-  , envAllocCompilers :: M.Map Space (AllocCompiler lore op)
+data Env lore r op = Env {
+    envExpCompiler :: ExpCompiler lore r op
+  , envStmsCompiler :: StmsCompiler lore r op
+  , envOpCompiler :: OpCompiler lore r op
+  , envCopyCompiler :: CopyCompiler lore r op
+  , envAllocCompilers :: M.Map Space (AllocCompiler lore r op)
   , envDefaultSpace :: Imp.Space
   , envVolatility :: Imp.Volatility
-  , envFunction :: Name
-    -- ^ Name of the function we are compiling.
+  , envEnv :: r
+    -- ^ User-extensible environment.
+  , envFunction :: Maybe Name
+    -- ^ Name of the function we are compiling, if any.
   }
 
-newEnv :: Operations lore op -> Imp.Space -> Name -> Env lore op
-newEnv ops ds fname =
+newEnv :: r -> Operations lore r op -> Imp.Space -> Env lore r op
+newEnv r ops ds =
   Env { envExpCompiler = opsExpCompiler ops
       , envStmsCompiler = opsStmsCompiler ops
       , envOpCompiler = opsOpCompiler ops
@@ -206,33 +209,34 @@
       , envAllocCompilers = mempty
       , envDefaultSpace = ds
       , envVolatility = Imp.Nonvolatile
-      , envFunction = fname
+      , envEnv = r
+      , envFunction = Nothing
       }
 
 -- | The symbol table used during compilation.
 type VTable lore = M.Map VName (VarEntry lore)
 
-data State lore op = State { stateVTable :: VTable lore
+data State lore r op = State { stateVTable :: VTable lore
                            , stateFunctions :: Imp.Functions op
                            , stateNameSource :: VNameSource
                            }
 
-newState :: VNameSource -> State lore op
+newState :: VNameSource -> State lore r op
 newState = State mempty mempty
 
-newtype ImpM lore op a = ImpM (RWS (Env lore op) (Imp.Code op) (State lore op) a)
+newtype ImpM lore r op a = ImpM (RWS (Env lore r op) (Imp.Code op) (State lore r op) a)
   deriving (Functor, Applicative, Monad,
-            MonadState (State lore op),
-            MonadReader (Env lore op),
+            MonadState (State lore r op),
+            MonadReader (Env lore r op),
             MonadWriter (Imp.Code op))
 
-instance MonadFreshNames (ImpM lore op) where
+instance MonadFreshNames (ImpM lore r op) where
   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
+instance HasScope SOACS (ImpM lore r op) where
   askScope = M.map (LetInfo . entryType) <$> gets stateVTable
     where entryType (MemVar _ memEntry) =
             Mem (entryMemSpace memEntry)
@@ -244,18 +248,18 @@
           entryType (ScalarVar _ scalarEntry) =
             Prim $ entryScalarType scalarEntry
 
-runImpM :: ImpM lore op a
-        -> Operations lore op -> Imp.Space -> Name -> State lore op
-        -> (a, State lore op, Imp.Code op)
-runImpM (ImpM m) ops space fname = runRWS m $ newEnv ops space fname
+runImpM :: ImpM lore r op a
+        -> r -> Operations lore r op -> Imp.Space -> State lore r op
+        -> (a, State lore r op, Imp.Code op)
+runImpM (ImpM m) r ops space = runRWS m $ newEnv r ops space
 
-subImpM_ :: Operations lore op' -> ImpM lore op' a
-         -> ImpM lore op (Imp.Code op')
-subImpM_ ops m = snd <$> subImpM ops m
+subImpM_ :: r' -> Operations lore r' op' -> ImpM lore r' op' a
+         -> ImpM lore r op (Imp.Code op')
+subImpM_ r ops m = snd <$> subImpM r ops m
 
-subImpM :: Operations lore op' -> ImpM lore op' a
-        -> ImpM lore op (a, Imp.Code op')
-subImpM ops (ImpM m) = do
+subImpM :: r' -> Operations lore r' op' -> ImpM lore r' op' a
+        -> ImpM lore r op (a, Imp.Code op')
+subImpM r ops (ImpM m) = do
   env <- ask
   s <- get
   let (x, s', code) =
@@ -264,6 +268,7 @@
                      , envCopyCompiler = opsCopyCompiler ops
                      , envOpCompiler = opsOpCompiler ops
                      , envAllocCompilers = opsAllocCompilers ops
+                     , envEnv = r
                      }
                  s { stateVTable = stateVTable s
                    , stateFunctions = mempty }
@@ -272,34 +277,34 @@
 
 -- | Execute a code generation action, returning the code that was
 -- emitted.
-collect :: ImpM lore op () -> ImpM lore op (Imp.Code op)
+collect :: ImpM lore r op () -> ImpM lore r op (Imp.Code op)
 collect m = pass $ do
   ((), code) <- listen m
   return (code, const mempty)
 
-collect' :: ImpM lore op a -> ImpM lore op (a, Imp.Code op)
+collect' :: ImpM lore r op a -> ImpM lore r op (a, Imp.Code op)
 collect' m = pass $ do
   (x, code) <- listen m
   return ((x, code), const mempty)
 
 -- | Execute a code generation action, wrapping the generated code
 -- within a 'Imp.Comment' with the given description.
-comment :: String -> ImpM lore op () -> ImpM lore op ()
+comment :: String -> ImpM lore r op () -> ImpM lore r op ()
 comment desc m = do code <- collect m
                     emit $ Imp.Comment desc code
 
 -- | Emit some generated imperative code.
-emit :: Imp.Code op -> ImpM lore op ()
+emit :: Imp.Code op -> ImpM lore r op ()
 emit = tell
 
 -- | Emit a function in the generated code.
-emitFunction :: Name -> Imp.Function op -> ImpM lore op ()
+emitFunction :: Name -> Imp.Function op -> ImpM lore r op ()
 emitFunction fname fun = do
   Imp.Functions fs <- gets stateFunctions
   modify $ \s -> s { stateFunctions = Imp.Functions $ (fname,fun) : fs }
 
 -- | Check if a function of a given name exists.
-hasFunction :: Name -> ImpM lore op Bool
+hasFunction :: Name -> ImpM lore r op Bool
 hasFunction fname = gets $ \s -> let Imp.Functions fs = stateFunctions s
                                  in isJust $ lookup fname fs
 
@@ -313,26 +318,25 @@
           M.singleton name $ memBoundToVarEntry (Just e) $ infoAttr info
 
 compileProg :: (ExplicitMemorish lore, FreeIn op, MonadFreshNames m) =>
-               Operations lore op -> Imp.Space
+               r -> Operations lore r op -> Imp.Space
             -> Prog lore -> m (Imp.Definitions op)
-compileProg ops space (Prog consts funs) =
+compileProg r ops space (Prog consts funs) =
   modifyNameSource $ \src ->
-  let s = (newState src) { stateVTable = constsVTable consts }
-      s' =
-        foldl' compileFunDef' s funs
-      free_in_funs =
-        freeIn (stateFunctions s')
-      (consts', s'', _) =
-        runImpM (compileConsts free_in_funs consts) ops space
-        (nameFromString "val") s'
-  in (Imp.Definitions consts' (stateFunctions s''),
+  let (consts', s', _) =
+        runImpM compile r ops space
+        (newState src) { stateVTable = constsVTable consts }
+  in (Imp.Definitions consts' (stateFunctions s'),
       stateNameSource s')
-  where compileFunDef' s fdef =
-          let ((), s', _) =
-                runImpM (compileFunDef fdef) ops space (funDefName fdef) s
-          in s'
+  where compile = do
+          mapM_ compileFunDef' funs
+          free_in_funs <- gets (freeIn . stateFunctions)
+          compileConsts free_in_funs consts
 
-compileConsts :: Names -> Stms lore -> ImpM lore op (Imp.Constants op)
+        compileFunDef' fdef =
+          local (\env -> env { envFunction = Just $ funDefName fdef }) $
+          compileFunDef fdef
+
+compileConsts :: Names -> Stms lore -> ImpM lore r op (Imp.Constants op)
 compileConsts used_consts stms = do
   code <- collect $ compileStms used_consts stms $ pure ()
   pure $ uncurry Imp.Constants $ first DL.toList $ extract code
@@ -352,7 +356,7 @@
           (mempty, s)
 
 compileInParam :: ExplicitMemorish lore =>
-                  FParam lore -> ImpM lore op (Either Imp.Param ArrayDecl)
+                  FParam lore -> ImpM lore r op (Either Imp.Param ArrayDecl)
 compileInParam fparam = case paramAttr fparam of
   MemPrim bt ->
     return $ Left $ Imp.ScalarParam name bt
@@ -370,7 +374,7 @@
 
 compileInParams :: ExplicitMemorish lore =>
                    [FParam lore] -> [EntryPointType]
-                -> ImpM lore op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue])
+                -> ImpM lore r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue])
 compileInParams params orig_epts = do
   let (ctx_params, val_params) =
         splitAt (length params - sum (map entryPointSize orig_epts)) params
@@ -419,7 +423,7 @@
 
 compileOutParams :: ExplicitMemorish lore =>
                     [RetType lore] -> [EntryPointType]
-                 -> ImpM lore op ([Imp.ExternalValue], [Imp.Param], Destination)
+                 -> ImpM lore r op ([Imp.ExternalValue], [Imp.Param], Destination)
 compileOutParams orig_rts orig_epts = do
   ((extvs, dests), (outparams,ctx_dests)) <-
     runWriterT $ evalStateT (mkExts orig_epts orig_rts) (M.empty, M.empty)
@@ -481,7 +485,7 @@
 
 compileFunDef :: ExplicitMemorish lore =>
                  FunDef lore
-              -> ImpM lore op ()
+              -> ImpM lore r op ()
 compileFunDef (FunDef entry fname rettype params body) = do
   ((outparams, inparams, results, args), body') <- collect' compile
   emitFunction fname $ Imp.Function (isJust entry) outparams inparams body' results args
@@ -499,18 +503,18 @@
 
           return (outparams, inparams, results, args)
 
-compileBody :: (ExplicitMemorish lore) => Pattern lore -> Body lore -> ImpM lore op ()
+compileBody :: (ExplicitMemorish lore) => Pattern lore -> Body lore -> ImpM lore r op ()
 compileBody pat (Body _ bnds ses) = do
   Destination _ dests <- destinationFromPattern pat
   compileStms (freeIn ses) bnds $
     forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se []
 
-compileBody' :: [Param attr] -> Body lore -> ImpM lore op ()
+compileBody' :: [Param attr] -> Body lore -> ImpM lore r op ()
 compileBody' params (Body _ bnds ses) =
   compileStms (freeIn ses) bnds $
     forM_ (zip params ses) $ \(param, se) -> copyDWIM (paramName param) [] se []
 
-compileLoopBody :: Typed attr => [Param attr] -> Body lore -> ImpM lore op ()
+compileLoopBody :: Typed attr => [Param attr] -> Body lore -> ImpM lore r 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
@@ -533,13 +537,13 @@
         _ -> return $ return ()
     sequence_ copy_to_merge_params
 
-compileStms :: Names -> Stms lore -> ImpM lore op () -> ImpM lore op ()
+compileStms :: Names -> Stms lore -> ImpM lore r op () -> ImpM lore r 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 -> Stms lore -> ImpM lore op () -> ImpM lore op ()
+                  Names -> Stms lore -> ImpM lore r op () -> ImpM lore r 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
@@ -570,13 +574,13 @@
                             Mem space -> Just (patElemName pe, space)
                             _         -> Nothing
 
-compileExp :: Pattern lore -> Exp lore -> ImpM lore op ()
+compileExp :: Pattern lore -> Exp lore -> ImpM lore r op ()
 compileExp pat e = do
   ec <- asks envExpCompiler
   ec pat e
 
 defCompileExp :: (ExplicitMemorish lore) =>
-                 Pattern lore -> Exp lore -> ImpM lore op ()
+                 Pattern lore -> Exp lore -> ImpM lore r op ()
 
 defCompileExp pat (If cond tbranch fbranch _) = do
   tcode <- collect $ compileBody pat tbranch
@@ -631,7 +635,7 @@
   opc pat op
 
 defCompileBasicOp :: ExplicitMemorish lore =>
-                     Pattern lore -> BasicOp lore -> ImpM lore op ()
+                     Pattern lore -> BasicOp lore -> ImpM lore r op ()
 
 defCompileBasicOp (Pattern _ [pe]) (SubExp se) =
   copyDWIM (patElemName pe) [] se []
@@ -753,7 +757,7 @@
   pretty pat ++ "\nfor expression\n  " ++ pretty e
 
 -- | Note: a hack to be used only for functions.
-addArrays :: [ArrayDecl] -> ImpM lore op ()
+addArrays :: [ArrayDecl] -> ImpM lore r op ()
 addArrays = mapM_ addArray
   where addArray (ArrayDecl name bt location) =
           addVar name $
@@ -764,52 +768,52 @@
 
 -- | Like 'dFParams', but does not create new declarations.
 -- Note: a hack to be used only for functions.
-addFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore op ()
+addFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore r op ()
 addFParams = mapM_ addFParam
   where addFParam fparam =
           addVar (paramName fparam) $
           memBoundToVarEntry Nothing $ noUniquenessReturns $ paramAttr fparam
 
 -- | Another hack.
-addLoopVar :: VName -> IntType -> ImpM lore op ()
+addLoopVar :: VName -> IntType -> ImpM lore r op ()
 addLoopVar i it = addVar i $ ScalarVar Nothing $ ScalarEntry $ IntType it
 
 dVars :: ExplicitMemorish lore =>
-            Maybe (Exp lore) -> [PatElem lore] -> ImpM lore op ()
+            Maybe (Exp lore) -> [PatElem lore] -> ImpM lore r op ()
 dVars e = mapM_ dVar
   where dVar = dScope e . scopeOfPatElem
 
-dFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore op ()
+dFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore r op ()
 dFParams = dScope Nothing . scopeOfFParams
 
-dLParams :: ExplicitMemorish lore => [LParam lore] -> ImpM lore op ()
+dLParams :: ExplicitMemorish lore => [LParam lore] -> ImpM lore r op ()
 dLParams = dScope Nothing . scopeOfLParams
 
-dPrimVol_ :: VName -> PrimType -> ImpM lore op ()
+dPrimVol_ :: VName -> PrimType -> ImpM lore r op ()
 dPrimVol_ name t = do
  emit $ Imp.DeclareScalar name Imp.Volatile t
  addVar name $ ScalarVar Nothing $ ScalarEntry t
 
-dPrim_ :: VName -> PrimType -> ImpM lore op ()
+dPrim_ :: VName -> PrimType -> ImpM lore r op ()
 dPrim_ name t = do
  emit $ Imp.DeclareScalar name Imp.Nonvolatile t
  addVar name $ ScalarVar Nothing $ ScalarEntry t
 
-dPrim :: String -> PrimType -> ImpM lore op VName
+dPrim :: String -> PrimType -> ImpM lore r op VName
 dPrim name t = do name' <- newVName name
                   dPrim_ name' t
                   return name'
 
-dPrimV_ :: VName -> Imp.Exp -> ImpM lore op ()
+dPrimV_ :: VName -> Imp.Exp -> ImpM lore r op ()
 dPrimV_ name e = do dPrim_ name $ primExpType e
                     name <-- e
 
-dPrimV :: String -> Imp.Exp -> ImpM lore op VName
+dPrimV :: String -> Imp.Exp -> ImpM lore r op VName
 dPrimV name e = do name' <- dPrim name $ primExpType e
                    name' <-- e
                    return name'
 
-dPrimVE :: String -> Imp.Exp -> ImpM lore op Imp.Exp
+dPrimVE :: String -> Imp.Exp -> ImpM lore r op Imp.Exp
 dPrimVE name e = do name' <- dPrim name $ primExpType e
                     name' <-- e
                     return $ Imp.var name' $ primExpType e
@@ -834,7 +838,7 @@
 infoAttr (IndexInfo it) = MemPrim $ IntType it
 
 dInfo :: Maybe (Exp lore) -> VName -> NameInfo ExplicitMemory
-         -> ImpM lore op ()
+         -> ImpM lore r op ()
 dInfo e name info = do
   let entry = memBoundToVarEntry e $ infoAttr info
   case entry of
@@ -846,19 +850,19 @@
       return ()
   addVar name entry
 
-dScope :: Maybe (Exp lore) -> Scope ExplicitMemory -> ImpM lore op ()
+dScope :: Maybe (Exp lore) -> Scope ExplicitMemory -> ImpM lore r op ()
 dScope e = mapM_ (uncurry $ dInfo e) . M.toList
 
-dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore op ()
+dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op ()
 dArray name bt shape membind =
   addVar name $
   memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind
 
-everythingVolatile :: ImpM lore op a -> ImpM lore op a
+everythingVolatile :: ImpM lore r op a -> ImpM lore r op a
 everythingVolatile = local $ \env -> env { envVolatility = Imp.Volatile }
 
 -- | Remove the array targets.
-funcallTargets :: Destination -> ImpM lore op [VName]
+funcallTargets :: Destination -> ImpM lore r op [VName]
 funcallTargets (Destination _ dests) =
   concat <$> mapM funcallTarget dests
   where funcallTarget (ScalarDestination name) =
@@ -872,7 +876,7 @@
 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
+  toExp :: a -> ImpM lore r op Imp.Exp
   -- | Compile where we know the type in advance.
   toExp' :: PrimType -> a -> Imp.Exp
 
@@ -892,20 +896,32 @@
   toExp = pure . fmap Imp.ScalarVar
   toExp' _ = fmap Imp.ScalarVar
 
-addVar :: VName -> VarEntry lore -> ImpM lore op ()
+addVar :: VName -> VarEntry lore -> ImpM lore r op ()
 addVar name entry =
   modify $ \s -> s { stateVTable = M.insert name entry $ stateVTable s }
 
+localDefaultSpace :: Imp.Space -> ImpM lore r op a -> ImpM lore r op a
+localDefaultSpace space = local (\env -> env { envDefaultSpace = space })
+
+askFunction :: ImpM lore r op (Maybe Name)
+askFunction = asks envFunction
+
+askEnv :: ImpM lore r op r
+askEnv = asks envEnv
+
+localEnv :: (r -> r) -> ImpM lore r op a -> ImpM lore r op a
+localEnv f = local $ \env -> env { envEnv = f $ envEnv env }
+
 -- | Get the current symbol table.
-getVTable :: ImpM lore op (VTable lore)
+getVTable :: ImpM lore r op (VTable lore)
 getVTable = gets stateVTable
 
-putVTable :: VTable lore -> ImpM lore op ()
+putVTable :: VTable lore -> ImpM lore r op ()
 putVTable vtable = modify $ \s -> s { stateVTable = vtable }
 
 -- | Run an action with a modified symbol table.  All changes to the
 -- symbol table will be reverted once the action is done!
-localVTable :: (VTable lore -> VTable lore) -> ImpM lore op a -> ImpM lore op a
+localVTable :: (VTable lore -> VTable lore) -> ImpM lore r op a -> ImpM lore r op a
 localVTable f m = do
   old_vtable <- getVTable
   putVTable $ f old_vtable
@@ -913,28 +929,28 @@
   putVTable old_vtable
   return a
 
-lookupVar :: VName -> ImpM lore op (VarEntry lore)
+lookupVar :: VName -> ImpM lore r op (VarEntry lore)
 lookupVar name = do
   res <- gets $ M.lookup name . stateVTable
   case res of
     Just entry -> return entry
     _ -> error $ "Unknown variable: " ++ pretty name
 
-lookupArray :: VName -> ImpM lore op ArrayEntry
+lookupArray :: VName -> ImpM lore r op ArrayEntry
 lookupArray name = do
   res <- lookupVar name
   case res of
     ArrayVar _ entry -> return entry
     _                -> error $ "ImpGen.lookupArray: not an array: " ++ pretty name
 
-lookupMemory :: VName -> ImpM lore op MemEntry
+lookupMemory :: VName -> ImpM lore r op MemEntry
 lookupMemory name = do
   res <- lookupVar name
   case res of
     MemVar _ entry -> return entry
     _              -> error $ "Unknown memory block: " ++ pretty name
 
-destinationFromPattern :: ExplicitMemorish lore => Pattern lore -> ImpM lore op Destination
+destinationFromPattern :: ExplicitMemorish lore => Pattern lore -> ImpM lore r op Destination
 destinationFromPattern pat =
   fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $
   patternElements pat
@@ -951,13 +967,13 @@
               return $ ScalarDestination name
 
 fullyIndexArray :: VName -> [Imp.Exp]
-                -> ImpM lore op (VName, Imp.Space, Count Elements Imp.Exp)
+                -> ImpM lore r op (VName, Imp.Space, Count Elements Imp.Exp)
 fullyIndexArray name indices = do
   arr <- lookupArray name
   fullyIndexArray' (entryArrayLocation arr) indices
 
 fullyIndexArray' :: MemLocation -> [Imp.Exp]
-                 -> ImpM lore op (VName, Imp.Space, Count Elements Imp.Exp)
+                 -> ImpM lore r op (VName, Imp.Space, Count Elements Imp.Exp)
 fullyIndexArray' (MemLocation mem _ ixfun) indices = do
   space <- entryMemSpace <$> lookupMemory mem
   let indices' = case space of
@@ -979,13 +995,13 @@
 
 -- More complicated read/write operations that use index functions.
 
-copy :: CopyCompiler lore op
+copy :: CopyCompiler lore r op
 copy bt pat src = do
   cc <- asks envCopyCompiler
   cc bt pat src
 
 -- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'.
-defaultCopy :: CopyCompiler lore op
+defaultCopy :: CopyCompiler lore r op
 defaultCopy bt dest src
   | Just destoffset <-
       IxFun.linearWithOffset destIxFun bt_size,
@@ -1008,7 +1024,7 @@
         isScalarSpace ScalarSpace{} = True
         isScalarSpace _ = False
 
-copyElementWise :: CopyCompiler lore op
+copyElementWise :: CopyCompiler lore r op
 copyElementWise bt dest src = do
     let bounds = map (toExp' int32) $ memLocationShape src
     is <- replicateM (length bounds) (newVName "i")
@@ -1025,7 +1041,7 @@
 copyArrayDWIM :: PrimType
               -> MemLocation -> [DimIndex Imp.Exp]
               -> MemLocation -> [DimIndex Imp.Exp]
-              -> ImpM lore op (Imp.Code op)
+              -> ImpM lore r op (Imp.Code op)
 copyArrayDWIM bt
   destlocation@(MemLocation _ destshape _) destslice
   srclocation@(MemLocation _ srcshape _) srcslice
@@ -1064,7 +1080,7 @@
 -- | Like 'copyDWIM', but the target is a 'ValueDestination'
 -- instead of a variable name.
 copyDWIMDest :: ValueDestination -> [DimIndex Imp.Exp] -> SubExp -> [DimIndex Imp.Exp]
-             -> ImpM lore op ()
+             -> ImpM lore r op ()
 
 copyDWIMDest _ _ (Constant v) (_:_) =
   error $
@@ -1151,7 +1167,7 @@
 -- This function will generally just Do What I Mean, and Do The Right
 -- Thing.  Both destination and source must be in scope.
 copyDWIM :: VName -> [DimIndex Imp.Exp] -> SubExp -> [DimIndex Imp.Exp]
-         -> ImpM lore op ()
+         -> ImpM lore r op ()
 copyDWIM dest dest_slice src src_slice = do
   dest_entry <- lookupVar dest
   let dest_target =
@@ -1167,7 +1183,7 @@
   copyDWIMDest dest_target dest_slice src src_slice
 
 -- | As 'copyDWIM', but implicitly 'DimFix'es the indexes.
-copyDWIMFix :: VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> ImpM lore op ()
+copyDWIMFix :: VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> ImpM lore r op ()
 copyDWIMFix dest dest_is src src_is =
   copyDWIM dest (map DimFix dest_is) src (map DimFix src_is)
 
@@ -1176,7 +1192,7 @@
 -- 'MemoryDestination',
 compileAlloc :: ExplicitMemorish lore =>
                 Pattern lore -> SubExp -> Space
-             -> ImpM lore op ()
+             -> ImpM lore r op ()
 compileAlloc (Pattern [] [mem]) e space = do
   e' <- Imp.bytes <$> toExp e
   allocator <- asks $ M.lookup space . envAllocCompilers
@@ -1194,13 +1210,13 @@
 
 --- Building blocks for constructing code.
 
-sFor' :: VName -> IntType -> Imp.Exp -> ImpM lore op () -> ImpM lore op ()
+sFor' :: VName -> IntType -> Imp.Exp -> ImpM lore r op () -> ImpM lore r op ()
 sFor' i it bound body = do
   addLoopVar i it
   body' <- collect body
   emit $ Imp.For i it bound body'
 
-sFor :: String -> Imp.Exp -> (Imp.Exp -> ImpM lore op ()) -> ImpM lore op ()
+sFor :: String -> Imp.Exp -> (Imp.Exp -> ImpM lore r op ()) -> ImpM lore r op ()
 sFor i bound body = do
   i' <- newVName i
   it <- case primExpType bound of
@@ -1210,59 +1226,59 @@
   body' <- collect $ body $ Imp.var i' $ IntType it
   emit $ Imp.For i' it bound body'
 
-sWhile :: Imp.Exp -> ImpM lore op () -> ImpM lore op ()
+sWhile :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op ()
 sWhile cond body = do
   body' <- collect body
   emit $ Imp.While cond body'
 
-sComment :: String -> ImpM lore op () -> ImpM lore op ()
+sComment :: String -> ImpM lore r op () -> ImpM lore r op ()
 sComment s code = do
   code' <- collect code
   emit $ Imp.Comment s code'
 
-sIf :: Imp.Exp -> ImpM lore op () -> ImpM lore op () -> ImpM lore op ()
+sIf :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -> ImpM lore r op ()
 sIf cond tbranch fbranch = do
   tbranch' <- collect tbranch
   fbranch' <- collect fbranch
   emit $ Imp.If cond tbranch' fbranch'
 
-sWhen :: Imp.Exp -> ImpM lore op () -> ImpM lore op ()
+sWhen :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op ()
 sWhen cond tbranch = sIf cond tbranch (return ())
 
-sUnless :: Imp.Exp -> ImpM lore op () -> ImpM lore op ()
+sUnless :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op ()
 sUnless cond = sIf cond (return ())
 
-sOp :: op -> ImpM lore op ()
+sOp :: op -> ImpM lore r op ()
 sOp = emit . Imp.Op
 
-sDeclareMem :: String -> Space -> ImpM lore op VName
+sDeclareMem :: String -> Space -> ImpM lore r op VName
 sDeclareMem name space = do
   name' <- newVName name
   emit $ Imp.DeclareMem name' space
   addVar name' $ MemVar Nothing $ MemEntry space
   return name'
 
-sAlloc_ :: VName -> Count Bytes Imp.Exp -> Space -> ImpM lore op ()
+sAlloc_ :: VName -> Count Bytes Imp.Exp -> Space -> ImpM lore r op ()
 sAlloc_ name' size' space = do
   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 Imp.Exp -> Space -> ImpM lore op VName
+sAlloc :: String -> Count Bytes Imp.Exp -> Space -> ImpM lore r op VName
 sAlloc name size space = do
   name' <- sDeclareMem name space
   sAlloc_ name' size space
   return name'
 
-sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore op VName
+sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op VName
 sArray name bt shape membind = do
   name' <- newVName name
   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 :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r 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
@@ -1271,12 +1287,12 @@
     ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm
 
 -- | Uses linear/iota index function.
-sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore op VName
+sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore r op VName
 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 :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore r op VName
 sStaticArray name space pt vs = do
   let num_elems = case vs of Imp.ArrayValues vs' -> length vs'
                              Imp.ArrayZeros n -> fromIntegral n
@@ -1286,21 +1302,21 @@
   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 ()
+sWrite :: VName -> [Imp.Exp] -> PrimExp Imp.ExpLeaf -> ImpM lore r op ()
 sWrite arr is v = do
   (mem, space, offset) <- fullyIndexArray arr is
   vol <- asks envVolatility
   emit $ Imp.Write mem offset (primExpType v) space vol v
 
-sUpdate :: VName -> Slice Imp.Exp -> SubExp -> ImpM lore op ()
+sUpdate :: VName -> Slice Imp.Exp -> SubExp -> ImpM lore r op ()
 sUpdate arr slice v = do
   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 ()
+          -> ([Imp.Exp] -> ImpM lore r op ())
+          -> ImpM lore r op ()
 sLoopNest = sLoopNest' [] . shapeDims
   where sLoopNest' is [] f = f $ reverse is
         sLoopNest' is (d:ds) f = do
@@ -1308,12 +1324,13 @@
           sFor "nest_i" d' $ \i -> sLoopNest' (i:is) ds f
 
 -- | ASsignment.
-(<--) :: VName -> Imp.Exp -> ImpM lore op ()
+(<--) :: VName -> Imp.Exp -> ImpM lore r op ()
 x <-- e = emit $ Imp.SetScalar x e
 infixl 3 <--
 
 -- | Constructing a non-entry point function.
-function :: [Imp.Param] -> [Imp.Param] -> ImpM lore op () -> ImpM lore op (Imp.Function op)
+function :: [Imp.Param] -> [Imp.Param] -> ImpM lore r op ()
+         -> ImpM lore r op (Imp.Function op)
 function outputs inputs m = do
   body <- collect $ do
     mapM_ addParam $ outputs ++ inputs
diff --git a/src/Futhark/CodeGen/ImpGen/CUDA.hs b/src/Futhark/CodeGen/ImpGen/CUDA.hs
--- a/src/Futhark/CodeGen/ImpGen/CUDA.hs
+++ b/src/Futhark/CodeGen/ImpGen/CUDA.hs
@@ -9,4 +9,4 @@
 import Futhark.MonadFreshNames
 
 compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m OpenCL.Program
-compileProg prog = kernelsToCUDA <$> ImpGenKernels.compileProg prog
+compileProg prog = kernelsToCUDA <$> ImpGenKernels.compileProgCUDA prog
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,12 +3,12 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE ConstraintKinds #-}
 module Futhark.CodeGen.ImpGen.Kernels
-  ( Futhark.CodeGen.ImpGen.Kernels.compileProg
+  ( compileProgOpenCL
+  , compileProgCUDA
   )
   where
 
 import Control.Monad.Except
-import Control.Monad.Reader
 import Data.Maybe
 import Data.List ()
 
@@ -19,7 +19,8 @@
 import Futhark.Representation.ExplicitMemory
 import qualified Futhark.CodeGen.ImpCode.Kernels as Imp
 import Futhark.CodeGen.ImpCode.Kernels (bytes)
-import Futhark.CodeGen.ImpGen
+import Futhark.CodeGen.ImpGen hiding (compileProg)
+import qualified Futhark.CodeGen.ImpGen
 import Futhark.CodeGen.ImpGen.Kernels.Base
 import Futhark.CodeGen.ImpGen.Kernels.SegMap
 import Futhark.CodeGen.ImpGen.Kernels.SegRed
@@ -30,7 +31,7 @@
 import Futhark.CodeGen.SetDefaultSpace
 import Futhark.Util.IntegralExp (quot, quotRoundingUp, IntegralExp)
 
-callKernelOperations :: Operations ExplicitMemory Imp.HostOp
+callKernelOperations :: Operations ExplicitMemory HostEnv Imp.HostOp
 callKernelOperations =
   Operations { opsExpCompiler = expCompiler
              , opsCopyCompiler = callKernelCopy
@@ -39,11 +40,29 @@
              , opsAllocCompilers = mempty
              }
 
-compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m Imp.Program
-compileProg prog =
+openclAtomics, cudaAtomics :: AtomicBinOp
+(openclAtomics, cudaAtomics) = (flip lookup opencl, flip lookup cuda)
+  where opencl = [ (Add Int32, Imp.AtomicAdd Int32)
+                 , (SMax Int32, Imp.AtomicSMax Int32)
+                 , (SMin Int32, Imp.AtomicSMin Int32)
+                 , (UMax Int32, Imp.AtomicUMax Int32)
+                 , (UMin Int32, Imp.AtomicUMin Int32)
+                 , (And Int32, Imp.AtomicAnd Int32)
+                 , (Or Int32, Imp.AtomicOr Int32)
+                 , (Xor Int32, Imp.AtomicXor Int32)
+                 ]
+        cuda = opencl ++ [(FAdd Float32, Imp.AtomicFAdd Float32)]
+
+compileProg :: MonadFreshNames m => HostEnv -> Prog ExplicitMemory -> m Imp.Program
+compileProg env prog =
   setDefaultSpace (Imp.Space "device") <$>
-  Futhark.CodeGen.ImpGen.compileProg callKernelOperations (Imp.Space "device") prog
+  Futhark.CodeGen.ImpGen.compileProg env callKernelOperations (Imp.Space "device") prog
 
+compileProgOpenCL, compileProgCUDA
+  :: MonadFreshNames m => Prog ExplicitMemory -> m Imp.Program
+compileProgOpenCL = compileProg $ HostEnv openclAtomics
+compileProgCUDA = compileProg $ HostEnv cudaAtomics
+
 opCompiler :: Pattern ExplicitMemory -> Op ExplicitMemory
            -> CallKernelGen ()
 
@@ -51,12 +70,12 @@
   compileAlloc dest e space
 
 opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSize key size_class))) = do
-  fname <- asks envFunction
+  fname <- askFunction
   sOp $ Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $
     sizeClassWithEntryPoint fname size_class
 
 opCompiler (Pattern _ [pe]) (Inner (SizeOp (CmpSizeLe key size_class x))) = do
-  fname <- asks envFunction
+  fname <- askFunction
   let size_class' = sizeClassWithEntryPoint fname size_class
   sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class'
     =<< toExp x
@@ -65,7 +84,7 @@
   sOp $ Imp.GetSizeMax (patElemName pe) size_class
 
 opCompiler (Pattern _ [pe]) (Inner (SizeOp (CalcNumGroups w64 max_num_groups_key group_size))) = do
-  fname <- asks envFunction
+  fname <- askFunction
   max_num_groups <- dPrim "max_num_groups" int32
   sOp $ Imp.GetSize max_num_groups (keyWithEntryPoint fname max_num_groups_key) $
     sizeClassWithEntryPoint fname SizeNumGroups
@@ -93,7 +112,7 @@
   compilerBugS $ "opCompiler: Invalid pattern\n  " ++
   pretty pat ++ "\nfor expression\n  " ++ pretty e
 
-sizeClassWithEntryPoint :: Name -> Imp.SizeClass -> Imp.SizeClass
+sizeClassWithEntryPoint :: Maybe Name -> Imp.SizeClass -> Imp.SizeClass
 sizeClassWithEntryPoint fname (Imp.SizeThreshold path) =
   Imp.SizeThreshold $ map f path
   where f (name, x) = (keyWithEntryPoint fname name, x)
@@ -111,7 +130,7 @@
 segOpCompiler pat segop =
   compilerBugS $ "segOpCompiler: unexpected " ++ pretty (segLevel segop) ++ " for rhs of pattern " ++ pretty pat
 
-expCompiler :: ExpCompiler ExplicitMemory Imp.HostOp
+expCompiler :: ExpCompiler ExplicitMemory HostEnv Imp.HostOp
 
 -- We generate a simple kernel for itoa and replicate.
 expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do
@@ -131,7 +150,7 @@
 expCompiler dest e =
   defCompileExp dest e
 
-callKernelCopy :: CopyCompiler ExplicitMemory Imp.HostOp
+callKernelCopy :: CopyCompiler ExplicitMemory HostEnv Imp.HostOp
 callKernelCopy bt
   destloc@(MemLocation destmem _ destIxFun)
   srcloc@(MemLocation srcmem srcshape srcIxFun)
@@ -163,10 +182,7 @@
 
 mapTransposeForType :: PrimType -> CallKernelGen Name
 mapTransposeForType bt = do
-  -- FIXME: 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
+  let fname = nameFromString $ "builtin#" <> mapTransposeName bt
 
   exists <- hasFunction fname
   unless exists $ emitFunction fname $ mapTransposeFunction bt
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
@@ -6,6 +6,8 @@
   , keyWithEntryPoint
   , CallKernelGen
   , InKernelGen
+  , HostEnv (..)
+  , KernelEnv (..)
   , computeThreadChunkSize
   , groupReduce
   , groupScan
@@ -23,6 +25,7 @@
   , groupCoverSpace
 
   , atomicUpdateLocking
+  , AtomicBinOp
   , Locking(..)
   , AtomicUpdate(..)
   , DoAtomicUpdate
@@ -30,7 +33,6 @@
   where
 
 import Control.Monad.Except
-import Control.Monad.Reader
 import Data.Maybe
 import qualified Data.Map.Strict as M
 import Data.List (elemIndex, find, nub, zip4)
@@ -47,9 +49,17 @@
 import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem)
 import Futhark.Util (chunks, maybeNth, mapAccumLM, takeLast, dropLast)
 
-type CallKernelGen = ImpM ExplicitMemory Imp.HostOp
-type InKernelGen = ImpM ExplicitMemory Imp.KernelOp
+newtype HostEnv = HostEnv
+  { hostAtomics :: AtomicBinOp }
 
+data KernelEnv = KernelEnv
+  { kernelAtomics :: AtomicBinOp
+  , kernelConstants :: KernelConstants
+  }
+
+type CallKernelGen = ImpM ExplicitMemory HostEnv Imp.HostOp
+type InKernelGen = ImpM ExplicitMemory KernelEnv Imp.KernelOp
+
 data KernelConstants = KernelConstants
                        { kernelGlobalThreadId :: Imp.Exp
                        , kernelLocalThreadId :: Imp.Exp
@@ -64,33 +74,32 @@
                        , kernelThreadActive :: Imp.Exp
                        }
 
-keyWithEntryPoint :: Name -> Name -> Name
+keyWithEntryPoint :: Maybe Name -> Name -> Name
 keyWithEntryPoint fname key =
-  nameFromString $ nameToString fname ++ "." ++ nameToString key
+  nameFromString $ maybe "" ((++".") . nameToString) fname ++ nameToString key
 
-allocLocal :: AllocCompiler ExplicitMemory Imp.KernelOp
+allocLocal :: AllocCompiler ExplicitMemory r Imp.KernelOp
 allocLocal mem size =
   sOp $ Imp.LocalAlloc mem size
 
-kernelAlloc :: KernelConstants
-            -> Pattern ExplicitMemory
+kernelAlloc :: Pattern ExplicitMemory
             -> SubExp -> Space
-            -> ImpM ExplicitMemory Imp.KernelOp ()
-kernelAlloc _ (Pattern _ [_]) _ ScalarSpace{} =
+            -> InKernelGen ()
+kernelAlloc (Pattern _ [_]) _ ScalarSpace{} =
   -- Handled by the declaration of the memory block, which is then
   -- translated to an actual scalar variable during C code generation.
   return ()
-kernelAlloc _ (Pattern _ [mem]) size (Space "local") = do
+kernelAlloc (Pattern _ [mem]) size (Space "local") = do
   size' <- toExp size
   allocLocal (patElemName mem) $ Imp.bytes size'
-kernelAlloc _ (Pattern _ [mem]) _ _ =
+kernelAlloc (Pattern _ [mem]) _ _ =
   compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel."
-kernelAlloc _ dest _ _ =
+kernelAlloc dest _ _ =
   error $ "Invalid target for in-kernel allocation: " ++ show dest
 
 splitSpace :: (ToExp w, ToExp i, ToExp elems_per_thread) =>
               Pattern ExplicitMemory -> SplitOrdering -> w -> i -> elems_per_thread
-           -> ImpM lore op ()
+           -> ImpM lore r op ()
 splitSpace (Pattern [] [size]) o w i elems_per_thread = do
   num_elements <- Imp.elements <$> toExp w
   i' <- toExp i
@@ -99,7 +108,7 @@
 splitSpace pat _ _ _ _ =
   error $ "Invalid target for splitSpace: " ++ pretty pat
 
-compileThreadExp :: ExpCompiler ExplicitMemory Imp.KernelOp
+compileThreadExp :: ExpCompiler ExplicitMemory KernelEnv Imp.KernelOp
 compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =
   forM_ (zip [0..] es) $ \(i,e) ->
   copyDWIMFix (patElemName dest) [fromIntegral (i::Int32)] e []
@@ -126,51 +135,52 @@
 -- | Assign iterations of a for-loop to threads in the workgroup.  The
 -- passed-in function is invoked with the (symbolic) iteration.  For
 -- multidimensional loops, use 'groupCoverSpace'.
-groupLoop :: KernelConstants -> Imp.Exp
+groupLoop :: Imp.Exp
           -> (Imp.Exp -> InKernelGen ()) -> InKernelGen ()
-groupLoop constants =
-  kernelLoop (kernelLocalThreadId constants) (kernelGroupSize constants)
+groupLoop n f = do
+  constants <- kernelConstants <$> askEnv
+  kernelLoop (kernelLocalThreadId constants) (kernelGroupSize constants) n f
 
 -- | Iterate collectively though a multidimensional space, such that
 -- all threads in the group participate.  The passed-in function is
 -- invoked with a (symbolic) point in the index space.
-groupCoverSpace :: KernelConstants -> [Imp.Exp]
+groupCoverSpace :: [Imp.Exp]
                 -> ([Imp.Exp] -> InKernelGen ()) -> InKernelGen ()
-groupCoverSpace constants ds f =
-  groupLoop constants (product ds) $ f . unflattenIndex ds
+groupCoverSpace ds f =
+  groupLoop (product ds) $ f . unflattenIndex ds
 
-groupCopy :: KernelConstants -> VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> InKernelGen ()
-groupCopy constants to to_is from from_is = do
+groupCopy :: VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> InKernelGen ()
+groupCopy to to_is from from_is = do
   ds <- mapM toExp . arrayDims =<< subExpType from
-  groupCoverSpace constants ds $ \is ->
+  groupCoverSpace ds $ \is ->
     copyDWIMFix to (to_is++ is) from (from_is ++ is)
 
-compileGroupExp :: KernelConstants -> ExpCompiler ExplicitMemory Imp.KernelOp
+compileGroupExp :: ExpCompiler ExplicitMemory KernelEnv Imp.KernelOp
 -- The static arrays stuff does not work inside kernels.
-compileGroupExp _ (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =
+compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) =
   forM_ (zip [0..] es) $ \(i,e) ->
   copyDWIMFix (patElemName dest) [fromIntegral (i::Int32)] e []
-compileGroupExp constants (Pattern _ [dest]) (BasicOp (Copy arr)) = do
-  groupCopy constants (patElemName dest) [] (Var arr) []
+compileGroupExp (Pattern _ [dest]) (BasicOp (Copy arr)) = do
+  groupCopy (patElemName dest) [] (Var arr) []
   sOp $ Imp.Barrier Imp.FenceLocal
-compileGroupExp constants (Pattern _ [dest]) (BasicOp (Manifest _ arr)) = do
-  groupCopy constants (patElemName dest) [] (Var arr) []
+compileGroupExp (Pattern _ [dest]) (BasicOp (Manifest _ arr)) = do
+  groupCopy (patElemName dest) [] (Var arr) []
   sOp $ Imp.Barrier Imp.FenceLocal
-compileGroupExp constants (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do
+compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do
   ds' <- mapM toExp $ shapeDims ds
-  groupCoverSpace constants ds' $ \is ->
+  groupCoverSpace ds' $ \is ->
     copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is)
   sOp $ Imp.Barrier Imp.FenceLocal
-compileGroupExp constants (Pattern _ [dest]) (BasicOp (Iota n e s _)) = do
+compileGroupExp (Pattern _ [dest]) (BasicOp (Iota n e s _)) = do
   n' <- toExp n
   e' <- toExp e
   s' <- toExp s
-  groupLoop constants n' $ \i' -> do
+  groupLoop n' $ \i' -> do
     x <- dPrimV "x" $ e' + i' * s'
     copyDWIMFix (patElemName dest) [i'] (Var x) []
   sOp $ Imp.Barrier Imp.FenceLocal
 
-compileGroupExp _ dest e =
+compileGroupExp dest e =
   defCompileExp dest e
 
 sanityCheckLevel :: SegLevel -> InKernelGen ()
@@ -178,29 +188,32 @@
 sanityCheckLevel SegGroup{} =
   error "compileGroupOp: unexpected group-level SegOp."
 
-compileGroupSpace :: KernelConstants -> SegLevel -> SegSpace -> InKernelGen ()
-compileGroupSpace constants lvl space = do
+compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen ()
+compileGroupSpace lvl space = do
   sanityCheckLevel lvl
 
   let (ltids, dims) = unzip $ unSegSpace space
   dims' <- mapM toExp dims
-  zipWithM_ dPrimV_ ltids $ unflattenIndex dims' $ kernelLocalThreadId constants
+  ltid <- kernelLocalThreadId . kernelConstants <$> askEnv
+  zipWithM_ dPrimV_ ltids $ unflattenIndex dims' ltid
 
-  dPrimV_ (segFlat space) $ kernelLocalThreadId constants
+  dPrimV_ (segFlat space) ltid
 
 -- Construct the necessary lock arrays for an intra-group histogram.
-prepareIntraGroupSegHist :: KernelConstants
-                           -> Count GroupSize SubExp
-                           -> [HistOp ExplicitMemory]
-                           -> InKernelGen [[Imp.Exp] -> InKernelGen ()]
-prepareIntraGroupSegHist constants group_size =
+prepareIntraGroupSegHist :: Count GroupSize SubExp
+                         -> [HistOp ExplicitMemory]
+                         -> InKernelGen [[Imp.Exp] -> InKernelGen ()]
+prepareIntraGroupSegHist group_size =
   fmap snd . mapAccumLM onOp Nothing
   where
     onOp l op = do
 
+      constants <- kernelConstants <$> askEnv
+      atomicBinOp <- kernelAtomics <$> askEnv
+
       let local_subhistos = histDest op
 
-      case (l, atomicUpdateLocking $ histOp op) of
+      case (l, atomicUpdateLocking atomicBinOp $ histOp op) of
         (_, AtomicPrim f) -> return (l, f (Space "local") local_subhistos)
         (_, AtomicCAS f) -> return (l, f (Space "local") local_subhistos)
         (Just l', AtomicLocking f) -> return (l, f l' (Space "local") local_subhistos)
@@ -220,31 +233,31 @@
             map (primExpFromSubExp int32) $ arrayDims locks_t
 
           sComment "All locks start out unlocked" $
-            groupCoverSpace constants [kernelGroupSize constants] $ \is ->
+            groupCoverSpace [kernelGroupSize constants] $ \is ->
             copyDWIMFix locks is (intConst Int32 0) []
 
           return (Just l', f l' (Space "local") local_subhistos)
 
-compileGroupOp :: KernelConstants -> OpCompiler ExplicitMemory Imp.KernelOp
+compileGroupOp :: OpCompiler ExplicitMemory KernelEnv Imp.KernelOp
 
-compileGroupOp constants pat (Alloc size space) =
-  kernelAlloc constants pat size space
+compileGroupOp pat (Alloc size space) =
+  kernelAlloc pat size space
 
-compileGroupOp _ pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) =
+compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) =
   splitSpace pat o w i elems_per_thread
 
-compileGroupOp constants pat (Inner (SegOp (SegMap lvl space _ body))) = do
-  void $ compileGroupSpace constants lvl space
+compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do
+  void $ compileGroupSpace lvl space
 
   sWhen (isActive $ unSegSpace space) $
     compileStms mempty (kernelBodyStms body) $
-    zipWithM_ (compileThreadResult space constants) (patternElements pat) $
+    zipWithM_ (compileThreadResult space) (patternElements pat) $
     kernelBodyResult body
 
   sOp $ Imp.ErrorSync Imp.FenceLocal
 
-compileGroupOp constants pat (Inner (SegOp (SegScan lvl space scan_op _ _ body))) = do
-  compileGroupSpace constants lvl space
+compileGroupOp pat (Inner (SegOp (SegScan lvl space scan_op _ _ body))) = do
+  compileGroupSpace lvl space
   let (ltids, dims) = unzip $ unSegSpace space
   dims' <- mapM toExp dims
 
@@ -259,11 +272,11 @@
 
   let segment_size = last dims'
       crossesSegment from to = (to-from) .>. (to `rem` segment_size)
-  groupScan constants (Just crossesSegment) (product dims') (product dims') scan_op $
+  groupScan (Just crossesSegment) (product dims') (product dims') scan_op $
     patternNames pat
 
-compileGroupOp constants pat (Inner (SegOp (SegRed lvl space ops _ body))) = do
-  compileGroupSpace constants lvl space
+compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do
+  compileGroupSpace lvl space
 
   let (ltids, dims) = unzip $ unSegSpace space
       (red_pes, map_pes) =
@@ -282,7 +295,7 @@
           splitAt (segRedResults ops) $ kernelBodyResult body
     forM_ (zip tmp_arrs red_res) $ \(dest, res) ->
       copyDWIMFix dest (map (`Imp.var` int32) ltids) (kernelResultSubExp res) []
-    zipWithM_ (compileThreadResult space constants) map_pes map_res
+    zipWithM_ (compileThreadResult space) map_pes map_res
 
   sOp $ Imp.ErrorSync Imp.FenceLocal
 
@@ -291,7 +304,7 @@
     -- handle directly with a group-level reduction.
     [dim'] -> do
       forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->
-        groupReduce constants dim' (segRedLambda op) tmps
+        groupReduce dim' (segRedLambda op) tmps
 
       sOp $ Imp.ErrorSync Imp.FenceLocal
 
@@ -306,7 +319,7 @@
           crossesSegment from to = (to-from) .>. (to `rem` segment_size)
 
       forM_ (zip ops tmps_for_ops) $ \(op, tmps) ->
-        groupScan constants (Just crossesSegment) (product dims') (product dims')
+        groupScan (Just crossesSegment) (product dims') (product dims')
         (segRedLambda op) tmps
 
       sOp $ Imp.ErrorSync Imp.FenceLocal
@@ -317,8 +330,8 @@
 
       sOp $ Imp.Barrier Imp.FenceLocal
 
-compileGroupOp constants pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do
-  compileGroupSpace constants lvl space
+compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do
+  compileGroupSpace lvl space
   let ltids = map fst $ unSegSpace space
 
   -- We don't need the red_pes, because it is guaranteed by our type
@@ -328,7 +341,7 @@
       (_red_pes, map_pes) =
         splitAt num_red_res $ patternElements pat
 
-  ops' <- prepareIntraGroupSegHist constants (segGroupSize lvl) ops
+  ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops
 
   -- Ensure that all locks have been initialised.
   sOp $ Imp.Barrier Imp.FenceLocal
@@ -337,7 +350,7 @@
     compileStms mempty (kernelBodyStms kbody) $ do
     let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody
         (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res
-    zipWithM_ (compileThreadResult space constants) map_pes map_res
+    zipWithM_ (compileThreadResult space) map_pes map_res
 
     let vs_per_op = chunks (map (length . histDest) ops) red_vs
 
@@ -359,15 +372,15 @@
 
   sOp $ Imp.ErrorSync Imp.FenceLocal
 
-compileGroupOp _ pat _ =
+compileGroupOp pat _ =
   compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat
 
-compileThreadOp :: KernelConstants -> OpCompiler ExplicitMemory Imp.KernelOp
-compileThreadOp constants pat (Alloc size space) =
-  kernelAlloc constants pat size space
-compileThreadOp _ pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) =
+compileThreadOp :: OpCompiler ExplicitMemory KernelEnv Imp.KernelOp
+compileThreadOp pat (Alloc size space) =
+  kernelAlloc pat size space
+compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) =
   splitSpace pat o w i elems_per_thread
-compileThreadOp _ pat _ =
+compileThreadOp pat _ =
   compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat
 
 -- | Locking strategy used for an atomic update.
@@ -388,28 +401,34 @@
 
 -- | A function for generating code for an atomic update.  Assumes
 -- that the bucket is in-bounds.
-type DoAtomicUpdate lore =
-  Space -> [VName] -> [Imp.Exp] -> ImpM lore Imp.KernelOp ()
+type DoAtomicUpdate lore r =
+  Space -> [VName] -> [Imp.Exp] -> ImpM lore r Imp.KernelOp ()
 
 -- | The mechanism that will be used for performing the atomic update.
 -- Approximates how efficient it will be.  Ordered from most to least
 -- efficient.
-data AtomicUpdate lore
-  = AtomicPrim (DoAtomicUpdate lore)
+data AtomicUpdate lore r
+  = AtomicPrim (DoAtomicUpdate lore r)
     -- ^ Supported directly by primitive.
-  | AtomicCAS (DoAtomicUpdate lore)
+  | AtomicCAS (DoAtomicUpdate lore r)
     -- ^ Can be done by efficient swaps.
-  | AtomicLocking (Locking -> DoAtomicUpdate lore)
+  | AtomicLocking (Locking -> DoAtomicUpdate lore r)
     -- ^ Requires explicit locking.
 
+-- | Is there an atomic 'BinOp' corresponding to this 'BinOp'?
+type AtomicBinOp =
+  BinOp ->
+  Maybe (VName -> VName -> Count Imp.Elements Imp.Exp -> Imp.Exp -> Imp.AtomicOp)
+
 -- | 'atomicUpdate', but where it is explicitly visible whether a
 -- locking strategy is necessary.
-atomicUpdateLocking :: ExplicitMemorish lore =>
-                       Lambda lore -> AtomicUpdate lore
+atomicUpdateLocking :: AtomicBinOp -> Lambda ExplicitMemory
+                    -> AtomicUpdate ExplicitMemory KernelEnv
 
-atomicUpdateLocking lam
+atomicUpdateLocking atomicBinOp lam
   | Just ops_and_ts <- splitOp lam,
-    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = AtomicPrim $ \space arrs bucket ->
+    all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts =
+    primOrCas ops_and_ts $ \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
@@ -429,12 +448,18 @@
 
   where opHasAtomicSupport space old arr' bucket' bop = do
           let atomic f = Imp.Atomic space . f old arr' bucket'
-          atomic <$> Imp.atomicBinOp bop
+          atomic <$> atomicBinOp bop
 
+        primOrCas ops
+          | all isPrim ops = AtomicPrim
+          | otherwise      = AtomicCAS
+
+        isPrim (op, _, _, _) = isJust $ atomicBinOp op
+
 -- If the operator functions purely on single 32-bit values, we can
 -- use an implementation based on CAS, no matter what the operator
 -- does.
-atomicUpdateLocking op
+atomicUpdateLocking _ op
   | [Prim t] <- lambdaReturnType op,
     [xp, _] <- lambdaParams op,
     primBitSize t == 32 = AtomicCAS $ \space [arr] bucket -> do
@@ -442,7 +467,7 @@
       atomicUpdateCAS space t arr old bucket (paramName xp) $
         compileBody' [xp] $ lambdaBody op
 
-atomicUpdateLocking op = AtomicLocking $ \locking space arrs bucket -> do
+atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do
   old <- dPrim "old" int32
   continue <- newVName "continue"
   dPrimVol_ continue Bool
@@ -455,13 +480,15 @@
   -- Critical section
   let try_acquire_lock =
         sOp $ Imp.Atomic space $
-        Imp.AtomicCmpXchg old locks' locks_offset (lockingIsUnlocked locking) (lockingToLock locking)
+        Imp.AtomicCmpXchg int32 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 space $
-        Imp.AtomicCmpXchg old locks' locks_offset (lockingToLock locking) (lockingToUnlock locking)
+        Imp.AtomicCmpXchg int32 old locks' locks_offset
+        (lockingToLock locking) (lockingToUnlock locking)
       break_loop = continue <-- false
 
   -- Preparing parameters. It is assumed that the caller has already
@@ -510,8 +537,8 @@
 atomicUpdateCAS :: Space -> PrimType
                 -> VName -> VName
                 -> [Imp.Exp] -> VName
-                -> ImpM lore Imp.KernelOp ()
-                -> ImpM lore Imp.KernelOp ()
+                -> InKernelGen ()
+                -> InKernelGen ()
 atomicUpdateCAS space t arr old bucket x do_op = do
   -- Code generation target:
   --
@@ -542,7 +569,7 @@
     do_op
     old_bits <- dPrim "old_bits" int32
     sOp $ Imp.Atomic space $
-      Imp.AtomicCmpXchg old_bits arr' bucket_offset
+      Imp.AtomicCmpXchg int32 old_bits arr' bucket_offset
       (toBits (Imp.var assumed t)) (toBits (Imp.var x t))
     old <-- fromBits (Imp.var old_bits int32)
     sWhen (toBits (Imp.var assumed t) .==. Imp.var old_bits int32)
@@ -590,9 +617,9 @@
                   | otherwise  -> return $ Just $ Imp.ScalarUse var bt
 
 isConstExp :: VTable ExplicitMemory -> Imp.Exp
-           -> ImpM lore op (Maybe Imp.KernelConstExp)
+           -> ImpM lore r op (Maybe Imp.KernelConstExp)
 isConstExp vtable size = do
-  fname <- asks envFunction
+  fname <- askFunction
   let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name
       onLeaf (Imp.SizeOf pt) _ = Just $ primByteSize pt
       onLeaf Imp.Index{} _ = Nothing
@@ -611,7 +638,7 @@
                        -> Imp.Count Imp.Elements Imp.Exp
                        -> Imp.Count Imp.Elements Imp.Exp
                        -> VName
-                       -> ImpM lore op ()
+                       -> ImpM lore r op ()
 computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do
   stride' <- toExp stride
   chunk_var <--
@@ -687,32 +714,50 @@
 -- kernel).
 makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a
 makeAllMemoryGlobal =
-  local (\env -> env { envDefaultSpace = Imp.Space "global" }) .
-  localVTable (M.map globalMemory)
+  localDefaultSpace (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
 
-groupReduce :: ExplicitMemorish lore =>
-               KernelConstants
-            -> Imp.Exp
-            -> Lambda lore
+groupReduce :: Imp.Exp
+            -> Lambda ExplicitMemory
             -> [VName]
-            -> ImpM lore Imp.KernelOp ()
-groupReduce constants w lam arrs = do
+            -> InKernelGen ()
+groupReduce w lam arrs = do
   offset <- dPrim "offset" int32
-  groupReduceWithOffset constants offset w lam arrs
+  groupReduceWithOffset offset w lam arrs
 
-groupReduceWithOffset :: ExplicitMemorish lore =>
-                         KernelConstants
-                      -> VName
+groupReduceWithOffset :: VName
                       -> Imp.Exp
-                      -> Lambda lore
+                      -> Lambda ExplicitMemory
                       -> [VName]
-                      -> ImpM lore Imp.KernelOp ()
-groupReduceWithOffset constants offset w lam arrs = do
+                      -> InKernelGen ()
+groupReduceWithOffset offset w lam arrs = do
+  constants <- kernelConstants <$> askEnv
+
+  let local_tid = kernelLocalThreadId constants
+      global_tid = kernelGlobalThreadId constants
+
+      barrier
+        | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal
+        | otherwise                           = sOp $ Imp.Barrier Imp.FenceGlobal
+
+      readReduceArgument param arr
+        | Prim _ <- paramType param = do
+            let i = local_tid + Imp.vi32 offset
+            copyDWIMFix (paramName param) [] (Var arr) [i]
+        | otherwise = do
+            let i = global_tid + Imp.vi32 offset
+            copyDWIMFix (paramName param) [] (Var arr) [i]
+
+      writeReduceOpResult param arr
+        | Prim _ <- paramType param =
+            copyDWIMFix arr [local_tid] (Var $ paramName param) []
+        | otherwise =
+            return ()
+
   let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam
 
   skip_waves <- dPrim "skip_waves" int32
@@ -769,35 +814,15 @@
 
   in_wave_reductions
   cross_wave_reductions
-  where local_tid = kernelLocalThreadId constants
-        global_tid = kernelGlobalThreadId constants
 
-        barrier
-          | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal
-          | otherwise                           = sOp $ Imp.Barrier Imp.FenceGlobal
-
-        readReduceArgument param arr
-          | Prim _ <- paramType param = do
-              let i = local_tid + Imp.vi32 offset
-              copyDWIMFix (paramName param) [] (Var arr) [i]
-          | otherwise = do
-              let i = global_tid + Imp.vi32 offset
-              copyDWIMFix (paramName param) [] (Var arr) [i]
-
-        writeReduceOpResult param arr
-          | Prim _ <- paramType param =
-              copyDWIMFix arr [local_tid] (Var $ paramName param) []
-          | otherwise =
-              return ()
-
-groupScan :: KernelConstants
-          -> Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp)
+groupScan :: Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp)
           -> Imp.Exp
           -> Imp.Exp
           -> Lambda ExplicitMemory
           -> [VName]
-          -> ImpM ExplicitMemory Imp.KernelOp ()
-groupScan constants seg_flag arrs_full_size w lam arrs = do
+          -> InKernelGen ()
+groupScan seg_flag arrs_full_size w lam arrs = do
+  constants <- kernelConstants <$> askEnv
   renamed_lam <- renameLambda lam
 
   let ltid = kernelLocalThreadId constants
@@ -1012,7 +1037,7 @@
 computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp)
 computeMapKernelGroups kernel_size = do
   group_size <- dPrim "group_size" int32
-  fname <- asks envFunction
+  fname <- askFunction
   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
@@ -1050,14 +1075,15 @@
 -- of memory expansion should be proportional to the number of
 -- *physical* threads (hardware parallelism), not the amount of
 -- application parallelism.
-virtualiseGroups :: KernelConstants
-                 -> SegVirt
+virtualiseGroups :: SegVirt
                  -> Imp.Exp
                  -> (VName -> InKernelGen ())
                  -> InKernelGen ()
-virtualiseGroups constants SegNoVirt _ m =
-  m $ kernelGroupIdVar constants
-virtualiseGroups constants SegVirt required_groups m = do
+virtualiseGroups SegNoVirt _ m = do
+  gid <- kernelGroupIdVar . kernelConstants <$> askEnv
+  m gid
+virtualiseGroups SegVirt required_groups m = do
+  constants <- kernelConstants <$> askEnv
   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`
@@ -1072,25 +1098,26 @@
 sKernelThread :: String
               -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp
               -> VName
-              -> (KernelConstants -> InKernelGen ())
+              -> InKernelGen ()
               -> CallKernelGen ()
 sKernelThread = sKernel threadOperations kernelGlobalThreadId
 
 sKernelGroup :: String
              -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp
              -> VName
-             -> (KernelConstants -> InKernelGen ())
+             -> InKernelGen ()
              -> CallKernelGen ()
 sKernelGroup = sKernel groupOperations kernelGroupId
 
 sKernelFailureTolerant :: Bool
-                       -> Operations ExplicitMemory Imp.KernelOp
+                       -> Operations ExplicitMemory KernelEnv Imp.KernelOp
                        -> KernelConstants
                        -> Name
-                       -> ImpM ExplicitMemory Imp.KernelOp a
+                       -> InKernelGen ()
                        -> CallKernelGen ()
 sKernelFailureTolerant tol ops constants name m = do
-  body <- makeAllMemoryGlobal $ subImpM_ ops m
+  HostEnv atomics <- askEnv
+  body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants) ops m
   uses <- computeKernelUses body mempty
   emit $ Imp.Op $ Imp.CallKernel Imp.Kernel
     { Imp.kernelBody = body
@@ -1101,23 +1128,23 @@
     , Imp.kernelFailureTolerant = tol
     }
 
-sKernel :: (KernelConstants -> Operations ExplicitMemory Imp.KernelOp)
+sKernel :: Operations ExplicitMemory KernelEnv Imp.KernelOp
         -> (KernelConstants -> Imp.Exp)
         -> String
         -> Count NumGroups Imp.Exp
         -> Count GroupSize Imp.Exp
         -> VName
-        -> (KernelConstants -> ImpM ExplicitMemory Imp.KernelOp a)
-        -> ImpM ExplicitMemory Imp.HostOp ()
+        -> InKernelGen ()
+        -> CallKernelGen ()
 sKernel ops flatf name num_groups group_size v f = do
   (constants, set_constants) <- kernelInitialisationSimple num_groups group_size
   let name' = nameFromString $ name ++ "_" ++ show (baseTag v)
-  sKernelFailureTolerant False (ops constants) constants name' $ do
+  sKernelFailureTolerant False ops constants name' $ do
     set_constants
     dPrimV_ v $ flatf constants
-    f constants
+    f
 
-copyInGroup :: CopyCompiler ExplicitMemory Imp.KernelOp
+copyInGroup :: CopyCompiler ExplicitMemory KernelEnv Imp.KernelOp
 copyInGroup pt destloc srcloc = do
   dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc)
   src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc)
@@ -1129,20 +1156,19 @@
   where isScalarMem ScalarSpace{} = True
         isScalarMem _ = False
 
-threadOperations, groupOperations :: KernelConstants
-                                  -> Operations ExplicitMemory Imp.KernelOp
-threadOperations constants =
-  (defaultOperations $ compileThreadOp constants)
+threadOperations, groupOperations :: Operations ExplicitMemory KernelEnv Imp.KernelOp
+threadOperations =
+  (defaultOperations compileThreadOp)
   { opsCopyCompiler = copyElementWise
   , opsExpCompiler = compileThreadExp
   , opsStmsCompiler = \_ -> defCompileStms mempty
   , opsAllocCompilers =
       M.fromList [ (Space "local", allocLocal) ]
   }
-groupOperations constants =
-  (defaultOperations $ compileGroupOp constants)
+groupOperations =
+  (defaultOperations compileGroupOp)
   { opsCopyCompiler = copyInGroup
-  , opsExpCompiler = compileGroupExp constants
+  , opsExpCompiler = compileGroupExp
   , opsStmsCompiler = \_ -> defCompileStms mempty
   , opsAllocCompilers =
       M.fromList [ (Space "local", allocLocal) ]
@@ -1162,7 +1188,7 @@
       name = nameFromString $ "replicate_" ++
              show (baseTag $ kernelGlobalThreadIdVar constants)
 
-  sKernelFailureTolerant True (threadOperations constants) constants name $ do
+  sKernelFailureTolerant True threadOperations constants name $ do
     set_constants
     sWhen (kernelThreadActive constants) $
       copyDWIMFix arr is' se $ drop (length ds) is'
@@ -1187,10 +1213,7 @@
 
 replicateForType :: PrimType -> CallKernelGen Name
 replicateForType bt = do
-  -- FIXME: The leading underscore is to avoid clashes with a
-  -- programmer-defined function of the same name (this is a bad
-  -- solution...).
-  let fname = nameFromString $ "_" <> replicateName bt
+  let fname = nameFromString $ "builtin#" <> replicateName bt
 
   exists <- hasFunction fname
   unless exists $ emitFunction fname =<< replicateFunction bt
@@ -1232,7 +1255,7 @@
   let name = nameFromString $ "iota_" ++
              show (baseTag $ kernelGlobalThreadIdVar constants)
 
-  sKernelFailureTolerant True (threadOperations constants) constants name $ do
+  sKernelFailureTolerant True threadOperations constants name $ do
     set_constants
     let gtid = kernelGlobalThreadId constants
     sWhen (kernelThreadActive constants) $ do
@@ -1260,7 +1283,7 @@
   let name = nameFromString $ "copy_" ++
              show (baseTag $ kernelGlobalThreadIdVar constants)
 
-  sKernelFailureTolerant True (threadOperations constants) constants name $ do
+  sKernelFailureTolerant True threadOperations constants name $ do
     set_constants
 
     let gtid = kernelGlobalThreadId constants
@@ -1275,12 +1298,16 @@
       Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile
 
 compileGroupResult :: SegSpace
-                   -> KernelConstants -> PatElem ExplicitMemory -> KernelResult
+                   -> PatElem ExplicitMemory -> KernelResult
                    -> InKernelGen ()
 
-compileGroupResult _ constants pe (TileReturns [(w,per_group_elems)] what) = do
+compileGroupResult _ pe (TileReturns [(w,per_group_elems)] what) = do
   n <- toExp . arraySize 0 =<< lookupType what
 
+  constants <- kernelConstants <$> askEnv
+  let ltid = kernelLocalThreadId constants
+      offset = toExp' int32 per_group_elems * kernelGroupId constants
+
   -- Avoid loop for the common case where each thread is statically
   -- known to write at most one element.
   if toExp' int32 per_group_elems == kernelGroupSize constants
@@ -1291,10 +1318,9 @@
       j <- fmap Imp.vi32 $ dPrimV "j" $
            kernelGroupSize constants * i + ltid
       sWhen (j .<. n) $ copyDWIMFix (patElemName pe) [j + offset] (Var what) [j]
-  where ltid = kernelLocalThreadId constants
-        offset = toExp' int32 per_group_elems * kernelGroupId constants
 
-compileGroupResult space constants pe (TileReturns dims what) = do
+compileGroupResult space pe (TileReturns dims what) = do
+  constants <- kernelConstants <$> askEnv
   let gids = map fst $ unSegSpace space
       out_tile_sizes = map (toExp' int32 . snd) dims
       local_is = unflattenIndex out_tile_sizes $ kernelLocalThreadId constants
@@ -1304,7 +1330,8 @@
   sWhen (isActive $ zip is_for_thread $ map fst dims) $
     copyDWIMFix (patElemName pe) (map Imp.vi32 is_for_thread) (Var what) local_is
 
-compileGroupResult space constants pe (Returns _ what) = do
+compileGroupResult space pe (Returns _ what) = do
+  constants <- kernelConstants <$> askEnv
   in_local_memory <- arrayInLocalMemory what
   let gids = map (Imp.vi32 . fst) $ unSegSpace space
 
@@ -1316,33 +1343,35 @@
       -- 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).
-      groupCopy constants (patElemName pe) gids what []
+      groupCopy (patElemName pe) gids what []
 
-compileGroupResult _ _ _ WriteReturns{} =
+compileGroupResult _ _ WriteReturns{} =
   compilerLimitationS "compileGroupResult: WriteReturns not handled yet."
 
-compileGroupResult _ _ _ ConcatReturns{} =
+compileGroupResult _ _ ConcatReturns{} =
   compilerLimitationS "compileGroupResult: ConcatReturns not handled yet."
 
 compileThreadResult :: SegSpace
-                    -> KernelConstants -> PatElem ExplicitMemory -> KernelResult
+                    -> PatElem ExplicitMemory -> KernelResult
                     -> InKernelGen ()
 
-compileThreadResult space _ pe (Returns _ what) = do
+compileThreadResult space pe (Returns _ what) = do
   let is = map (Imp.vi32 . fst) $ unSegSpace space
   copyDWIMFix (patElemName pe) is what []
 
-compileThreadResult _ constants pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do
+compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do
+  constants <- kernelConstants <$> askEnv
+  let offset = toExp' int32 per_thread_elems * kernelGlobalThreadId constants
   n <- toExp' int32 . arraySize 0 <$> lookupType what
   copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) []
-  where offset = toExp' int32 per_thread_elems * kernelGlobalThreadId constants
 
-compileThreadResult _ constants pe (ConcatReturns (SplitStrided stride) _ _ what) = do
+compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do
+  offset <- kernelGlobalThreadId . kernelConstants <$> askEnv
   n <- toExp' int32 . arraySize 0 <$> lookupType what
   copyDWIM (patElemName pe) [DimSlice offset n $ toExp' int32 stride] (Var what) []
-  where offset = kernelGlobalThreadId constants
 
-compileThreadResult _ constants pe (WriteReturns rws _arr dests) = do
+compileThreadResult _ pe (WriteReturns rws _arr dests) = do
+  constants <- kernelConstants <$> askEnv
   rws' <- mapM toExp rws
   forM_ dests $ \(is, e) -> do
     is' <- mapM toExp is
@@ -1351,7 +1380,7 @@
                 zipWith condInBounds is' rws'
     sWhen write $ copyDWIMFix (patElemName pe) (map (toExp' int32) is) e []
 
-compileThreadResult _ _ _ TileReturns{} =
+compileThreadResult _ _ TileReturns{} =
   compilerBugS "compileThreadResult: TileReturns unhandled."
 
 arrayInLocalMemory :: SubExp -> InKernelGen Bool
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs
@@ -42,7 +42,6 @@
   where
 
 import Control.Monad.Except
-import Control.Monad.Reader
 import Data.Maybe
 import Data.List (foldl', genericLength, zip4, zip6)
 
@@ -71,7 +70,7 @@
                    { slugOp :: HistOp ExplicitMemory
                    , slugNumSubhistos :: VName
                    , slugSubhistos :: [SubhistosInfo]
-                   , slugAtomicUpdate :: AtomicUpdate ExplicitMemory
+                   , slugAtomicUpdate :: AtomicUpdate ExplicitMemory KernelEnv
                    }
 
 histoSpaceUsage :: HistOp ExplicitMemory
@@ -139,10 +138,12 @@
   let h = histoSpaceUsage op
       segmented_h = h * product (map (Imp.bytes . toExp' int32) $ init $ segSpaceDims space)
 
+  atomics <- hostAtomics <$> askEnv
+
   return (h,
           segmented_h,
           SegHistSlug op num_subhistos subhisto_infos $
-          atomicUpdateLocking $ histOp op)
+          atomicUpdateLocking atomics $ histOp op)
 
 prepareAtomicUpdateGlobal :: Maybe Locking -> [VName] -> SegHistSlug
                           -> CallKernelGen (Maybe Locking,
@@ -213,11 +214,12 @@
 
   -- Querying L2 cache size is not reliable.  Instead we provide a
   -- tunable knob with a hopefully sane default.
-  let hist_L2_def = 5632 * 1024
+  let hist_L2_def = 4 * 1024 * 1024
   hist_L2 <- dPrim "L2_size" int32
-  entry <- asks envFunction
+  entry <- askFunction
   -- Equivalent to F_L2*L2 in paper.
-  sOp $ Imp.GetSize hist_L2 (entry <> "." <> nameFromString (pretty hist_L2)) $
+  sOp $ Imp.GetSize hist_L2
+    (keyWithEntryPoint entry $ nameFromString (pretty hist_L2)) $
     Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def
 
   let hist_L2_ln_sz = 16*4 -- L2 cache line size approximation
@@ -302,8 +304,10 @@
 
       -- Number of subhistograms per result histogram.
       hist_M <- dPrimVE "hist_M" $
-                Imp.BinOpExp (SMax Int32) hist_M_min $
-                t64 $ r64 hist_T / hist_C
+        case slugAtomicUpdate slug of
+          AtomicPrim{} -> 1
+          _ -> Imp.BinOpExp (SMax Int32) hist_M_min $
+               t64 $ r64 hist_T / hist_C
 
       emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just hist_k_max
       emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just hist_M
@@ -345,9 +349,9 @@
                      -> SegSpace
                      -> [SegHistSlug]
                      -> KernelBody ExplicitMemory
-                     -> [[Imp.Exp] -> ImpM ExplicitMemory Imp.KernelOp ()]
+                     -> [[Imp.Exp] -> InKernelGen ()]
                      -> Imp.Exp -> Imp.Exp
-                     -> ImpM ExplicitMemory Imp.HostOp ()
+                     -> CallKernelGen ()
 histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do
 
   let (space_is, space_sizes) = unzip $ unSegSpace space
@@ -358,7 +362,9 @@
     w' <- toExp w
     dPrimVE "hist_H_chk" $ w' `quotRoundingUp` hist_S
 
-  sKernelThread "seghist_global" num_groups group_size (segFlat space) $ \constants -> do
+  sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
+
     -- Compute subhistogram index for each thread, per histogram.
     subhisto_inds <- forM slugs $ \slug ->
       dPrimVE "subhisto_ind" $
@@ -442,7 +448,7 @@
     histograms hist_S chk_i
 
 type InitLocalHistograms = [([VName],
-                              (KernelConstants, SubExp) ->
+                              SubExp ->
                               InKernelGen ([VName],
                                             [Imp.Exp] -> InKernelGen ()))]
 
@@ -466,7 +472,7 @@
         case do_op of
           AtomicPrim f -> return $ const $ return f
           AtomicCAS f -> return $ const $ return f
-          AtomicLocking f -> return $ \(constants, hist_H_chk) -> do
+          AtomicLocking f -> return $ \hist_H_chk -> do
             let lock_shape =
                   Shape $ Var num_subhistos_per_group :
                   shapeDims (histShape op) ++
@@ -477,14 +483,14 @@
             locks <- sAllocArray "locks" int32 lock_shape $ Space "local"
 
             sComment "All locks start out unlocked" $
-              groupCoverSpace constants dims $ \is ->
+              groupCoverSpace dims $ \is ->
               copyDWIMFix locks is (intConst Int32 0) []
 
             return $ f $ Locking locks 0 1 0 id
 
       -- Initialise local-memory sub-histograms.  These are
       -- represented as two-dimensional arrays.
-      let init_local_subhistos (constants, hist_H_chk) = do
+      let init_local_subhistos hist_H_chk = do
             local_subhistos <-
               forM (histType op) $ \t -> do
                 let sub_local_shape =
@@ -493,7 +499,7 @@
                 sAllocArray "subhistogram_local"
                   (elemType t) sub_local_shape (Space "local")
 
-            do_op' <- mk_op (constants, hist_H_chk)
+            do_op' <- mk_op hist_H_chk
 
             return (local_subhistos, do_op' (Space "local") local_subhistos)
 
@@ -529,9 +535,11 @@
     w' <- toExp w
     dPrimV "hist_H_chk" $ w' `quotRoundingUp` hist_S
 
-  sKernelThread "seghist_local" num_groups group_size (segFlat space) $ \constants ->
-    virtualiseGroups constants SegVirt (unCount groups_per_segment * num_segments) $ \group_id_var -> do
+  sKernelThread "seghist_local" num_groups group_size (segFlat space) $
+    virtualiseGroups SegVirt (unCount groups_per_segment * num_segments) $ \group_id_var -> do
 
+    constants <- kernelConstants <$> askEnv
+
     let group_id = Imp.vi32 group_id_var
 
     flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` unCount groups_per_segment
@@ -549,7 +557,7 @@
 
     histograms <- forM (zip init_histograms hist_H_chks) $
                   \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do
-      (local_subhistos, do_op) <- init_local_subhistos (constants, Var hist_H_chk)
+      (local_subhistos, do_op) <- init_local_subhistos $ Var hist_H_chk
       return (zip glob_subhistos local_subhistos, hist_H_chk, do_op)
 
     -- Find index of local subhistograms updated by this thread.  We
@@ -724,9 +732,9 @@
 slugMaxLocalMemPasses :: SegHistSlug -> Int
 slugMaxLocalMemPasses slug =
   case slugAtomicUpdate slug of
-    AtomicPrim _ -> 2
+    AtomicPrim _ -> 3
     AtomicCAS _  -> 4
-    AtomicLocking _ -> 5
+    AtomicLocking _ -> 6
 
 localMemoryCase :: [PatElem ExplicitMemory]
                 -> Imp.Exp
@@ -734,8 +742,8 @@
                 -> Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp
                 -> [SegHistSlug]
                 -> KernelBody ExplicitMemory
-                -> CallKernelGen (Imp.Exp, ImpM ExplicitMemory Imp.HostOp ())
-localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody = do
+                -> CallKernelGen (Imp.Exp, CallKernelGen ())
+localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N _ slugs kbody = do
   let space_sizes = segSpaceDims space
       segment_dims = init space_sizes
       segmented = not $ null segment_dims
@@ -755,7 +763,6 @@
       t64 = ConvOpExp (FPToSI Float64 Int32)
       i32_to_i64 = ConvOpExp (SExt Int32 Int64)
       i64_to_i32 = ConvOpExp (SExt Int64 Int32)
-      f64ceil x = t64 $ FunExp "round64" [x] $ FloatType Float64
 
   -- M approximation.
   hist_m' <- dPrimVE "hist_m_prime" $
@@ -764,50 +771,12 @@
                   (hist_N `quotRoundingUp` unCount num_groups'))
              / r64 hist_H
 
-  hist_m <- dPrimVE "hist_m" $
-            Imp.BinOpExp (FMax Float64) (r64 1) hist_m'
-
-  -- FIXME: query the lockstep width at runtime.
-  hist_W <- dPrimVE "hist_W" 32
-
-  hist_RFC <- dPrimVE "hist_RFC" $
-              Imp.BinOpExp (FMin Float64) (r64 hist_RF) $
-              r64 hist_W * Imp.BinOpExp (FPow Float64) (r64 hist_RF / r64 hist_W) (1.0 / 3.0)
-
   let hist_B = unCount group_size'
-  hist_f' <- dPrimVE "hist_f_prime" $
-             (r64 hist_B * hist_RFC) /
-             (hist_m * hist_m * r64 hist_H)
 
-  let casOp AtomicCAS{} = True
-      casOp _ = False
-      lockOp AtomicLocking{} = True
-      lockOp _ = False
-
-      hwdCase =
-        dPrimVE "hist_M0" $
-        Imp.BinOpExp (SMax Int32) 1 $
-        Imp.BinOpExp (SMin Int32) (t64 hist_m') hist_B
-
-      casCase = do
-        hist_f_cas <- dPrimVE "hist_f_cas" $ Imp.BinOpExp (SMax Int32) 1 $ f64ceil hist_f'
-        dPrimVE "hist_M0" $
-          Imp.BinOpExp (SMax Int32) 1 $
-          Imp.BinOpExp (SMin Int32) (t64 $ hist_m * r64 hist_f_cas) hist_B
-
-      lockCase = do
-        hist_f_cas <- dPrimVE "hist_f_lock" $ Imp.BinOpExp (SMax Int32) 1 $ f64ceil hist_f'
-        dPrimVE "hist_M0" $
-          Imp.BinOpExp (SMax Int32) 1 $
-          Imp.BinOpExp (SMin Int32) (t64 $ hist_m' * r64 hist_f_cas) hist_B
-
   -- M in the paper, but not adjusted for asymptotic efficiency.
-  hist_M0 <-
-    if any (lockOp . slugAtomicUpdate) slugs
-    then lockCase
-    else if any (casOp . slugAtomicUpdate) slugs
-    then casCase
-    else hwdCase
+  hist_M0 <- dPrimVE "hist_M0" $
+             Imp.BinOpExp (SMax Int32) 1 $
+             Imp.BinOpExp (SMin Int32) (t64 hist_m') hist_B
 
   -- Minimal sequential chunking factor.
   let q_small = 2
@@ -1006,7 +975,7 @@
               [(subhistogram_id, Var num_histos)]
 
         let segred_op = SegRedOp Commutative (histOp op) (histNeutral op) mempty
-        compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \_ red_cont ->
+        compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \red_cont ->
           red_cont $ flip map subhistos $ \subhisto ->
             (Var subhisto, map Imp.vi32 $
               map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids)
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs
@@ -32,25 +32,25 @@
     SegThread{} -> do
       emit $ Imp.DebugPrint "\n# SegMap" Nothing
       let virt_num_groups = product dims' `quotRoundingUp` unCount group_size'
-      sKernelThread "segmap" num_groups' group_size' (segFlat space) $ \constants ->
-        virtualiseGroups constants (segVirt lvl) virt_num_groups $ \group_id -> do
-        let global_tid = Imp.vi32 group_id * unCount group_size' +
-                         kernelLocalThreadId constants
+      sKernelThread "segmap" num_groups' group_size' (segFlat space) $
+        virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do
+        local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv
+        let global_tid = Imp.vi32 group_id * unCount group_size' + local_tid
 
         zipWithM_ dPrimV_ is $ unflattenIndex dims' global_tid
 
         sWhen (isActive $ unSegSpace space) $
           compileStms mempty (kernelBodyStms kbody) $
-          zipWithM_ (compileThreadResult space constants) (patternElements pat) $
+          zipWithM_ (compileThreadResult space) (patternElements pat) $
           kernelBodyResult kbody
 
     SegGroup{} ->
-      sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ \constants -> do
+      sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do
       let virt_num_groups = product dims'
-      virtualiseGroups constants (segVirt lvl) virt_num_groups $ \group_id -> do
+      virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do
 
         zipWithM_ dPrimV_ is $ unflattenIndex dims' $ Imp.vi32 group_id
 
         compileStms mempty (kernelBodyStms kbody) $
-          zipWithM_ (compileGroupResult space constants) (patternElements pat) $
+          zipWithM_ (compileGroupResult space) (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
@@ -69,7 +69,7 @@
 maxNumOps :: Int32
 maxNumOps = 10
 
-type DoSegBody = (KernelConstants -> ([(SubExp, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen ())
+type DoSegBody = ([(SubExp, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen ()
 
 -- | Compile 'SegRed' instance to host-level code with calls to
 -- various kernels.
@@ -79,13 +79,13 @@
               -> KernelBody ExplicitMemory
               -> CallKernelGen ()
 compileSegRed pat lvl space reds body =
-  compileSegRed' pat lvl space reds $ \constants red_cont ->
+  compileSegRed' pat lvl space reds $ \red_cont ->
   compileStms mempty (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_ (compileThreadResult space constants) map_arrs map_res
+    zipWithM_ (compileThreadResult space) map_arrs map_res
 
   red_cont $ zip (map kernelResultSubExp red_res) $ repeat []
 
@@ -178,7 +178,8 @@
 
   emit $ Imp.DebugPrint "\n# SegRed" Nothing
 
-  sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ \constants -> do
+  sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
     sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
     reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds
 
@@ -237,14 +238,14 @@
   emit $ Imp.DebugPrint "segments_per_group" $ Just segments_per_group
   emit $ Imp.DebugPrint "required_groups" $ Just required_groups
 
-  sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ \constants -> do
-
+  sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
     reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds
 
     -- We probably do not have enough actual workgroups to cover the
     -- entire iteration space.  Some groups thus have to perform double
     -- duty; we put an outer loop to accomplish this.
-    virtualiseGroups constants SegVirt required_groups $ \group_id_var' -> do
+    virtualiseGroups SegVirt 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
@@ -262,7 +263,7 @@
             copyDWIMFix arr [ltid] ne []
 
           in_bounds =
-            body constants $ \red_res ->
+            body $ \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)) ->
@@ -279,7 +280,7 @@
       sWhen (segment_size .>. 0) $
         sComment "perform segmented scan to imitate reduction" $
         forM_ (zip reds reds_arrs) $ \(SegRedOp _ red_op _ _, red_arrs) ->
-        groupScan constants (Just crossesSegment) (Imp.vi32 num_threads)
+        groupScan (Just crossesSegment) (Imp.vi32 num_threads)
         (segment_size*segments_per_group) red_op red_arrs
 
       sOp $ Imp.Barrier Imp.FenceLocal
@@ -350,15 +351,15 @@
     sStaticArray "counter" (Space "device") int32 $
     Imp.ArrayZeros num_counters
 
-  sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ \constants -> do
-
+  sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
     reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds
     sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local"
 
     -- We probably do not have enough actual workgroups to cover the
     -- entire iteration space.  Some groups thus have to perform double
     -- duty; we put an outer loop to accomplish this.
-    virtualiseGroups constants SegVirt (Imp.vi32 virt_num_groups) $ \group_id_var -> do
+    virtualiseGroups SegVirt (Imp.vi32 virt_num_groups) $ \group_id_var -> do
       let segment_gtids = init gtids
           group_id = Imp.vi32 group_id_var
           flat_segment_id = group_id `quot` groups_per_segment
@@ -503,7 +504,7 @@
 
           sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier.
 
-          groupReduce constants (kernelGroupSize constants) slug_op_renamed (slugArrs slug)
+          groupReduce (kernelGroupSize constants) slug_op_renamed (slugArrs slug)
 
           sOp $ Imp.Barrier Imp.FenceLocal
 
@@ -535,7 +536,7 @@
              kernelGroupSize constants
 
     check_bounds $ sComment "apply map function" $
-      body constants $ \all_red_res -> do
+      body $ \all_red_res -> do
 
       let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res
 
@@ -620,7 +621,7 @@
     sOp $ Imp.MemFence Imp.FenceGlobal
     -- Increment the counter, thus stating that our result is
     -- available.
-    sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset 1
+    sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd Int32 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.
     sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1
@@ -637,7 +638,8 @@
     -- with an atomic to avoid warnings about write/write
     -- races in oclgrind.
     sWhen (local_tid .==. 0) $
-      sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd old_counter counter_mem counter_offset $
+      sOp $ Imp.Atomic DefaultSpace $
+      Imp.AtomicAdd Int32 old_counter counter_mem counter_offset $
       negate groups_per_segment
     sLoopNest (slugShape slug) $ \vec_is -> do
       comment "read in the per-group-results" $
@@ -656,7 +658,7 @@
       sOp $ Imp.Barrier Imp.FenceLocal
 
       sComment "reduce the per-group results" $ do
-        groupReduce constants group_size red_op_renamed red_arrs
+        groupReduce group_size red_op_renamed red_arrs
 
         sComment "and back to memory with the final result" $
           sWhen (local_tid .==. 0) $
diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs
--- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs
+++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs
@@ -77,7 +77,8 @@
             (to-from) .>. (to `rem` segment_size)
           _ -> Nothing
 
-  sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ \constants -> do
+  sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
     local_arrs <- makeLocalArrays group_size (Var num_threads) nes scan_op
 
     -- The variables from scan_op will be used for the carry and such
@@ -129,7 +130,7 @@
 
       sOp $ Imp.ErrorSync fence
 
-      groupScan constants crossesSegment'
+      groupScan crossesSegment'
         (Imp.vi32 num_threads)
         (kernelGroupSize constants) scan_op_renamed local_arrs
 
@@ -187,7 +188,8 @@
         Just $ \from to ->
           f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1)
 
-  sKernelThread  "scan_stage2" 1 group_size' (segFlat space) $ \constants -> do
+  sKernelThread  "scan_stage2" 1 group_size' (segFlat space) $ do
+    constants <- kernelConstants <$> askEnv
     local_arrs <- makeLocalArrays group_size (Var stage1_num_threads) nes scan_op
 
     flat_idx <- dPrimV "flat_idx" $
@@ -208,7 +210,7 @@
 
     barrier
 
-    groupScan constants crossesSegment'
+    groupScan crossesSegment'
       (Imp.vi32 stage1_num_threads) (kernelGroupSize constants) scan_op local_arrs
 
     sComment "threads in bounds write scanned carries" $
@@ -231,8 +233,10 @@
   required_groups <- dPrimVE "required_groups" $
                      product dims' `quotRoundingUp` unCount group_size'
 
-  sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $ \constants ->
-    virtualiseGroups constants SegVirt required_groups $ \virt_group_id -> do
+  sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $
+    virtualiseGroups SegVirt required_groups $ \virt_group_id -> do
+    constants <- kernelConstants <$> askEnv
+
     -- Compute our logical index.
     flat_idx <- dPrimVE "flat_idx" $
                 Imp.vi32 virt_group_id * unCount group_size' +
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
@@ -1,5 +1,6 @@
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TemplateHaskell #-}
 -- | This module defines a translation from imperative code with
 -- kernels to imperative code with OpenCL calls.
 module Futhark.CodeGen.ImpGen.Kernels.ToOpenCL
@@ -11,6 +12,7 @@
 import Control.Monad.State
 import Control.Monad.Identity
 import Control.Monad.Reader
+import Data.FileEmbed
 import Data.Maybe
 import qualified Data.Set as S
 import qualified Data.Map.Strict as M
@@ -27,6 +29,7 @@
 import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL
 import Futhark.MonadFreshNames
 import Futhark.Util (zEncodeString)
+import Futhark.Util.Pretty (prettyOneLine)
 
 kernelsToCUDA, kernelsToOpenCL :: ImpKernels.Program -> ImpOpenCL.Program
 kernelsToCUDA = translateKernels TargetCUDA
@@ -261,7 +264,7 @@
 constDef (ConstUse v e) = Just ([C.citem|$escstm:def|],
                                 [C.citem|$escstm:undef|])
   where e' = compilePrimExp e
-        def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ pretty e' ++ ")"
+        def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ prettyOneLine e' ++ ")"
         undef = "#undef " ++ pretty (C.toIdent v mempty)
 constDef _ = Nothing
 
@@ -272,6 +275,9 @@
           [[C.cedecl|$func:kernel_func|] |
            kernel_func <- kernels ]
 
+atomicsDefs :: String
+atomicsDefs = $(embedStringFile "rts/c/atomics.h")
+
 genOpenClPrelude :: S.Set PrimType -> [C.Definition]
 genOpenClPrelude ts =
   -- Clang-based OpenCL implementations need this for 'static' to work.
@@ -317,39 +323,18 @@
 |] ++
   cIntOps ++ cFloat32Ops ++ cFloat32Funs ++
   (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else [])
+  ++ [[C.cedecl|$esc:atomicsDefs|]]
   where uses_float64 = FloatType Float64 `S.member` ts
 
-
-cudaAtomicOps :: [C.Definition]
-cudaAtomicOps = (mkOp <$> opNames <*> types) ++ extraOps
-  where
-    mkOp (clName, cuName) t =
-      [C.cedecl|static inline $ty:t $id:clName(volatile $ty:t *p, $ty:t val) {
-                 return $id:cuName(($ty:t *)p, val);
-               }|]
-    types = [ [C.cty|int|]
-            , [C.cty|unsigned int|]
-            , [C.cty|unsigned long long|]
-            ]
-    opNames = [ ("atomic_add",  "atomicAdd")
-              , ("atomic_max",  "atomicMax")
-              , ("atomic_min",  "atomicMin")
-              , ("atomic_and",  "atomicAnd")
-              , ("atomic_or",   "atomicOr")
-              , ("atomic_xor",  "atomicXor")
-              , ("atomic_xchg", "atomicExch")
-              ]
-    extraOps =
-      [ [C.cedecl|static inline $ty:t atomic_cmpxchg(volatile $ty:t *p, $ty:t cmp, $ty:t val) {
-                  return atomicCAS(($ty:t *)p, cmp, val);
-                }|] | t <- types]
-
 genCUDAPrelude :: [C.Definition]
 genCUDAPrelude =
-  cudafy ++ cudaAtomicOps ++ ops
+  cudafy ++ ops
   where ops = cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops
-                ++ cFloat64Funs ++ cFloatConvOps
+              ++ cFloat64Funs ++ cFloatConvOps
+              ++ [[C.cedecl|$esc:atomicsDefs|]]
         cudafy = [CUDAC.cunit|
+$esc:("#define FUTHARK_CUDA")
+
 typedef char int8_t;
 typedef short int16_t;
 typedef int int32_t;
@@ -445,6 +430,7 @@
 static inline void mem_fence_global() {
   __threadfence();
 }
+
 $esc:("#define NAN (0.0/0.0)")
 $esc:("#define INFINITY (1.0/0.0)")
 extern volatile __shared__ char shared_mem[];
@@ -541,48 +527,57 @@
                              _            -> pointerQuals "global"
           return [C.cty|$tyquals:(volatile++quals) $ty:t|]
 
-        doAtomic s old arr ind val op ty = do
+        atomicSpace (Space sid) = sid
+        atomicSpace _           = "global"
+
+        doAtomic s t old arr ind val op ty = do
           ind' <- GenericC.compileExp $ unCount ind
           val' <- GenericC.compileExp val
           cast <- atomicCast s ty
-          GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|]
+          GenericC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|]
+          where op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s
 
-        atomicOps s (AtomicAdd old arr ind val) =
-          doAtomic s old arr ind val "atomic_add" [C.cty|int|]
+        atomicOps s (AtomicAdd t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_add" [C.cty|int|]
 
-        atomicOps s (AtomicSMax old arr ind val) =
-          doAtomic s old arr ind val "atomic_max" [C.cty|int|]
+        atomicOps s (AtomicFAdd t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_fadd" [C.cty|float|]
 
-        atomicOps s (AtomicSMin old arr ind val) =
-          doAtomic s old arr ind val "atomic_min" [C.cty|int|]
+        atomicOps s (AtomicSMax t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_smax" [C.cty|int|]
 
-        atomicOps s (AtomicUMax old arr ind val) =
-          doAtomic s old arr ind val "atomic_max" [C.cty|unsigned int|]
+        atomicOps s (AtomicSMin t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_smin" [C.cty|int|]
 
-        atomicOps s (AtomicUMin old arr ind val) =
-          doAtomic s old arr ind val "atomic_min" [C.cty|unsigned int|]
+        atomicOps s (AtomicUMax t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|]
 
-        atomicOps s (AtomicAnd old arr ind val) =
-          doAtomic s old arr ind val "atomic_and" [C.cty|unsigned int|]
+        atomicOps s (AtomicUMin t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|]
 
-        atomicOps s (AtomicOr old arr ind val) =
-          doAtomic s old arr ind val "atomic_or" [C.cty|unsigned int|]
+        atomicOps s (AtomicAnd t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_and" [C.cty|int|]
 
-        atomicOps s (AtomicXor old arr ind val) =
-          doAtomic s old arr ind val "atomic_xor" [C.cty|unsigned int|]
+        atomicOps s (AtomicOr t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_or" [C.cty|int|]
 
-        atomicOps s (AtomicCmpXchg old arr ind cmp val) = do
+        atomicOps s (AtomicXor t old arr ind val) =
+          doAtomic s t old arr ind val "atomic_xor" [C.cty|int|]
+
+        atomicOps s (AtomicCmpXchg t old arr ind cmp val) = do
           ind' <- GenericC.compileExp $ unCount ind
           cmp' <- GenericC.compileExp cmp
           val' <- GenericC.compileExp 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');|]
+          GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|]
+          where op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s
 
-        atomicOps s (AtomicXchg old arr ind val) = do
+        atomicOps s (AtomicXchg t old arr ind val) = do
           ind' <- GenericC.compileExp $ unCount ind
           val' <- GenericC.compileExp val
           cast <- atomicCast s [C.cty|int|]
-          GenericC.stm [C.cstm|$id:old = atomic_xchg(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|]
+          GenericC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|]
+          where op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s
 
         cannotAllocate :: GenericC.Allocate KernelOp KernelState
         cannotAllocate _ =
@@ -627,7 +622,7 @@
                 | has_communication = [C.citems|local_failure = true;
                                                 goto $id:label;|]
                 | otherwise         = [C.citems|return;|]
-          GenericC.stm [C.cstm|{ if (atomic_cmpxchg(global_failure, -1, $int:n) == -1)
+          GenericC.stm [C.cstm|{ if (atomic_cmpxchg_i32_global(global_failure, -1, $int:n) == -1)
                                  { $stms:argstms; }
                                  $items:what_next
                                }|]
diff --git a/src/Futhark/CodeGen/ImpGen/OpenCL.hs b/src/Futhark/CodeGen/ImpGen/OpenCL.hs
--- a/src/Futhark/CodeGen/ImpGen/OpenCL.hs
+++ b/src/Futhark/CodeGen/ImpGen/OpenCL.hs
@@ -9,4 +9,4 @@
 import Futhark.MonadFreshNames
 
 compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m OpenCL.Program
-compileProg prog = kernelsToOpenCL <$> ImpGenKernels.compileProg prog
+compileProg prog = kernelsToOpenCL <$> ImpGenKernels.compileProgOpenCL prog
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
@@ -10,9 +10,9 @@
 import Futhark.MonadFreshNames
 
 compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m 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 :: ImpGen.OpCompiler ExplicitMemory () Imp.Sequential
         opCompiler dest (Alloc e space) =
           ImpGen.compileAlloc dest e space
         opCompiler _ (Inner _) =
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
@@ -989,8 +989,6 @@
 defuncValBind (ValBind entry name _ (Info (rettype, retext)) tparams params body _ loc)
   | Scalar Arrow{} <- rettype = do
       (body_pats, body', rettype') <- etaExpand (fromStruct rettype) body
-      -- FIXME: we should also handle non-constant size annotations
-      -- here.
       defuncValBind $ ValBind entry name Nothing
         (Info (rettype', retext))
         tparams (params <> body_pats) body' Nothing loc
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
@@ -474,6 +474,7 @@
 
   -- 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 resulting kernel.
+  use_scope <- (<>scopeOf rem_bnds) <$> askScope
   (_,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
@@ -493,7 +494,7 @@
                     cons_no_out_transf = SOAC.nullTransforms $ outputTransform ker
 
                 consumer_ok   <- do let consumer_bnd   = rem_bnds !! bnd_ind
-                                    maybesoac <- SOAC.fromExp $ stmExp consumer_bnd
+                                    maybesoac <- runReaderT (SOAC.fromExp $ stmExp consumer_bnd) use_scope
                                     case maybesoac of
                                       -- check that consumer's lambda body does not use
                                       -- directly the produced arrays (e.g., see noFusion3.fut).
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
@@ -2,6 +2,8 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE Strict #-}
+{-# LANGUAGE StrictData #-}
 -- |
 --
 -- Perform general rule-based simplification based on data dependency
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
@@ -16,6 +16,7 @@
 
 import Control.Monad.RWS.Strict
 import Data.Maybe
+import qualified Data.Map as M
 import Data.List (find)
 
 import Futhark.Pass.ExtractKernels.Distribution
@@ -69,7 +70,7 @@
         pat' = Pattern [] $ rearrangeShape perm $ patternValueElements pat
 
     return $
-      SeqLoop [0..patternSize pat-1] pat' merge_expanded form $
+      SeqLoop perm pat' merge_expanded form $
       mkBody (pre_copy_bnds<>oneStm map_bnd) res
   where free_in_body = freeIn body
 
@@ -123,7 +124,8 @@
                  -> m (Stms SOACS)
 interchangeLoops nest loop = do
   (loop', bnds) <-
-    runBinder $ foldM (interchangeLoop isMapParameter) loop $ reverse $ kernelNestLoops nest
+    runBinder $ foldM (interchangeLoop isMapParameter) loop $
+    reverse $ kernelNestLoops nest
   return $ bnds <> oneStm (seqLoopStm loop')
   where isMapParameter v =
           fmap snd $ find ((==v) . paramName . fst) $
@@ -150,13 +152,13 @@
           Pattern [] $ map (fmap (`arrayOfRow` w)) $ patternElements branch_pat
 
         mkBranch branch = (renameBody=<<) $ do
-          branch' <- if null $ bodyStms branch
-                     then runBodyBinder $
-                          -- XXX: We need a temporary dummy binding to
-                          -- prevent an empty map body.  The kernel
-                          -- extractor does not like empty map bodies.
-                          resultBody <$> mapM dummyBind (bodyResult branch)
-                     else return branch
+          let bound_in_branch = scopeOf (bodyStms branch)
+          branch' <-
+            -- XXX: We may need dummys binding to prevent identity
+            -- mappings.  The kernel extractor does not like identity
+            -- mappings.
+            runBodyBinder $
+            resultBody <$> (mapM (dummyBindIfNotIn bound_in_branch) =<< bodyBind branch)
           let lam = Lambda params branch' lam_ret
               res = map Var $ patternNames branch_pat'
               map_bnd = Let branch_pat' (StmAux cs ()) $ Op $ Screma w (mapSOAC lam) arrs
@@ -170,6 +172,11 @@
           dummy <- newVName "dummy"
           letBindNames_ [dummy] (BasicOp $ SubExp se)
           return $ Var dummy
+
+        dummyBindIfNotIn bound_in_branch se
+          | Var v <- se,
+            v `M.member` bound_in_branch = return se
+          | otherwise = dummyBind se
 
 interchangeBranch :: (MonadFreshNames m, HasScope SOACS m) =>
                      KernelNest -> Branch -> m (Stms SOACS)
diff --git a/src/Futhark/Representation/AST/Syntax.hs b/src/Futhark/Representation/AST/Syntax.hs
--- a/src/Futhark/Representation/AST/Syntax.hs
+++ b/src/Futhark/Representation/AST/Syntax.hs
@@ -1,4 +1,6 @@
 {-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, StandaloneDeriving #-}
+{-# LANGUAGE Strict #-}
+{-# LANGUAGE StrictData #-}
 -- | Futhark core language skeleton.  Concrete representations further
 -- extend this skeleton by defining a "lore", which specifies concrete
 -- annotations ("Futhark.Representation.AST.Annotations") and
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
@@ -1,5 +1,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE Strict #-}
+{-# LANGUAGE StrictData #-}
 -- | The most primitive ("core") aspects of the AST.  Split out of
 -- "Futhark.Representation.AST.Syntax" in order for
 -- "Futhark.Representation.AST.Annotations" to use these definitions.  This
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
@@ -40,7 +40,14 @@
   blockers { Engine.blockHoistBranch = blockAllocs }
   where blockAllocs vtable _ (Let _ _ (Op Alloc{})) =
           not $ ST.simplifyMemory vtable
-        blockAllocs _ _ _ = False
+        -- Do not hoist statements that produce arrays.  This is
+        -- because in the ExplicitMemory representation, multiple
+        -- arrays can be located in the same memory block, and moving
+        -- their creation out of a branch can thus cause memory
+        -- corruption.  At this point in the compiler we have probably
+        -- already moved all the array creations that matter.
+        blockAllocs _ _ (Let pat _ _) =
+          not $ all primType $ patternTypes pat
 
 simplifyStms :: (HasScope ExplicitMemory m, MonadFreshNames m) =>
                 Stms ExplicitMemory -> m (ST.SymbolTable (Wise ExplicitMemory),
diff --git a/src/Language/Futhark/Core.hs b/src/Language/Futhark/Core.hs
--- a/src/Language/Futhark/Core.hs
+++ b/src/Language/Futhark/Core.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE Strict #-}
+{-# LANGUAGE StrictData #-}
 -- | This module contains very basic definitions for Futhark - so basic,
 -- that they can be shared between the internal and external
 -- representation.
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
@@ -2,6 +2,8 @@
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE StandaloneDeriving         #-}
+{-# LANGUAGE StrictData                 #-}
+{-# LANGUAGE Strict                     #-}
 -- | This is an ever-changing syntax representation for Futhark.  Some
 -- types, such as @Exp@, are parametrised by type and name
 -- representation.  See the @https://futhark.readthedocs.org@ for a
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
@@ -566,8 +566,11 @@
 nastyType _ = True
 
 nastyReturnType :: Monoid als => Maybe (TypeExp VName) -> TypeBase dim als -> Bool
-nastyReturnType _ (Scalar (Arrow _ _ t1 t2)) =
+nastyReturnType Nothing (Scalar (Arrow _ _ t1 t2)) =
   nastyType t1 || nastyReturnType Nothing t2
+nastyReturnType (Just (TEArrow _ te1 te2 _)) (Scalar (Arrow _ _ t1 t2)) =
+  (not (niceTypeExp te1) && nastyType t1) ||
+  nastyReturnType (Just te2) t2
 nastyReturnType (Just te) _
   | niceTypeExp te = False
 nastyReturnType te t
diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs
--- a/src/Language/Futhark/TypeChecker/Unify.hs
+++ b/src/Language/Futhark/TypeChecker/Unify.hs
@@ -320,6 +320,10 @@
 type UnifyDims m =
   BreadCrumbs -> [VName] -> (VName -> Maybe Int) -> DimDecl VName -> DimDecl VName -> m ()
 
+flipUnifyDims :: UnifyDims m -> UnifyDims m
+flipUnifyDims onDims bcs bound nonrigid t1 t2 =
+  onDims bcs bound nonrigid t2 t1
+
 unifyWith :: MonadUnify m =>
              UnifyDims m -> Usage -> BreadCrumbs
           -> StructType -> StructType -> m ()
@@ -340,9 +344,14 @@
 
           -- Remove any of the intermediate dimensions we added just
           -- for unification purposes.
-          link v lvl = linkVarToType onDims usage bcs v lvl . applySubst unbind
+          link ord' v lvl =
+            linkVarToType linkDims usage bcs v lvl . applySubst unbind
             where unbind d | d `elem` bound = Just $ SizeSubst AnyDim
                            | otherwise      = Nothing
+                  -- We may have to flip the order of future calls to
+                  -- onDims inside linkVarToType.
+                  linkDims | ord' = flipUnifyDims onDims
+                           | otherwise = onDims
 
           unifyTypeArg bcs' (TypeArgDim d1 _) (TypeArgDim d2 _) =
             onDims' bcs' (swap ord d1 d2)
@@ -379,18 +388,18 @@
          Scalar (TypeVar _ _ (TypeName [] v2) [])) ->
           case (nonrigid v1, nonrigid v2) of
             (Nothing, Nothing) -> failure
-            (Just lvl1, Nothing) -> link v1 lvl1 t2'
-            (Nothing, Just lvl2) -> link v2 lvl2 t1'
+            (Just lvl1, Nothing) -> link ord v1 lvl1 t2'
+            (Nothing, Just lvl2) -> link (not ord) v2 lvl2 t1'
             (Just lvl1, Just lvl2)
-              | lvl1 <= lvl2 -> link v1 lvl1 t2'
-              | otherwise    -> link v2 lvl2 t1'
+              | lvl1 <= lvl2 -> link ord v1 lvl1 t2'
+              | otherwise    -> link (not ord) v2 lvl2 t1'
 
         (Scalar (TypeVar _ _ (TypeName [] v1) []), _)
           | Just lvl <- nonrigid v1 ->
-              link v1 lvl t2'
+              link ord v1 lvl t2'
         (_, Scalar (TypeVar _ _ (TypeName [] v2) []))
           | Just lvl <- nonrigid v2 ->
-              link v2 lvl t1'
+              link (not ord) v2 lvl t1'
 
         (Scalar (Arrow _ p1 a1 b1),
          Scalar (Arrow _ p2 a2 b2)) -> do
@@ -858,8 +867,8 @@
 unifyMostCommon usage t1 t2 = do
   -- We are ignoring the dimensions here, because any mismatches
   -- should be turned into fresh size variables.
-  unify usage (toStruct (anySizes t1))
-              (toStruct (anySizes t2))
+  let allOK _ _ _ _ _ = return ()
+  unifyWith allOK usage noBreadCrumbs (toStruct t1) (toStruct t2)
   t1' <- normTypeFully t1
   t2' <- normTypeFully t2
   newDimOnMismatch (srclocOf usage) t1' t2'