diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,29 @@
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
+## [0.26.2]
+
+### Added
+
+* New server commands: `zip` and `unzip`, matching the corresponding C APIs.
+
+* The restrictions for passing tuples as consumed function parameters have been
+  loosened. (#2456)
+
+* The reverse-mode AD transformation now supports custom adjoints through a new
+  prelude function, `with_vjp`.
+
+### Fixed
+
+* A regression in fusion (#2444).
+
+* A potential compiler crash in register tiling (#2441).
+
+* The type checker would disregard uniqueness annotations on local functions.
+  (#2459)
+
+* Some loops would have aliases inferred incorrectly. (#2461)
+
 ## [0.26.1]
 
 ### Changed
diff --git a/docs/server-protocol.rst b/docs/server-protocol.rst
--- a/docs/server-protocol.rst
+++ b/docs/server-protocol.rst
@@ -222,6 +222,22 @@
 *v1*, which must be an array of rank *N*, at position *[i0]...[iN-1]*, where
 each *i* is an integer. Fails if the index is out of bounds.
 
+``zip`` *v0* *type* *v1* ... *vN*
+.................................
+
+Create a new variable *v0* of type *type*, which must be an array of records
+where the elements have *N* fields, where *v1* to *vN* are variables that are
+arrays of the corresponding field types. The order in which the arrays must be
+passed are given by the ``fields`` command on *type*.
+
+``unzip`` *v0* *v1* ... *vN*
+............................
+
+Unzip an array of records into new variables. The variable *v0* must be an
+array whose element type is a record with *N* fields. The order of constructed
+arrays corresponds to the field order given by the ``fields`` command on the
+type of *v0*.
+
 Record Commands
 ~~~~~~~~~~~~~~~
 
diff --git a/docs/usage.rst b/docs/usage.rst
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -360,14 +360,6 @@
   with a type abbreviation to give it a specific name, otherwise one
   will be generated.
 
-Return types follow these rules, with one addition:
-
-* If the return type is an *m*-element tuple, then the function
-  returns *m* values, mapped according to the rules above (but not
-  including this one - nested tuples are not mapped directly).  This
-  rule does not apply when the entry point has been given a return
-  type ascription that is not syntactically a tuple type.
-
 .. _api-consumption:
 
 Consumption and Aliasing
@@ -391,9 +383,8 @@
    Further, any *aliases* of that value are also considered consumed
    and may not be used.
 
-2. Each entry point output is either *unique* or *nonunique*.  A
-   unique output has no aliases.  A nonunique output aliases *every*
-   nonconsuming input parameter.
+2. The entry point output iseither *unique* or *nonunique*. A unique output has
+   no aliases. A nonunique output aliases *every* nonconsuming input parameter.
 
 Note that these distinctions are currently usually not visible in the
 generated API, and so correct usage requires knowledge of the original
diff --git a/futhark.cabal b/futhark.cabal
--- a/futhark.cabal
+++ b/futhark.cabal
@@ -1,6 +1,6 @@
 cabal-version: 3.0
 name:           futhark
-version:        0.26.1
+version:        0.26.2
 synopsis:       An optimising compiler for a functional, array-oriented language.
 
 description:    Futhark is a small programming language designed to be compiled to
@@ -479,7 +479,7 @@
     , filepath >=1.4.1.1
     , free >=5.1.10
     , futhark-data >= 1.1.3.0
-    , futhark-server >= 1.4.0.0
+    , futhark-server >= 1.4.1.0
     , futhark-manifest == 1.8.0.0
     , githash >=0.1.6.1
     , half >= 0.3
diff --git a/prelude/ad.fut b/prelude/ad.fut
--- a/prelude/ad.fut
+++ b/prelude/ad.fut
@@ -12,6 +12,21 @@
 -- purpose where you might need derivatives, such as for example
 -- computing surface normals for signed distance functions.
 --
+-- Futhark's AD support includes the following:
+--
+--   * Differentiation operators for forward-mode (`jvp`) and reverse-mode
+--     (`vjp`).
+--
+--   * Arbitrary control flow in differentiable code.
+--
+--   * Higher order derivatives by nesting differentiation operators, including
+--     arbitrary mixing of forward- and reverse mode (although using multiple
+--     rounds of reverse mode is rarely useful and often slow).
+--
+--   * Custom derivatives (`with_vjp`).
+--
+--   * Checkpointing of sequential loops.
+--
 -- ## Jacobians
 --
 -- For a differentiable function *f* whose input comprise *n* scalars
@@ -115,3 +130,26 @@
 -- | Vector-Jacobian Product ("reverse mode").
 def vjp 'a 'b (f: a -> b) (x: a) (y': b) : a =
   (vjp2 f x y').1
+
+-- | Provide custom reverse-mode adjoint code for a given function. This is
+-- useful when the adjoint synthesised by AD is not as good as one that is known
+-- analytically.
+--
+-- The function `f` returns a result of type `b`. In the return sweep, the
+-- function `f'` is invoked first with the result of `f` and second with the
+-- cotangents of the result (be careful not to mix up the order), and must
+-- return the sensitivity with respect to the input.
+--
+-- A common pattern is that `b` is a tuple where some part is the intended
+-- primal result of `with_vjp`, and some part is only used in `f'`.
+--
+-- **Beware:** if `f` uses any free variables, these will not be taken into
+-- **account when computing the adjoint. Make these part of the argument
+-- **instead.
+def with_vjp 'a 'b (f: a -> b) (f': (res: b) -> (b_adj: b) -> a) (x: a) : b =
+  intrinsics.with_vjp f f' x
+
+-- | A variant of `with_vjp` where the intermediate result necessary for the
+-- adjoint (`c`) is explicitly separated from the primal result (`b`).
+def with_vjp_tape 'a 'b 'c (f: a -> (c, b)) (f': (c, b) -> a) (x: a) : b =
+  (with_vjp f (\(tape, _) (_, adj) -> f' (tape, adj)) x).1
diff --git a/rts/c/backends/c.h b/rts/c/backends/c.h
--- a/rts/c/backends/c.h
+++ b/rts/c/backends/c.h
@@ -6,7 +6,7 @@
   int profiling;
   int logging;
   char *cache_fname;
-  struct tuning_param tuning_params[NUM_TUNING_PARAMS];
+  struct tuning_param tuning_params[NUM_TUNING_PARAMS+1];
 };
 
 static void backend_context_config_setup(struct futhark_context_config* cfg) {
diff --git a/rts/c/backends/hip.h b/rts/c/backends/hip.h
--- a/rts/c/backends/hip.h
+++ b/rts/c/backends/hip.h
@@ -309,12 +309,6 @@
       (device_query(ctx->dev, hipDeviceAttributePhysicalMultiProcessorCount) *
        device_query(ctx->dev, hipDeviceAttributeMaxThreadsPerMultiProcessor))
       / cfg->gpu.default_block_size;
-
-    // XXX: this is a hack due to the inability of two-pass scan to handle a
-    // grid size that is larger than the maximum block size.
-    if (cfg->gpu.default_grid_size > ctx->max_thread_block_size) {
-      cfg->gpu.default_grid_size = ctx->max_thread_block_size;
-    }
   }
 
   for (int i = 0; i < NUM_TUNING_PARAMS; i++) {
diff --git a/rts/c/backends/opencl.h b/rts/c/backends/opencl.h
--- a/rts/c/backends/opencl.h
+++ b/rts/c/backends/opencl.h
@@ -811,12 +811,6 @@
     ctx->max_shared_memory = max_shared_memory;
   }
 
-  // XXX: this is a hack due to the inability of two-pass scan to handle a
-  // grid size that is larger than the maximum block size.
-  if (ctx->cfg->gpu.default_grid_size > max_thread_block_size) {
-    ctx->cfg->gpu.default_grid_size = max_thread_block_size;
-  }
-
   // Now we go through all the sizes, clamp them to the valid range,
   // or set them to the default.
   for (int i = 0; i < NUM_TUNING_PARAMS; i++) {
diff --git a/rts/c/cache.h b/rts/c/cache.h
--- a/rts/c/cache.h
+++ b/rts/c/cache.h
@@ -44,7 +44,7 @@
 }
 
 #define CACHE_HEADER_SIZE 8
-static const char cache_header[CACHE_HEADER_SIZE] = "FUTHARK\0";
+static const char cache_header[CACHE_HEADER_SIZE] = "FUTHARK";
 
 static int cache_restore(const char *fname, const struct cache_hash *hash,
                          unsigned char **buf, size_t *buflen) {
diff --git a/rts/c/server.h b/rts/c/server.h
--- a/rts/c/server.h
+++ b/rts/c/server.h
@@ -21,6 +21,9 @@
 typedef int (*array_set_fn)(struct futhark_context *, const void*, const void*, const int64_t*);
 typedef const int64_t* (*array_shape_fn)(struct futhark_context*, const void*);
 typedef int (*array_index_fn)(struct futhark_context*, void*, const void*, const int64_t*);
+typedef int (*array_values_fn)(struct futhark_context*, const void*, void*);
+typedef int (*array_free_data_fn)(struct futhark_context*, void*);
+typedef int (*array_zip_fn)(struct futhark_context*, void*, const void*[]);
 typedef int (*project_fn)(struct futhark_context*, void*, const void*);
 typedef int (*variant_fn)(struct futhark_context*, const void*);
 typedef int (*new_fn)(struct futhark_context*, void**, const void*[]);
@@ -34,19 +37,26 @@
   OPAQUE
 };
 
+struct field {
+  const char *name;
+  const struct type *type;
+  project_fn project;
+};
+
 struct array {
   int rank;
   const struct type *element_type;
+  const struct primtype_info_t* info;
+  const char *name;
+  int num_fields;
+  const struct field *fields;
   array_new_fn new;
   array_set_fn set;
   array_shape_fn shape;
   array_index_fn index;
-};
-
-struct field {
-  const char *name;
-  const struct type *type;
-  project_fn project;
+  array_values_fn values;
+  array_free_data_fn free;
+  array_zip_fn zip;
 };
 
 struct record {
@@ -763,7 +773,8 @@
     return;
   }
 
-  int64_t* dims = alloca(a->rank * sizeof(int64_t));
+  int64_t* dims = alloca((size_t)a->rank * sizeof(int64_t));
+
   int64_t n_values = 1;
 
   for (int i = 0; i < a->rank; ++i) {
@@ -788,14 +799,40 @@
     return;
   }
 
-  const void** value_ptrs = alloca(n_values * sizeof(void*));
+  char *values = NULL;
+  const void **value_ptrs = NULL;
 
+  if (n_values < 0) {
+    failure();
+    printf("Invalid array size.\n");
+    return;
+  }
+
+  if (a->info != NULL) {
+    size_t values_size = (size_t)n_values * a->info->size;
+    values = malloc(values_size);
+    if (values == NULL) {
+      failure();
+      printf("Out of memory.\n");
+      return;
+    }
+  } else {
+    value_ptrs = malloc((size_t)n_values * sizeof(void*));
+    if (value_ptrs == NULL) {
+      failure();
+      printf("Out of memory.\n");
+      return;
+    }
+  }
+
   for (int64_t i = 0; i < n_values; i++) {
     struct variable* v = get_variable(s, args[2+a->rank+i]);
 
     if (v == NULL) {
       failure();
       printf("Unknown variable: %s\n", args[2+a->rank+i]);
+      free(value_ptrs);
+      free(values);
       return;
     }
 
@@ -803,13 +840,21 @@
       failure();
       printf("Value %d mismatch: expected type %s, got %s\n",
              (int)i, a->element_type->name, v->value.type->name);
+      free(value_ptrs);
+      free(values);
       return;
     }
 
-    value_ptrs[i] = value_ptr(&v->value);
+    if (a->info != NULL) {
+      memcpy(values + i * a->info->size, value_ptr(&v->value), a->info->size);
+    } else {
+      value_ptrs[i] = value_ptr(&v->value);
+    }
   }
 
-  a->new(s->ctx, value_ptr(&to->value), value_ptrs, dims);
+  a->new(s->ctx, value_ptr(&to->value), a->info != NULL ? (void*)values : value_ptrs, dims);
+  free(value_ptrs);
+  free(values);
 }
 
 void cmd_set(struct server_state *s, const char *args[]) {
@@ -935,6 +980,134 @@
   a->index(s->ctx, value_ptr(&to->value), from->value.value.v_ptr, indices);
 }
 
+void cmd_zip(struct server_state *s, const char *args[]) {
+  const char *to_name = get_arg(args, 0);
+  const char *type_name = get_arg(args, 1);
+  const struct type *type = get_type(s, type_name);
+
+  if (type->kind != ARRAY) {
+    failure();
+    printf("Not an array type\n");
+    return;
+  }
+
+  const struct array *a = type->info;
+  if (a->zip == NULL || a->fields == NULL) {
+    failure();
+    printf("Cannot zip this array type\n");
+    return;
+  }
+
+  int num_args = 0;
+  for (int i = 2; arg_exists(args, i); i++) {
+    num_args++;
+  }
+
+  if (num_args != a->num_fields) {
+    failure();
+    printf("%d arrays expected but %d values provided.\n", a->num_fields, num_args);
+    return;
+  }
+
+  const void** value_ptrs = alloca(num_args * sizeof(void*));
+
+  for (int i = 0; i < num_args; i++) {
+    struct variable* v = get_variable(s, args[2+i]);
+
+    if (v == NULL) {
+      failure();
+      printf("Unknown variable: %s\n", args[2+i]);
+      return;
+    }
+
+    if (strcmp(v->value.type->name, a->fields[i].type->name) != 0) {
+      failure();
+      printf("Field %s mismatch: expected type %s, got %s\n",
+             a->fields[i].name, a->fields[i].type->name, v->value.type->name);
+      return;
+    }
+
+    value_ptrs[i] = v->value.value.v_ptr;
+  }
+
+  struct variable *to = create_variable(s, to_name, type);
+
+  if (to == NULL) {
+    failure();
+    printf("Variable already exists: %s\n", to_name);
+    return;
+  }
+
+  int err = a->zip(s->ctx, value_ptr(&to->value), value_ptrs);
+  err |= futhark_context_sync(s->ctx);
+  error_check(s, err);
+  if (err != 0) {
+    drop_variable(to);
+  }
+}
+
+void cmd_unzip(struct server_state *s, const char *args[]) {
+  const char *from_name = get_arg(args, 0);
+  struct variable* from = get_variable(s, from_name);
+
+  if (from == NULL) {
+    failure();
+    printf("Unknown variable: %s\n", from_name);
+    return;
+  }
+
+  if (from->value.type->kind != ARRAY) {
+    failure();
+    printf("Not an array type\n");
+    return;
+  }
+
+  const struct array *a = from->value.type->info;
+  if (a->fields == NULL) {
+    failure();
+    printf("Cannot unzip this array type\n");
+    return;
+  }
+
+  int num_args = 0;
+  for (int i = 1; arg_exists(args, i); i++) {
+    num_args++;
+  }
+
+  if (num_args != a->num_fields) {
+    failure();
+    printf("%d arrays expected but %d values provided.\n", a->num_fields, num_args);
+    return;
+  }
+
+  struct variable **outs = alloca(num_args * sizeof(struct variable*));
+  for (int i = 0; i < num_args; i++) {
+    const char *to_name = get_arg(args, i+1);
+    struct variable *to = create_variable(s, to_name, a->fields[i].type);
+    if (to == NULL) {
+      failure();
+      printf("Variable already exists: %s\n", to_name);
+      for (int j = 0; j < i; j++) {
+        drop_variable(outs[j]);
+      }
+      return;
+    }
+    outs[i] = to;
+  }
+
+  int err = 0;
+  for (int i = 0; i < num_args; i++) {
+    err |= a->fields[i].project(s->ctx, value_ptr(&outs[i]->value), from->value.value.v_ptr);
+  }
+  err |= futhark_context_sync(s->ctx);
+  error_check(s, err);
+  if (err != 0) {
+    for (int i = 0; i < num_args; i++) {
+      drop_variable(outs[i]);
+    }
+  }
+}
+
 void cmd_fields(struct server_state *s, const char *args[]) {
   const char *type = get_arg(args, 0);
   const struct type *t = get_type(s, type);
@@ -1335,6 +1508,10 @@
     cmd_set(s, tokens+1);
   } else if (strcmp(command, "index") == 0) {
     cmd_index(s, tokens+1);
+  } else if (strcmp(command, "zip") == 0) {
+    cmd_zip(s, tokens+1);
+  } else if (strcmp(command, "unzip") == 0) {
+    cmd_unzip(s, tokens+1);
   } else if (strcmp(command, "fields") == 0) {
     cmd_fields(s, tokens+1);
   } else if (strcmp(command, "variants") == 0) {
@@ -1393,59 +1570,39 @@
 // The aux struct lets us write generic method implementations without
 // code duplication.
 
-typedef void* (*aux_array_new_fn)(struct futhark_context*, const void**, const int64_t*);
-typedef const int64_t* (*aux_array_shape_fn)(struct futhark_context*, void*);
-typedef int (*aux_array_index_fn)(struct futhark_context*, void*, const void*, const int64_t*);
-typedef int (*aux_array_values_fn)(struct futhark_context*, void*, void*);
-typedef int (*aux_array_free_fn)(struct futhark_context*, void*);
-
-struct array_aux {
-  int rank;
-  const struct primtype_info_t* info;
-  const char *name;
-  aux_array_new_fn new;
-  aux_array_shape_fn shape;
-  aux_array_values_fn values;
-  aux_array_free_fn free;
-};
-
-int restore_array(const struct array_aux *aux, FILE *f,
+int restore_array(const struct array *a, FILE *f,
                   struct futhark_context *ctx, void *p) {
   void *data = NULL;
-  int64_t shape[aux->rank];
-  if (read_array(f, aux->info, &data, shape, aux->rank) != 0) {
+  int64_t shape[a->rank];
+  if (read_array(f, a->info, &data, shape, a->rank) != 0) {
     return 1;
   }
 
-  void *arr = aux->new(ctx, data, shape);
-  if (arr == NULL) {
-    return 1;
-  }
-  int err = futhark_context_sync(ctx);
-  *(void**)p = arr;
+  int err = a->new(ctx, p, data, shape);
+  err |= futhark_context_sync(ctx);
   free(data);
   return err;
 }
 
-void store_array(const struct array_aux *aux, FILE *f,
+void store_array(const struct array *a, FILE *f,
                  struct futhark_context *ctx, void *p) {
   void *arr = *(void**)p;
-  const int64_t *shape = aux->shape(ctx, arr);
-  int64_t size = sizeof(aux->info->size);
-  for (int i = 0; i < aux->rank; i++) {
+  const int64_t *shape = a->shape(ctx, arr);
+  int64_t size = a->info->size;
+  for (int i = 0; i < a->rank; i++) {
     size *= shape[i];
   }
   int32_t *data = malloc(size);
-  assert(aux->values(ctx, arr, data) == 0);
+  assert(a->values(ctx, arr, data) == 0);
   assert(futhark_context_sync(ctx) == 0);
-  assert(write_array(f, 1, aux->info, data, shape, aux->rank) == 0);
+  assert(write_array(f, 1, a->info, data, shape, a->rank) == 0);
   free(data);
 }
 
-int free_array(const struct array_aux *aux,
+int free_array(const struct array *a,
                struct futhark_context *ctx, void *p) {
   void *arr = *(void**)p;
-  return aux->free(ctx, arr);
+  return a->free(ctx, arr);
 }
 
 typedef void* (*opaque_restore_fn)(struct futhark_context*, void*);
diff --git a/rts/c/values.h b/rts/c/values.h
--- a/rts/c/values.h
+++ b/rts/c/values.h
@@ -529,7 +529,7 @@
 //// Types
 
 struct primtype_info_t {
-  const char binname[4]; // Used for parsing binary data.
+  const char binname[5]; // Used for parsing binary data.
   const char* type_name; // Same name as in Futhark.
   const int64_t size; // in bytes
   const writer write_str; // Write in text format.
diff --git a/rts/javascript/server.js b/rts/javascript/server.js
--- a/rts/javascript/server.js
+++ b/rts/javascript/server.js
@@ -256,6 +256,9 @@
       case 'types': this._cmd_types(args); break
       case 'fields': this._cmd_fields(args); break
       case 'project': this._cmd_project(args); break
+          // XXX: these should be implemented.
+      case 'attributes': this._cmd_dummy(args); break
+      case 'entry_points': this._cmd_dummy(args); break
       default: throw "Unknown command: " + cmd;
       }
     }
diff --git a/rts/python/server.py b/rts/python/server.py
--- a/rts/python/server.py
+++ b/rts/python/server.py
@@ -40,12 +40,12 @@
 
     def _cmd_inputs(self, args):
         entry = self._get_arg(args, 0)
-        for t in self._get_entry_point(entry)[1]:
+        for t in self._get_entry_point(entry)["inputs"]:
             print(t)
 
     def _cmd_output(self, args):
         entry = self._get_arg(args, 0)
-        print(self._get_entry_point(entry)[2])
+        print(self._get_entry_point(entry)["output"])
 
     def _cmd_dummy(self, args):
         pass
@@ -65,8 +65,8 @@
 
     def _cmd_call(self, args):
         entry = self._get_entry_point(self._get_arg(args, 0))
-        entry_fname = entry[0]
-        num_ins = len(entry[1])
+        entry_fname = entry["name"]
+        num_ins = len(entry["inputs"])
         exp_len = 2 + num_ins
 
         if len(args) != exp_len:
@@ -177,6 +177,9 @@
         # FIXME: assuming a tuple.
         self._vars[dst] = self._vars[src].data[int(field)]
 
+    def _cmd_attributes(self, args):
+        return self._get_entry_point(self._get_arg(args, 0))["attributes"]
+
     def _cmd_entry_points(self, args):
         for k in self._ctx.entry_points.keys():
             print(k)
@@ -197,6 +200,7 @@
         "entry_points": _cmd_entry_points,
         "fields": _cmd_fields,
         "project": _cmd_project,
+        "attributes": _cmd_attributes,
     }
 
     def _process_line(self, line):
diff --git a/src-testing/Language/Futhark/TypeChecker/ConsumptionTests.hs b/src-testing/Language/Futhark/TypeChecker/ConsumptionTests.hs
--- a/src-testing/Language/Futhark/TypeChecker/ConsumptionTests.hs
+++ b/src-testing/Language/Futhark/TypeChecker/ConsumptionTests.hs
@@ -29,7 +29,7 @@
               [Id "x_1" (Info "[2]i32") mempty]
               "[2]i32"
               ( second
-                  (const (S.singleton (AliasBound "x_1")))
+                  (const (S.singleton (AliasBound "x_1" [])))
                   ("[2]i32" :: StructType)
               )
               @?= "[2]i32",
@@ -39,7 +39,7 @@
               [Id "x_1" (Info "[2]i32") mempty]
               "([2]i32, [2]i32)"
               ( second
-                  (const (S.singleton (AliasFree "y_2")))
+                  (const (S.singleton (AliasFree "y_2" [])))
                   ("([2]i32,[2]i32)" :: StructType)
               )
               @?= "([2]i32, [2]i32)",
@@ -49,7 +49,7 @@
              in inferReturnUniqueness
                   [Id "n_1" (Info "i64") mempty]
                   t
-                  (second (const (S.singleton (AliasFree "y_3"))) t)
+                  (second (const (S.singleton (AliasFree "y_3" []))) t)
                   @?= (t `setUniqueness` Nonunique),
           --
           testCase "*opaque" $
diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs
--- a/src/Futhark/AD/Fwd.hs
+++ b/src/Futhark/AD/Fwd.hs
@@ -331,6 +331,12 @@
             histNeutral = interleave nes nes_tan,
             histOp = op'
           }
+fwdSOAC pat aux (WithVJP args lam _) = do
+  -- You have a custom adjoint? Too bad we are in tangent land.
+  (mapM_ fwdStm <=< runBuilder_) $ do
+    lam_res <- auxing aux $ eLambda lam $ map eSubExp args
+    forM (zip (patNames pat) lam_res) $ \(v, SubExpRes cs se) ->
+      certifying cs $ letBindNames [v] $ BasicOp $ SubExp se
 fwdSOAC _ _ JVP {} =
   error "fwdSOAC: nested JVP not allowed."
 fwdSOAC _ _ VJP {} =
diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs
--- a/src/Futhark/AD/Rev/SOAC.hs
+++ b/src/Futhark/AD/Rev/SOAC.hs
@@ -189,6 +189,18 @@
 vjpSOAC ops pat aux (Stream w as accs lam) m = do
   stms <- collectStms_ $ auxing aux $ sequentialStreamWholeArray pat w accs lam as
   foldr (vjpStm ops) m stms
+vjpSOAC _ops pat aux (WithVJP args lam lam_adj) m = do
+  lam_res <- auxing aux (eLambda lam (map eSubExp args))
+  forM_ (zip (patNames pat) lam_res) $ \(v, SubExpRes cs se) ->
+    certifying cs $ letBindNames [v] $ BasicOp $ SubExp se
+  m
+  pat_adj <- mapM lookupAdjVal $ patNames pat
+  contribs <-
+    eLambda lam_adj (map (eSubExp . resSubExp) lam_res ++ map (eSubExp . Var) pat_adj)
+  forM_ (zip args contribs) $ \(arg, contrib) ->
+    (updateSubExpAdj arg <=< letExp "contrib") $
+      BasicOp . SubExp . resSubExp $
+        contrib
 vjpSOAC _ _ _ soac _ =
   error $ "vjpSOAC unhandled:\n" ++ prettyString soac
 
diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs
@@ -6,7 +6,7 @@
   )
 where
 
-import Data.List (unzip5)
+import Data.List (intersperse, unzip5)
 import Data.Map qualified as M
 import Data.Text qualified as T
 import Futhark.CodeGen.Backends.GenericC.Options
@@ -20,6 +20,7 @@
   )
 import Futhark.CodeGen.RTS.C (tuningH, valuesH)
 import Futhark.Manifest
+import Futhark.Util (showText)
 import Futhark.Util.Pretty (prettyString)
 import Language.C.Quote.OpenCL qualified as C
 import Language.C.Syntax qualified as C
@@ -262,6 +263,15 @@
             [C.cstm|assert($id:(opaqueFree ops)(ctx, $id:result) == 0);|]
           )
 
+recordFieldCType :: Manifest -> RecordField -> C.Type
+recordFieldCType manifest field =
+  case M.lookup t $ manifestTypes manifest of
+    Nothing -> uncurry primAPIType $ scalarToPrim t
+    Just (TypeArray tname _ _ _) -> [C.cty|typename $id:tname|]
+    Just (TypeOpaque tname _ _) -> [C.cty|typename $id:tname|]
+  where
+    t = recordFieldType field
+
 -- | Return a statement printing the given external value.
 printStm :: Manifest -> T.Text -> C.Exp -> C.Stm
 printStm manifest tname e =
@@ -269,6 +279,22 @@
     Nothing ->
       let info = tname <> "_info"
        in [C.cstm|write_scalar(stdout, binary_output, &$id:info, &$exp:e);|]
+    Just (TypeOpaque _ _ (Just (OpaqueRecord record)))
+      | map recordFieldName fields == take (length fields) (map showText [0 :: Int ..]) ->
+          [C.cstm|{$stms:(intersperse newline (map getField fields))}|]
+      where
+        fields = recordFields record
+        printField field =
+          printStm manifest (recordFieldType field) [C.cexp|field|]
+        newline = [C.cstm|puts("");|]
+        getField field =
+          [C.cstm|{$ty:(recordFieldCType manifest field) field;
+                   if ($id:(recordFieldProject field)(ctx, &field, $exp:e) != FUTHARK_SUCCESS) {
+                     futhark_panic(1, "Failed to project field %s from result\n", $string:(T.unpack (recordFieldName field)));
+                   } else {
+                     $stm:(printField field)
+                   }
+                   }|]
     Just (TypeOpaque desc _ _) ->
       [C.cstm|{
          fprintf(stderr, "Values of type \"%s\" have no external representation.\n", $string:(T.unpack desc));
@@ -295,7 +321,7 @@
 printResult :: Manifest -> [(T.Text, C.Exp)] -> [C.Stm]
 printResult manifest = concatMap f
   where
-    f (v, e) = [printStm manifest v e, [C.cstm|printf("\n");|]]
+    f (v, e) = [printStm manifest v e, [C.cstm|puts("");|]]
 
 cliEntryPoint ::
   Manifest -> T.Text -> EntryPoint -> (C.Definition, C.Initializer)
diff --git a/src/Futhark/CodeGen/Backends/GenericC/Server.hs b/src/Futhark/CodeGen/Backends/GenericC/Server.hs
--- a/src/Futhark/CodeGen/Backends/GenericC/Server.hs
+++ b/src/Futhark/CodeGen/Backends/GenericC/Server.hs
@@ -147,9 +147,7 @@
       element_c_type = cType manifest et
       type_name = typeStructName tname
       array_name = type_name <> "_array"
-      aux_name = type_name <> "_aux"
       info_name = et <> "_info"
-      aux_array_new_wrap = arrayNew ops <> "_aux_wrap"
       array_new_wrap = arrayNew ops <> "_wrap"
       array_set = arrayNew ops <> "_set"
       array_index_wrap = arrayIndex ops <> "_wrap"
@@ -158,24 +156,12 @@
    in ( [C.cedecl|const struct type $id:type_name;|],
         [C.cinit|&$id:type_name|],
         [C.cunit|
-              void* $id:aux_array_new_wrap(struct futhark_context *ctx,
-                                           const void* p,
-                                           const typename int64_t* shape) {
-                return $id:(arrayNew ops)(ctx, p, $args:shape_args);
-              }
-              int $id:array_new_wrap(struct futhark_context* ctx,
-                                     typename $id:c_type_name* outp,
-                                     $ty:element_c_type *ps[],
+               int $id:array_new_wrap(struct futhark_context* ctx,
+                                     void** outp,
+                                     const void* p,
                                      const typename int64_t* shape) {
-                typename int64_t n_values = 1;
-                for (int i = 0; i < $int:rank; ++i) {
-                  n_values *= shape[i];
-                }
-                $ty:element_c_type *values = alloca(n_values * sizeof($ty:element_c_type));
-                for (typename int64_t i = 0; i < n_values; ++i) {
-                  values[i] = *ps[i];
-                }
-                *outp = $id:(arrayNew ops)(ctx, values, $args:shape_args);
+                typename $id:c_type_name *out = (typename $id:c_type_name*) outp;
+                *out = $id:(arrayNew ops)(ctx, p, $args:shape_args);
                 return 0;
               }
               int $id:array_set(struct futhark_context *ctx,
@@ -200,26 +186,24 @@
               const struct array $id:array_name = {
                 .rank = $int:rank,
                 .element_type = &$id:element_type_name,
+                .info = &$id:info_name,
+                .name = $string:(T.unpack tname),
+                .num_fields = 0,
+                .fields = NULL,
                 .new = (typename array_new_fn)$id:array_new_wrap,
                 .set = (typename array_set_fn)$id:array_set,
                 .shape = (typename array_shape_fn)$id:(arrayShape ops),
                 .index = (typename array_index_fn)$id:array_index_wrap,
-              };
-              const struct array_aux $id:aux_name = {
-                .name = $string:(T.unpack tname),
-                .rank = $int:rank,
-                .info = &$id:info_name,
-                .new = (typename aux_array_new_fn)$id:aux_array_new_wrap,
-                .free = (typename aux_array_free_fn)$id:(arrayFree ops),
-                .shape = (typename aux_array_shape_fn)$id:(arrayShape ops),
-                .values = (typename aux_array_values_fn)$id:(arrayValues ops)
+                .values = (typename array_values_fn)$id:(arrayValues ops),
+                .free = (typename array_free_data_fn)$id:(arrayFree ops),
+                .zip = NULL,
               };
               const struct type $id:type_name = {
                 .name = $string:(T.unpack tname),
                 .restore = (typename restore_fn)restore_array,
                 .store = (typename store_fn)store_array,
                 .free = (typename free_fn)free_array,
-                .aux = &$id:aux_name,
+                .aux = &$id:array_name,
                 .kind = $exp:(cKind Array),
                 .info = &$id:array_name
               };|]
@@ -335,11 +319,39 @@
             [C.cinit|&$id:sum_name|],
             Sum
           )
-    transparentDefs type_name (Just (OpaqueArray ops')) = opaqueArrayDefs type_name (opaqueArrayRank ops') (opaqueArrayElemType ops') (opaqueArrayNew ops') (opaqueArraySet ops') (opaqueArrayShape ops') (opaqueArrayIndex ops')
-    transparentDefs type_name (Just (OpaqueRecordArray ops')) = opaqueArrayDefs type_name (recordArrayRank ops') (recordArrayElemType ops') (recordArrayNew ops') (recordArraySet ops') (recordArrayShape ops') (recordArrayIndex ops')
+    transparentDefs type_name (Just (OpaqueArray ops')) =
+      opaqueArrayDefs
+        type_name
+        (opaqueArrayRank ops')
+        (opaqueArrayElemType ops')
+        (opaqueArrayNew ops')
+        (opaqueArraySet ops')
+        (opaqueArrayShape ops')
+        (opaqueArrayIndex ops')
+        Nothing
+    transparentDefs type_name (Just (OpaqueRecordArray ops')) =
+      opaqueArrayDefs
+        type_name
+        (recordArrayRank ops')
+        (recordArrayElemType ops')
+        (recordArrayNew ops')
+        (recordArraySet ops')
+        (recordArrayShape ops')
+        (recordArrayIndex ops')
+        $ Just (recordArrayFields ops', recordArrayZip ops')
     transparentDefs _ _ = ([], [C.cinit|NULL|], Opaque)
 
-    opaqueArrayDefs type_name rank et new set shape index =
+    opaqueArrayDefs ::
+      T.Text ->
+      Int ->
+      T.Text ->
+      CFuncName ->
+      CFuncName ->
+      CFuncName ->
+      CFuncName ->
+      Maybe ([RecordField], CFuncName) ->
+      ([C.Definition], C.Initializer, Kind)
+    opaqueArrayDefs type_name rank et new set shape index maybe_fields =
       let array_name = type_name <> "_array"
           element_type_name = typeStructName et
           element_c_type = cType manifest et
@@ -348,10 +360,48 @@
           index_wrap = index <> "_wrap"
           shape_args = [[C.cexp|shape[$int:i]|] | i <- [0 .. rank - 1]]
           is_args = [[C.cexp|is[$int:i]|] | i <- [0 .. rank - 1]]
-       in ( [C.cunit|
-              int $id:new_wrap(struct futhark_context *ctx,
-                               typename $id:c_type_name *outp,
-                               $ty:element_c_type *ps[],
+          (zip_defs, fields_init, num_fields, zip_init) =
+            case maybe_fields of
+              Nothing ->
+                ( [],
+                  [C.cexp|NULL|],
+                  0 :: Int,
+                  [C.cexp|NULL|]
+                )
+              Just (fields, zip_f) ->
+                let fields_name = type_name <> "_zip_fields"
+                    zip_wrap = zip_f <> "_aux_wrap"
+                    onField i field =
+                      let field_type = cType manifest $ recordFieldType field
+                          field_v = "f" <> prettyText i
+                       in ( [C.cinit|{ .name = $string:(T.unpack (recordFieldName field)),
+                                       .type = &$id:(typeStructName (recordFieldType field)),
+                                       .project = (typename project_fn)$id:(recordFieldProject field)
+                                     }|],
+                            [C.citem|const $ty:field_type $id:field_v =
+                                         (const $ty:field_type)f[$int:i];|],
+                            [C.cexp|$id:field_v|]
+                          )
+                    (field_inits, get_fields, zip_args) = unzip3 $ zipWith onField [0 :: Int ..] fields
+                 in ( [C.cunit|const struct field $id:fields_name[] = {
+                                  $inits:field_inits
+                                };|]
+                        ++ [C.cunit|int $id:zip_wrap(struct futhark_context *ctx,
+                                                     void *outp,
+                                                     const void *f[]) {
+                                      typename $id:c_type_name *out = (typename $id:c_type_name*)outp;
+                                      $items:get_fields
+                                      return $id:zip_f(ctx, out, $args:zip_args);
+                                    }|],
+                      [C.cexp|$id:fields_name|],
+                      length fields,
+                      [C.cexp|(typename array_zip_fn)$id:zip_wrap|]
+                    )
+       in ( zip_defs
+              ++ [C.cunit|
+               int $id:new_wrap(struct futhark_context *ctx,
+                                typename $id:c_type_name *outp,
+                                $ty:element_c_type *ps[],
                                const typename int64_t shape[]) {
                 typename int64_t n_values = 1;
                 for (int i = 0; i < $int:rank; ++i) {
@@ -369,20 +419,27 @@
                                const typename int64_t *is) {
                 return $id:set(ctx, arr, *val, $args:is_args);
               }
-              int $id:index_wrap(struct futhark_context *ctx,
-                                 void *dest,
-                                 typename $id:c_type_name arr,
-                                 const typename int64_t *is) {
-                return $id:index(ctx, dest, arr, $args:is_args);
-              }
-              const struct array $id:array_name = {
-                .rank = $int:rank,
-                .element_type = &$id:element_type_name,
-                .new = (typename array_new_fn)$id:new_wrap,
-                .set = (typename array_set_fn)$id:set_wrap,
-                .shape = (typename array_shape_fn)$id:shape,
-                .index = (typename array_index_fn)$id:index_wrap,
-              };|],
+               int $id:index_wrap(struct futhark_context *ctx,
+                                  void *dest,
+                                  typename $id:c_type_name arr,
+                                  const typename int64_t *is) {
+                 return $id:index(ctx, dest, arr, $args:is_args);
+               }
+               const struct array $id:array_name = {
+                 .rank = $int:rank,
+                 .element_type = &$id:element_type_name,
+                 .info = NULL,
+                 .name = $string:(T.unpack tname),
+                 .num_fields = $int:num_fields,
+                 .fields = $exp:fields_init,
+                 .new = (typename array_new_fn)$id:new_wrap,
+                 .set = (typename array_set_fn)$id:set_wrap,
+                 .shape = (typename array_shape_fn)$id:shape,
+                 .index = (typename array_index_fn)$id:index_wrap,
+                 .values = NULL,
+                 .free = NULL,
+                 .zip = $exp:zip_init,
+                };|],
             [C.cinit|&$id:array_name|],
             Array
           )
diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs
--- a/src/Futhark/CodeGen/Backends/GenericPython.hs
+++ b/src/Futhark/CodeGen/Backends/GenericPython.hs
@@ -43,9 +43,11 @@
 
 import Control.Monad
 import Control.Monad.RWS hiding (reader, writer)
+import Data.Bifunctor (first)
 import Data.Char (isAlpha, isAlphaNum)
 import Data.Map qualified as M
 import Data.Maybe
+import Data.Set qualified as S
 import Data.Text qualified as T
 import Futhark.CodeGen.Backends.GenericPython.AST
 import Futhark.CodeGen.Backends.GenericPython.Options
@@ -438,14 +440,14 @@
 
       case mode of
         ToLibrary -> do
-          (entry_points, entry_point_types) <-
+          (entry_points, entry_point_info) <-
             unzip . catMaybes <$> mapM (compileEntryFun sync DoNotReturnTiming) funs
           pure
             [ ClassDef $
                 Class class_name $
                   [ Assign
                       (Var "entry_points")
-                      (Dict entry_point_types),
+                      (strDict entry_point_info),
                     opaques_def,
                     Assign
                       (Var "sizes")
@@ -454,7 +456,7 @@
                     <> map FunDef (constructor' : definitions ++ entry_points)
             ]
         ToServer -> do
-          (entry_points, entry_point_types) <-
+          (entry_points, entry_point_info) <-
             unzip . catMaybes <$> mapM (compileEntryFun sync ReturnTiming) funs
           pure $
             parse_options_server
@@ -462,7 +464,7 @@
                      ( Class class_name $
                          [ Assign
                              (Var "entry_points")
-                             (Dict entry_point_types),
+                             (strDict entry_point_info),
                            opaques_def,
                            Assign
                              (Var "sizes")
@@ -519,8 +521,8 @@
 
     selectEntryPoint entry_point_names entry_points =
       [ Assign (Var "entry_points") $
-          Dict $
-            zip (map String entry_point_names) entry_points,
+          strDict $
+            zip entry_point_names entry_points,
         Assign (Var "entry_point_fun") $
           simpleCall "entry_points.get" [Var "entry_point"],
         If
@@ -899,11 +901,15 @@
 
 data ReturnTiming = ReturnTiming | DoNotReturnTiming
 
+-- | Construct a dictionary with string keys.
+strDict :: [(T.Text, PyExp)] -> PyExp
+strDict = Dict . map (first String)
+
 compileEntryFun ::
   [PyStmt] ->
   ReturnTiming ->
   (Name, Imp.Function op) ->
-  CompilerM op s (Maybe (PyFunDef, (PyExp, PyExp)))
+  CompilerM op s (Maybe (PyFunDef, (T.Text, PyExp)))
 compileEntryFun sync timing fun
   | Just entry <- Imp.functionEntry $ snd fun = do
       let ename = Imp.entryPointName entry
@@ -939,11 +945,16 @@
         Just
           ( Def (T.unpack (escapeName ename)) ("self" : params) $
               prepareIn ++ do_run ++ prepareOut ++ sync ++ [ret],
-            ( String (nameToText ename),
-              Tuple
-                [ String (escapeName ename),
-                  List (map String pts),
-                  String rts
+            ( nameToText ename,
+              strDict
+                [ ("name", String (escapeName ename)),
+                  ("inputs", List (map String pts)),
+                  ("output", String rts),
+                  ( "attributes",
+                    List . map (String . prettyText) $
+                      S.toList . Imp.unAttrs . Imp.functionAttrs $
+                        snd fun
+                  )
                 ]
             )
           )
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs
@@ -178,7 +178,7 @@
           compileReduction (chunk_v, chunk_const) nonsegmentedReduction
         _ -> do
           let segment_size = pe64 $ last $ segSpaceDims space
-              use_small_segments = segment_size * 2 .<. pe64 (unCount tblock_size) * tvExp chunk_v
+              use_small_segments = segment_size * 2 .<. pe64 (unCount tblock_size)
           sIf
             use_small_segments
             (compileReduction (chunk_v, chunk_const) smallSegmentsReduction)
diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
--- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
+++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs
@@ -328,77 +328,173 @@
   TV Int32 ->
   Imp.TExp Int64 ->
   Count NumBlocks SubExp ->
+  Count BlockSize SubExp ->
   CrossesSegment ->
   SegSpace ->
   [SegBinOp GPUMem] ->
   CallKernelGen ()
-scanStage2 scan_out stage1_num_threads elems_per_group num_tblocks crossesSegment space scans = do
+scanStage2 scan_out stage1_num_threads elems_per_group stage1_num_tblocks stage2_tblock_size crossesSegment space scans = do
   let (gtids, dims) = unzip $ unSegSpace space
       dims' = map pe64 dims
-
-  -- Our group size is the number of groups for the stage 1 kernel.
-  let tblock_size = Count $ unCount num_tblocks
+      stage1_num_tblocks_e = pe64 $ unCount stage1_num_tblocks
+      stage2_tblock_size_e = pe64 $ unCount stage2_tblock_size
 
-  let crossesSegment' = do
-        f <- crossesSegment
-        Just $ \from to ->
-          f
-            ((sExt64 from + 1) * elems_per_group - 1)
-            ((sExt64 to + 1) * elems_per_group - 1)
+  -- Number of chunks needed to cover all stage-1 blocks.
+  num_chunks <-
+    dPrimVE "stage2_num_chunks" $
+      stage1_num_tblocks_e `divUp` stage2_tblock_size_e
 
-  sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) tblock_size) $ do
+  sKernelThread "scan_stage2" (segFlat space) (defKernelAttrs (Count (intConst Int64 1)) stage2_tblock_size) $ do
     constants <- kernelConstants <$> askEnv
-    per_scan_local_arrs <- makeLocalArrays tblock_size (tvSize stage1_num_threads) scans
+    per_scan_local_arrs <- makeLocalArrays stage2_tblock_size (tvSize stage1_num_threads) scans
     let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans
         per_scan_pes = segBinOpChunks scans scan_out
 
-    flat_idx <-
-      dPrimV "flat_idx" $
-        (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1
-    -- Construct segment indices.
-    zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx
+    -- Declare lambda params and initialise carries (xParams) to the
+    -- neutral element.  For scalar scans these persist across chunk
+    -- iterations as registers; for array scans they are reloaded from
+    -- global memory at the start of each chunk.
+    forM_ scans $ \scan -> do
+      dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan
+      forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->
+        copyDWIMFix (paramName p) [] ne []
 
-    forM_ (L.zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $
-      \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) ->
-        sLoopNest vec_shape $ \vec_is -> do
-          let glob_is = map Imp.le64 gtids ++ vec_is
+    sFor "chunk_id" num_chunks $ \chunk_id -> do
+      let chunk_offset = chunk_id * stage2_tblock_size_e
 
-              in_bounds =
-                foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims'
+      flat_idx <-
+        dPrimV "flat_idx" $
+          (chunk_offset + sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1
+      -- Construct segment indices.
+      zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx
 
-              when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) ->
-                copyDWIMFix
-                  arr
-                  [localArrayIndex constants t]
-                  (Var pe)
-                  glob_is
+      forM_ (L.zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $
+        \(scan@(SegBinOp _ scan_op nes vec_shape), local_arrs, rets, pes) ->
+          sComment "do one stage-2 scan chunk" $ do
+            let (array_scan, fence, barrier) = barrierFor scan_op
+                scan_x_params = xParams scan
+                scan_y_params = yParams scan
+                -- Scalar scans with a non-trivial vec_shape need per-vec-element
+                -- carries reloaded from global memory, just like array scans.
+                -- The scalar carry register cannot hold independent carries for
+                -- each vector element across chunk iterations.
+                use_global_carry = array_scan || not (null (shapeDims vec_shape))
 
-              when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) ->
-                copyDWIMFix arr [localArrayIndex constants t] ne []
-              (_, _, barrier) =
-                barrierFor scan_op
+            when use_global_carry (sOp $ Imp.Barrier Imp.FenceGlobal)
 
-          sComment "threads in bound read carries; others get neutral element" $
-            sIf in_bounds when_in_bounds when_out_of_bounds
+            sLoopNest vec_shape $ \vec_is -> do
+              let glob_is = map Imp.le64 gtids ++ vec_is
+                  in_bounds =
+                    foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims'
+                  -- Element index of the last element of the previous chunk,
+                  -- used to load the inter-chunk carry from global memory.
+                  prev_chunk_last = chunk_offset * elems_per_group - 1
 
-          barrier
+              -- For array scans and scalar scans with non-trivial vec_shape,
+              -- reload carry (xParams) from the scan output written by the
+              -- previous chunk.  Thread 0 reads the last element of the
+              -- previous chunk, unless a segment boundary falls between the
+              -- chunks.
+              when use_global_carry $ do
+                crosses_seg <-
+                  dPrimVE "crosses_seg" $
+                    case crossesSegment of
+                      Nothing -> false
+                      Just f -> f prev_chunk_last (prev_chunk_last + 1)
+                sIf
+                  (chunk_id .>. 0 .&&. kernelLocalThreadId constants .==. 0 .&&. bNot crosses_seg)
+                  ( do
+                      let carry_is = unflattenIndex dims' prev_chunk_last
+                      forM_ (zip scan_x_params pes) $ \(p, pe) ->
+                        copyDWIMFix (paramName p) [] (Var pe) (carry_is ++ vec_is)
+                  )
+                  ( forM_ (zip scan_x_params nes) $ \(p, ne) ->
+                      copyDWIMFix (paramName p) [] ne []
+                  )
 
-          blockScan
-            crossesSegment'
-            (sExt64 $ tvExp stage1_num_threads)
-            (sExt64 $ kernelBlockSize constants)
-            scan_op
-            local_arrs
+              -- Load the stage-1 partial-scan result for this thread's
+              -- block into yParams (or the neutral element when out of
+              -- bounds).
+              sIf
+                in_bounds
+                ( forM_ (zip scan_y_params pes) $ \(p, pe) ->
+                    copyDWIMFix (paramName p) [] (Var pe) glob_is
+                )
+                ( forM_ (zip scan_y_params nes) $ \(p, ne) ->
+                    copyDWIMFix (paramName p) [] ne []
+                )
 
-          sComment "threads in bounds write scanned carries" $
-            sWhen in_bounds $
-              forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->
-                copyDWIMFix
-                  pe
-                  glob_is
-                  (Var arr)
-                  [localArrayIndex constants t]
+              -- Combine carry (xParams) with new value (yParams) and
+              -- write the result to shared/local memory.  Thread 0
+              -- incorporates the inter-chunk carry; other threads have
+              -- neutral in xParams, so the op is a no-op for them.
+              compileStms mempty (bodyStms $ lambdaBody scan_op) $
+                forM_ (zip3 rets local_arrs $ map resSubExp $ bodyResult $ lambdaBody scan_op) $
+                  \(t, arr, se) ->
+                    copyDWIMFix arr [localArrayIndex constants t] se []
 
+              sOp $ Imp.ErrorSync fence
+
+              -- crossesSegment' maps block-local thread IDs to element
+              -- indices, adjusting for the current chunk offset.
+              let crossesSegment' =
+                    crossesSegment >>= \f ->
+                      Just $ \from to ->
+                        f
+                          ((chunk_offset + sExt64 from + 1) * elems_per_group - 1)
+                          ((chunk_offset + sExt64 to + 1) * elems_per_group - 1)
+
+              scan_op_renamed <- renameLambda scan_op
+              blockScan
+                crossesSegment'
+                (sExt64 $ tvExp stage1_num_threads)
+                (sExt64 $ kernelBlockSize constants)
+                scan_op_renamed
+                local_arrs
+
+              sComment "threads in bounds write scanned carries" $
+                sWhen in_bounds $
+                  forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->
+                    copyDWIMFix
+                      pe
+                      glob_is
+                      (Var arr)
+                      [localArrayIndex constants t]
+
+              barrier
+
+              -- For scalar scans with trivial vec_shape (no vec loop), update
+              -- the carry register (xParams) so the next chunk can use it.
+              -- For array scans and scalar scans with non-trivial vec_shape,
+              -- the carry is reloaded from global memory at the start of each
+              -- chunk (above), so no register update is needed here.
+              unless use_global_carry $ do
+                let next_chunk_start = chunk_offset + stage2_tblock_size_e
+                crosses_seg2 <-
+                  dPrimVE "crosses_seg" $
+                    case crossesSegment of
+                      Nothing -> false
+                      Just f ->
+                        f
+                          (next_chunk_start * elems_per_group - 1)
+                          ((next_chunk_start + 1) * elems_per_group - 1)
+                let should_load_carry =
+                      kernelLocalThreadId constants .==. 0 .&&. bNot crosses_seg2
+                    load_carry =
+                      forM_ (zip local_arrs scan_x_params) $ \(arr, p) ->
+                        copyDWIMFix
+                          (paramName p)
+                          []
+                          (Var arr)
+                          [sExt64 (kernelBlockSize constants) - 1]
+                    load_neutral =
+                      forM_ (zip nes scan_x_params) $ \(ne, p) ->
+                        copyDWIMFix (paramName p) [] ne []
+                sWhen should_load_carry load_carry
+                sUnless should_load_carry load_neutral
+
+              barrier
+
 scanStage3 ::
   Pat LetDecMem ->
   [VName] ->
@@ -542,18 +638,18 @@
 compileSegScan pat lvl space ts scans kbody post_op = do
   attrs <- lvlKernelAttrs lvl
 
-  -- Since stage 2 involves a group size equal to the number of groups
-  -- used for stage 1, we have to cap this number to the maximum group
-  -- size.
-  stage1_max_num_tblocks <- dPrim "stage1_max_num_tblocks"
-  sOp $ Imp.GetSizeMax (tvVar stage1_max_num_tblocks) SizeThreadBlock
+  -- Stage 2 uses loop virtualization, so stage1_num_tblocks is no
+  -- longer capped by the maximum thread block size.
+  let stage1_num_tblocks = kAttrNumBlocks attrs
 
-  stage1_num_tblocks <-
+  -- The stage-2 block size is a tunable/user-settable parameter.  It
+  -- is independent of stage1_num_tblocks so the autotuner can treat
+  -- it as a fixed knob.
+  stage2_tblock_size_param <- getSize "segscan_stage2_tblock_size" SizeThreadBlock
+  stage2_tblock_size <-
     fmap (Imp.Count . tvSize) $
-      dPrimV "stage1_num_tblocks" $
-        sMin64 (tvExp stage1_max_num_tblocks) $
-          pe64 . Imp.unCount . kAttrNumBlocks $
-            attrs
+      dPrimV "stage2_tblock_size" $
+        tvExp stage2_tblock_size_param
 
   let shpT op = (segBinOpShape op,) <$> lambdaReturnType (segBinOpLambda op)
       scan_ts = concatMap shpT scans
@@ -590,5 +686,5 @@
 
   emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group
 
-  scanStage2 scan_out stage1_num_threads elems_per_group stage1_num_tblocks crossesSegment space scans
+  scanStage2 scan_out stage1_num_threads elems_per_group stage1_num_tblocks stage2_tblock_size crossesSegment space scans
   scanStage3 pat scan_out map_out (kAttrNumBlocks attrs) (kAttrBlockSize attrs) elems_per_group crossesSegment space scans post_op
diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs
--- a/src/Futhark/Doc/Generator.hs
+++ b/src/Futhark/Doc/Generator.hs
@@ -11,7 +11,6 @@
 import Data.List (find, groupBy, inits, intersperse, isPrefixOf, partition, sort, sortOn, tails)
 import Data.Map qualified as M
 import Data.Maybe
-import Data.Ord
 import Data.Set qualified as S
 import Data.String (fromString)
 import Data.Text qualified as T
@@ -635,15 +634,15 @@
     pure $ "?" <> mconcat (map (brackets . renderName . baseName) dims) <> "." <> t'
 
 qualNameHtml :: QualName VName -> DocM Html
-qualNameHtml (QualName names vname@(VName name tag)) =
-  if tag <= maxIntrinsicTag
-    then pure $ renderName name
+qualNameHtml (QualName names vname) =
+  if isIntrinsic vname
+    then pure $ renderName $ baseName vname
     else f <$> ref
   where
     prefix :: Html
     prefix = mapM_ ((<> ".") . renderName . baseName) names
-    f (Just s) = H.a ! A.href (fromString s) $ prefix <> renderName name
-    f Nothing = prefix <> renderName name
+    f (Just s) = H.a ! A.href (fromString s) $ prefix <> renderName (baseName vname)
+    f Nothing = prefix <> renderName (baseName vname)
 
     ref = do
       boring <- asks $ S.member vname . ctxNoLink
diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs
--- a/src/Futhark/IR/SOACS.hs
+++ b/src/Futhark/IR/SOACS.hs
@@ -54,6 +54,7 @@
     lamUsesAD = bodyUsesAD . lambdaBody
     expUsesAD (Op JVP {}) = True
     expUsesAD (Op VJP {}) = True
+    expUsesAD (Op WithVJP {}) = True
     expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam
     expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds post_lam))) =
       lamUsesAD lam
diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs
--- a/src/Futhark/IR/SOACS/SOAC.hs
+++ b/src/Futhark/IR/SOACS/SOAC.hs
@@ -82,6 +82,8 @@
     JVP [SubExp] [SubExp] (Lambda rep)
   | -- FIXME: this should not be here
     VJP [SubExp] [SubExp] (Lambda rep)
+  | -- FIXME: this should not be here
+    WithVJP [SubExp] (Lambda rep) (Lambda rep)
   | -- | A combination of scan, reduction, and map.  The first
     -- t'SubExp' is the size of the input arrays.
     Screma SubExp [VName] (ScremaForm rep)
@@ -409,6 +411,11 @@
     <$> mapM (mapOnSOACSubExp tv) args
     <*> mapM (mapOnSOACSubExp tv) vec
     <*> mapOnSOACLambda tv lam
+mapSOACM tv (WithVJP args lam0 lam1) =
+  WithVJP
+    <$> mapM (mapOnSOACSubExp tv) args
+    <*> mapOnSOACLambda tv lam0
+    <*> mapOnSOACLambda tv lam1
 mapSOACM tv (Stream size arrs accs lam) =
   Stream
     <$> mapOnSOACSubExp tv size
@@ -506,6 +513,8 @@
   lambdaReturnType lam ++ lambdaReturnType lam
 soacType (VJP _ _ lam) =
   lambdaReturnType lam ++ map paramType (lambdaParams lam)
+soacType (WithVJP _ lam _) =
+  lambdaReturnType lam
 soacType (Stream outersize _ accs lam) =
   map (substNamesInType substs) rtp
   where
@@ -526,6 +535,7 @@
 
   consumedInOp JVP {} = mempty
   consumedInOp VJP {} = mempty
+  consumedInOp WithVJP {} = mempty
   -- Only map functions can consume anything.  The operands to scan
   -- and reduce functions are always considered "fresh".
   consumedInOp (Screma _ arrs (ScremaForm map_lam _ _ _)) =
@@ -555,6 +565,11 @@
     JVP args vec (Alias.analyseLambda aliases lam)
   addOpAliases aliases (VJP args vec lam) =
     VJP args vec (Alias.analyseLambda aliases lam)
+  addOpAliases aliases (WithVJP args lam lam_adj) =
+    WithVJP
+      args
+      (Alias.analyseLambda aliases lam)
+      (Alias.analyseLambda aliases lam_adj)
   addOpAliases aliases (Stream size arr accs lam) =
     Stream size arr accs $ Alias.analyseLambda aliases lam
   addOpAliases aliases (Hist w arrs ops bucket_fun) =
@@ -614,6 +629,12 @@
       lam
       (zipWith (<>) (map depsOf' args) (map depsOf' vec))
       <> map (const $ freeIn args <> freeIn lam) (lambdaParams lam)
+  opDependencies (WithVJP args lam _lam_adj) =
+    lambdaDependencies
+      mempty
+      lam
+      (map depsOf' args)
+      <> map (const $ freeIn args <> freeIn lam) (lambdaParams lam)
   opDependencies (Screma w arrs (ScremaForm map_lam scans reds post_lam)) =
     let (scans_in, reds_in, map_deps) =
           splitAt3 (scanResults scans) (redResults reds) $
@@ -707,6 +728,17 @@
         </> PP.indent 2 (pretty $ map TC.argType args')
         </> "does not match type of seed vector"
         </> PP.indent 2 (pretty vec_ts)
+typeCheckSOAC (WithVJP args lam lam_adj) = do
+  args' <- mapM TC.checkArg args
+  TC.checkLambda lam $ map TC.noArgAliases args'
+  TC.checkLambda lam_adj $
+    map (,mempty) (lambdaReturnType lam <> lambdaReturnType lam)
+  unless (lambdaReturnType lam_adj == map TC.argType args') $
+    TC.bad . TC.TypeError . docText $
+      "Adjoint lambda return type"
+        </> PP.indent 2 (pretty $ lambdaReturnType lam_adj)
+        </> "does not match type of arguments"
+        </> PP.indent 2 (pretty $ map TC.argType args')
 typeCheckSOAC (Stream size arrexps accexps lam) = do
   TC.require (Prim int64) size
   accargs <- mapM TC.checkArg accexps
@@ -823,6 +855,8 @@
     VJP args vec <$> rephraseLambda r lam
   rephraseInOp r (JVP args vec lam) =
     JVP args vec <$> rephraseLambda r lam
+  rephraseInOp r (WithVJP args lam lam_adj) =
+    WithVJP args <$> rephraseLambda r lam <*> rephraseLambda r lam_adj
   rephraseInOp r (Stream w arrs acc lam) =
     Stream w arrs acc <$> rephraseLambda r lam
   rephraseInOp r (Hist w arrs ops lam) =
@@ -852,6 +886,9 @@
     inside "VJP" $ lambdaMetrics lam
   opMetrics (JVP _ _ lam) =
     inside "JVP" $ lambdaMetrics lam
+  opMetrics (WithVJP _ lam lam_adj) = do
+    inside "WithVJP" $ lambdaMetrics lam
+    inside "WithVJP" $ lambdaMetrics lam_adj
   opMetrics (Stream _ _ _ lam) =
     inside "Stream" $ lambdaMetrics lam
   opMetrics (Hist _ _ ops bucket_fun) =
@@ -879,6 +916,14 @@
             PP.braces (commasep $ map pretty args)
               <> comma </> PP.braces (commasep $ map pretty vec)
               <> comma </> pretty lam
+        )
+  pretty (WithVJP args lam lam_adj) =
+    "with_vjp"
+      <> parens
+        ( PP.align $
+            PP.braces (commasep $ map pretty args)
+              <> comma </> pretty lam
+              <> comma </> pretty lam_adj
         )
   pretty (Stream size arrs acc lam) =
     ppStream size arrs acc lam
diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs
--- a/src/Futhark/IR/SOACS/Simplify.hs
+++ b/src/Futhark/IR/SOACS/Simplify.hs
@@ -101,6 +101,11 @@
   arr' <- mapM Engine.simplify arr
   vec' <- mapM Engine.simplify vec
   pure (JVP arr' vec' lam', hoisted)
+simplifySOAC (WithVJP args lam lam_adj) = do
+  args' <- mapM Engine.simplify args
+  (lam', hoisted) <- Engine.simplifyLambda mempty lam
+  (lam_adj', hoisted_adj) <- Engine.simplifyLambda mempty lam_adj
+  pure (WithVJP args' lam' lam_adj', hoisted <> hoisted_adj)
 simplifySOAC (Stream outerdim arr nes lam) = do
   outerdim' <- Engine.simplify outerdim
   nes' <- mapM Engine.simplify nes
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
@@ -245,7 +245,7 @@
       pure sv
     Nothing -- If the variable is unknown, it may refer to the 'intrinsics'
     -- module, which we will have to treat specially.
-      | baseTag x <= maxIntrinsicTag -> pure IntrinsicSV
+      | isIntrinsic x -> pure IntrinsicSV
       | otherwise ->
           -- Anything not in scope is going to be an existential size.
           pure $ Dynamic $ Scalar $ Prim $ Signed Int64
diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs
--- a/src/Futhark/Internalise/Exps.hs
+++ b/src/Futhark/Internalise/Exps.hs
@@ -365,14 +365,14 @@
       case () of
         ()
           -- Short-circuiting operators are magical.
-          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
+          | isIntrinsic (qualLeaf qfname),
             baseName (qualLeaf qfname) == "&&",
             [(x, _), (y, _)] <- args ->
               internaliseExp desc $
                 E.AppExp
                   (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)
                   (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])
-          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
+          | isIntrinsic (qualLeaf qfname),
             baseName (qualLeaf qfname) == "||",
             [(x, _), (y, _)] <- args ->
               internaliseExp desc $
@@ -389,7 +389,7 @@
               internalise =<< mapM prepareArg args
           | Just internalise <- isIntrinsicFunction qfname (map fst args) ->
               internalise desc
-          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,
+          | isIntrinsic (qualLeaf qfname),
             Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do
               let tag ses = [(se, I.Observe) | se <- ses]
               args' <- reverse <$> mapM (internaliseArg arg_desc) (reverse args)
@@ -1672,7 +1672,7 @@
   Name ->
   Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])
 isOverloadedFunction qname desc = do
-  guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag
+  guard $ isIntrinsic $ qualLeaf qname
   handle $ baseName $ qualLeaf qname
   where
     -- Handle equality and inequality specially, to treat the case of
@@ -1757,7 +1757,7 @@
   [E.Exp] ->
   Maybe (Name -> InternaliseM [SubExp])
 isIntrinsicFunction qname args = do
-  guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag
+  guard $ isIntrinsic $ qualLeaf qname
   let handlers =
         [ handleSign,
           handleOps,
@@ -1858,6 +1858,13 @@
             case fname of
               "jvp2" -> JVP x' v' lam
               _ -> VJP x' v' lam
+    handleAD [f, f_adj, x] "with_vjp" = Just $ \desc -> do
+      x' <- internaliseExp "ad_x" x
+      lam <- internaliseLambdaCoerce f =<< mapM subExpType x'
+      lam_adj <-
+        internaliseLambdaCoerce f_adj $
+          lambdaReturnType lam ++ lambdaReturnType lam
+      fmap (map I.Var) . letTupExp desc . Op $ WithVJP x' lam lam_adj
     handleAD _ _ = Nothing
 
     handleRest [a, si, v] "scatter" = Just $ scatterF 1 a si v
diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs
--- a/src/Futhark/Internalise/Monomorphise.hs
+++ b/src/Futhark/Internalise/Monomorphise.hs
@@ -90,15 +90,13 @@
 canCalculate scope mapping = do
   filter
     ( (`S.isSubsetOf` scope)
-        . S.filter notIntrisic
+        . S.filter (not . isIntrinsic)
         . fvVars
         . freeInExp
         . unReplaced
         . fst
     )
     mapping
-  where
-    notIntrisic vn = baseTag vn > maxIntrinsicTag
 
 -- Replace some expressions by a parameter.
 expReplace :: ExpReplacements -> Exp -> Exp
@@ -250,9 +248,7 @@
 -- that is arguments not currently in scope.
 askIntros :: S.Set VName -> MonoM (S.Set VName)
 askIntros argset =
-  (S.filter notIntrisic argset `S.difference`) <$> askScope
-  where
-    notIntrisic vn = baseTag vn > maxIntrinsicTag
+  (S.filter (not . isIntrinsic) argset `S.difference`) <$> askScope
 
 -- | Gets and removes expressions that could not be calculated when
 -- the arguments set will be unscoped.
@@ -394,7 +390,7 @@
 transformFName loc fname ft = do
   t' <- transformType ft
   let mono_t = monoType ft
-  if baseTag (qualLeaf fname) <= maxIntrinsicTag
+  if isIntrinsic (qualLeaf fname)
     then pure $ var fname t'
     else do
       maybe_fname <- lookupLifted (qualLeaf fname) mono_t
@@ -948,11 +944,10 @@
       Sum <$> (traverse . traverse) (arrowArgType env) cs
     arrowArgScalar (scope', dimsToPush) (Arrow as argName d argT retT) =
       pass $ do
-        let intros = S.filter notIntrisic argset' `S.difference` scope'
+        let intros = S.filter (not . isIntrinsic) argset' `S.difference` scope'
         retT' <- arrowArgRetType (scope', filter (`S.notMember` intros) dimsToPush) fullArgset retT
         pure (Arrow as argName d argT retT', bimap (intros `S.union`) (const mempty))
       where
-        notIntrisic vn = baseTag vn > maxIntrinsicTag
         argset' = fvVars $ freeInType argT
         fullArgset =
           case argName of
@@ -1131,7 +1126,7 @@
   runWriterT $ fst <$> execStateT (sub orig_t1 orig_t2) (mempty, mempty)
   where
     subRet (Scalar (TypeVar _ v _)) rt =
-      unless (baseTag (qualLeaf v) <= maxIntrinsicTag) $
+      unless (isIntrinsic (qualLeaf v)) $
         addSubst v rt
     subRet t1 (RetType _ t2) =
       sub t1 t2
@@ -1144,7 +1139,7 @@
         _ -> pure ()
       sub (stripArray 1 t1) (stripArray 1 t2)
     sub (Scalar (TypeVar _ v _)) t =
-      unless (baseTag (qualLeaf v) <= maxIntrinsicTag) $
+      unless (isIntrinsic (qualLeaf v)) $
         addSubst v $
           RetType [] t
     sub (Scalar (Record fields1)) (Scalar (Record fields2)) =
diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs
--- a/src/Futhark/Internalise/TypesValues.hs
+++ b/src/Futhark/Internalise/TypesValues.hs
@@ -256,7 +256,7 @@
       | otherwise ->
           concat <$> mapM (internaliseTypeM exts . snd) (E.sortFields ets)
     E.Scalar (E.TypeVar u tn [E.TypeArgType arr_t])
-      | baseTag (E.qualLeaf tn) <= E.maxIntrinsicTag,
+      | E.isIntrinsic (E.qualLeaf tn),
         baseName (E.qualLeaf tn) == "acc" -> do
           ts <-
             foldMap (toList . fmap (fromDecl . onAccType))
diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs
--- a/src/Futhark/Optimise/BlkRegTiling.hs
+++ b/src/Futhark/Optimise/BlkRegTiling.hs
@@ -678,7 +678,13 @@
         code1,
     -- identify load_A, load_B
     tmp_stms <- mapMaybe (`M.lookup` tab_inv_stm) arrs,
-    length tmp_stms == length arrs =
+    length tmp_stms == length arrs,
+    -- If any tiled-loop input array is also used in the postlude code,
+    -- we cannot safely apply this optimization (the array won't be in
+    -- scope in the epilogue).  See issue #2467.
+    not $
+      namesFromList (M.keys tab_inv_stm)
+        `namesIntersect` freeIn (code2'' <> code2) =
       let zip_AB = zip3 tmp_stms arrs [map_t1_0, map_t2_0]
           [(load_A, inp_A, map_t1), (load_B, inp_B, map_t2)] =
             if var_dims == [0, 1]
diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs
--- a/src/Futhark/Optimise/Fusion/GraphRep.hs
+++ b/src/Futhark/Optimise/Fusion/GraphRep.hs
@@ -394,6 +394,7 @@
   Futhark.Stream w is nes lam -> inputs is <> freeClassifications (w, nes, lam)
   Futhark.JVP {} -> freeClassifications soac
   Futhark.VJP {} -> freeClassifications soac
+  Futhark.WithVJP {} -> freeClassifications soac
   where
     inputs = S.fromList . (`zip` repeat SOACInput)
 expInputs e
diff --git a/src/Futhark/Optimise/Fusion/Screma.hs b/src/Futhark/Optimise/Fusion/Screma.hs
--- a/src/Futhark/Optimise/Fusion/Screma.hs
+++ b/src/Futhark/Optimise/Fusion/Screma.hs
@@ -255,11 +255,18 @@
                     <> varsRes (map paramName post_forward_params)
                 )
           }
+  -- Deduplicate inp_r: when inp_p and inp_c_real share inputs, the
+  -- combined input list contains duplicates.  Remove them by keeping
+  -- only one copy of each unique input and replacing the extra lambda
+  -- parameters in lam1 with let-bindings.
+  let nil_post = Lambda [] [] (mkBody mempty [])
+      (inp_r', form_lam1') = dedupInput inp_r (ScremaForm lam1 [] [] nil_post)
+      lam1' = scremaLambda form_lam1'
   pure
     ( SuperScrema
         w
-        inp_r
-        lam1
+        inp_r'
+        lam1'
         (scremaScans form_p)
         (scremaReduces form_p)
         lam2
diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs
--- a/src/Futhark/Optimise/Simplify/Rules.hs
+++ b/src/Futhark/Optimise/Simplify/Rules.hs
@@ -64,6 +64,7 @@
 removeUnnecessaryCopy :: (BuilderOps rep) => BottomUpRuleBasicOp rep
 removeUnnecessaryCopy (vtable, used) (Pat [d]) aux (Replicate (Shape []) (Var v))
   | not (v `UT.isConsumed` used),
+    allNames (not . (`UT.isConsumed` used)) $ ST.lookupAliases v vtable,
     -- This two first clauses below are too conservative, but the
     -- problem is that 'v' might not look like it has been consumed if
     -- it is consumed in an outer scope.  This is because the
diff --git a/src/Futhark/Pass/AD.hs b/src/Futhark/Pass/AD.hs
--- a/src/Futhark/Pass/AD.hs
+++ b/src/Futhark/Pass/AD.hs
@@ -35,48 +35,53 @@
   forM_ (zip (patNames pat) res) $ \(v, SubExpRes cs se) ->
     certifying cs $ letBindNames [v] $ BasicOp $ SubExp se
 
-onStm :: Mode -> Scope SOACS -> Stm SOACS -> PassM (Stms SOACS)
-onStm mode scope (Let pat aux (Op (VJP args vec lam))) = do
-  lam' <- onLambda mode scope lam
+onStm :: Bool -> Mode -> Scope SOACS -> Stm SOACS -> PassM (Stms SOACS)
+onStm _ mode scope (Let pat aux (Op (VJP args vec lam))) = do
+  lam' <- onLambda True mode scope lam
   if mode == All || lam == lam'
     then do
       lam'' <- (`runReaderT` scope) . simplifyLambda =<< revVJP scope lam'
       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope
     else pure $ oneStm $ Let pat aux $ Op $ VJP args vec lam'
-onStm mode scope (Let pat aux (Op (JVP args vec lam))) = do
-  lam' <- onLambda mode scope lam
+onStm _ mode scope (Let pat aux (Op (JVP args vec lam))) = do
+  lam' <- onLambda True mode scope lam
   if mode == All || lam == lam'
     then do
       lam'' <- fwdJVP scope lam'
       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope
     else pure $ oneStm $ Let pat aux $ Op $ JVP args vec lam'
-onStm mode scope (Let pat aux e) = oneStm . Let pat aux <$> mapExpM mapper e
+--
+-- This corresponds to a WithVJP that is not inside of a differential operator.
+-- FIXME: this assumption will go bad when we don't inline so much.
+onStm False _ scope (Let pat aux (Op (WithVJP args lam _))) =
+  runBuilderT_ (bindLambda pat aux lam args) scope
+onStm ad mode scope (Let pat aux e) = oneStm . Let pat aux <$> mapExpM mapper e
   where
     mapper =
       (identityMapper @SOACS)
-        { mapOnBody = \bscope -> onBody mode (bscope <> scope),
+        { mapOnBody = \bscope -> onBody ad mode (bscope <> scope),
           mapOnOp = mapSOACM soac_mapper
         }
-    soac_mapper = identitySOACMapper {mapOnSOACLambda = onLambda mode scope}
+    soac_mapper = identitySOACMapper {mapOnSOACLambda = onLambda ad mode scope}
 
-onStms :: Mode -> Scope SOACS -> Stms SOACS -> PassM (Stms SOACS)
-onStms mode scope stms = mconcat <$> mapM (onStm mode scope') (stmsToList stms)
+onStms :: Bool -> Mode -> Scope SOACS -> Stms SOACS -> PassM (Stms SOACS)
+onStms ad mode scope stms = mconcat <$> mapM (onStm ad mode scope') (stmsToList stms)
   where
     scope' = scopeOf stms <> scope
 
-onBody :: Mode -> Scope SOACS -> Body SOACS -> PassM (Body SOACS)
-onBody mode scope body = do
-  stms <- onStms mode scope $ bodyStms body
+onBody :: Bool -> Mode -> Scope SOACS -> Body SOACS -> PassM (Body SOACS)
+onBody ad mode scope body = do
+  stms <- onStms ad mode scope $ bodyStms body
   pure $ body {bodyStms = stms}
 
-onLambda :: Mode -> Scope SOACS -> Lambda SOACS -> PassM (Lambda SOACS)
-onLambda mode scope lam = do
-  body <- onBody mode (scopeOfLParams (lambdaParams lam) <> scope) $ lambdaBody lam
+onLambda :: Bool -> Mode -> Scope SOACS -> Lambda SOACS -> PassM (Lambda SOACS)
+onLambda ad mode scope lam = do
+  body <- onBody ad mode (scopeOfLParams (lambdaParams lam) <> scope) $ lambdaBody lam
   pure $ lam {lambdaBody = body}
 
 onFun :: Mode -> Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)
 onFun mode consts fd = do
-  body <- onBody mode (scopeOf consts <> scopeOf fd) $ funDefBody fd
+  body <- onBody False mode (scopeOf consts <> scopeOf fd) $ funDefBody fd
   pure $ fd {funDefBody = body}
 
 applyAD :: Pass SOACS SOACS
@@ -86,7 +91,7 @@
       passDescription = "Apply AD operators",
       passFunction =
         intraproceduralTransformationWithConsts
-          (onStms All mempty)
+          (onStms False All mempty)
           (onFun All)
     }
 
@@ -97,6 +102,6 @@
       passDescription = "Apply innermost AD operators",
       passFunction =
         intraproceduralTransformationWithConsts
-          (onStms Innermost mempty)
+          (onStms False Innermost mempty)
           (onFun Innermost)
     }
diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs
--- a/src/Futhark/Pass/ExtractMulticore.hs
+++ b/src/Futhark/Pass/ExtractMulticore.hs
@@ -213,6 +213,8 @@
   error "transformSOAC: unhandled JVP"
 transformSOAC _ _ VJP {} =
   error "transformSOAC: unhandled VJP"
+transformSOAC _ _ WithVJP {} =
+  error "transformSOAC: unhandled WithVJP"
 transformSOAC pat _ (Screma w arrs form)
   | Just lam <- isMapSOAC form = do
       seq_op <- transformMap DoNotRename sequentialiseBody w lam arrs
diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs
--- a/src/Futhark/Transform/FirstOrderTransform.hs
+++ b/src/Futhark/Transform/FirstOrderTransform.hs
@@ -235,6 +235,8 @@
   error "transformSOAC: unhandled JVP"
 transformSOAC _ VJP {} =
   error "transformSOAC: unhandled VJP"
+transformSOAC _ WithVJP {} =
+  error "transformSOAC: unhandled WithVJP"
 transformSOAC pat (Screma w arrs form) =
   transformScrema pat w arrs form
 transformSOAC pat (Stream w arrs nes lam) = do
diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs
--- a/src/Language/Futhark/Interpreter.hs
+++ b/src/Language/Futhark/Interpreter.hs
@@ -2156,6 +2156,10 @@
     def "manifest" = Just $ fun1 pure
     def "jvp2" = Just $ fun3 doJVP2
     def "vjp2" = Just $ fun3 doVJP2
+    def "with_vjp" = Just $ fun3 $ \f _ arg ->
+      -- XXX? We simply ignore the custom derivative. This is correct, but makes
+      -- it more of a hassle to test them.
+      apply noLoc mempty f arg
     def "acc" = Nothing
     def s | nameFromText s `M.member` namesToPrimTypes = Nothing
     def s = error $ "Missing intrinsic: " ++ T.unpack s
diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs
--- a/src/Language/Futhark/Prop.hs
+++ b/src/Language/Futhark/Prop.hs
@@ -8,6 +8,7 @@
     intrinsics,
     intrinsicVar,
     maxIntrinsicTag,
+    isIntrinsic,
     namesToPrimTypes,
     qualName,
     qualify,
@@ -966,6 +967,19 @@
                   $ RetType []
                   $ Scalar
                   $ tupleRecord [Scalar $ t_b Nonunique, Scalar $ t_a Nonunique]
+              ),
+              ( "with_vjp",
+                IntrinsicPolyFun
+                  [tp_a, tp_b]
+                  [ Scalar (t_a NoUniqueness) `arr` Scalar (t_b Nonunique),
+                    Scalar (t_b NoUniqueness)
+                      `arr` ( Scalar (t_b NoUniqueness)
+                                `arr` Scalar (t_a Nonunique)
+                            ),
+                    Scalar (t_a Observe)
+                  ]
+                  $ RetType []
+                  $ Scalar (t_b Nonunique)
               )
             ]
               ++
@@ -1214,6 +1228,10 @@
 -- determine whether a 'VName' refers to an intrinsic or a user-defined name.
 maxIntrinsicTag :: Int
 maxIntrinsicTag = maxinum $ map baseTag $ M.keys intrinsics
+
+-- | Is this the name of an intrinsic?
+isIntrinsic :: VName -> Bool
+isIntrinsic = (<= maxIntrinsicTag) . baseTag
 
 -- | Create a name with no qualifiers from a name.
 qualName :: v -> QualName v
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
@@ -390,7 +390,7 @@
   (v', env) <- lookupMod loc v
   when
     ( baseName (qualLeaf v') == nameFromString "intrinsics"
-        && baseTag (qualLeaf v') <= maxIntrinsicTag
+        && isIntrinsic (qualLeaf v')
     )
     $ typeError loc mempty "The 'intrinsics' module may not be used in module expressions."
   pure (mempty, MTy mempty env, ModVar v' loc)
diff --git a/src/Language/Futhark/TypeChecker/Consumption.hs b/src/Language/Futhark/TypeChecker/Consumption.hs
--- a/src/Language/Futhark/TypeChecker/Consumption.hs
+++ b/src/Language/Futhark/TypeChecker/Consumption.hs
@@ -36,14 +36,17 @@
 -- an equivalence class.  See uniqueness-error18.fut for an example of
 -- why this is necessary.
 data Alias
-  = AliasBound {aliasVar :: VName}
-  | AliasFree {aliasVar :: VName}
+  = AliasBound {aliasVar :: VName, aliasFields :: [Name]}
+  | AliasFree {aliasVar :: VName, aliasFields :: [Name]}
   deriving (Eq, Ord, Show)
 
 instance Pretty Alias where
-  pretty (AliasBound v) = prettyName v
-  pretty (AliasFree v) = "~" <> prettyName v
+  pretty (AliasBound v fs) = prettyAlias v fs
+  pretty (AliasFree v fs) = "~" <> prettyAlias v fs
 
+prettyAlias :: VName -> [Name] -> Doc ann
+prettyAlias v fs = prettyName v <> mconcat (map (("." <>) . prettyName) fs)
+
 instance Pretty (S.Set Alias) where
   pretty = braces . commasep . map pretty . S.toList
 
@@ -76,6 +79,20 @@
 aliases :: TypeAliases -> Aliases
 aliases = bifoldMap (const mempty) id
 
+selfAliasType :: VName -> TypeBase Size asf -> TypeAliases
+selfAliasType v = insertSelfAliases v . second (const mempty)
+
+insertSelfAliases :: VName -> TypeAliases -> TypeAliases
+insertSelfAliases v = onPath []
+  where
+    onPath fs (Array als shape et) = Array (S.insert (AliasBound v fs) als) shape et
+    onPath fs (Scalar st) = Scalar $ onPath' fs st
+    onPath' fs (TypeVar als tn args) = TypeVar (S.insert (AliasBound v fs) als) tn args
+    onPath' fs (Record ts) = Record $ M.mapWithKey (\f -> onPath (fs ++ [f])) ts
+    onPath' fs (Sum cs) = Sum $ fmap (map (onPath fs)) cs
+    onPath' fs (Arrow als mn d ps rt) = Arrow (S.insert (AliasBound v fs) als) mn d ps rt
+    onPath' _ et@Prim {} = et
+
 updateAliases :: TypeAliases -> [UpdateStep Info VName] -> TypeAliases -> TypeAliases
 updateAliases _ [] ve_als =
   ve_als
@@ -196,7 +213,7 @@
     checkReturnAlias' params' seen (Unique, names) = do
       when (any (`S.member` S.map snd seen) $ S.toList names) $
         uniqueReturnAliased loc
-      notAliasesParam params' names
+      notAliasesParam params' $ S.map aliasVar names
       pure $ seen `S.union` tag Unique names
     checkReturnAlias' _ seen (Nonunique, names) = do
       when (any (`S.member` seen) $ S.toList $ tag Unique names) $
@@ -218,7 +235,7 @@
     returnAliases (Scalar (Record ets1)) (Scalar (Record ets2)) =
       concat $ M.elems $ M.intersectionWith returnAliases ets1 ets2
     returnAliases expected got =
-      [(uniqueness expected, S.map aliasVar $ aliases got)]
+      [(uniqueness expected, aliases got)]
 
     consumableParamType (Array u _ _) = u == Consume
     consumableParamType (Scalar Prim {}) = True
@@ -230,8 +247,8 @@
 unscope :: [VName] -> Aliases -> Aliases
 unscope bound = S.map f
   where
-    f (AliasFree v) = AliasFree v
-    f (AliasBound v) = if v `elem` bound then AliasFree v else AliasBound v
+    f (AliasFree v fs) = AliasFree v fs
+    f (AliasBound v fs) = if v `elem` bound then AliasFree v fs else AliasBound v fs
 
 -- | Figure out the aliases of each bound name in a pattern.
 matchPat :: Pat t -> TypeAliases -> DL.DList (VName, (t, TypeAliases))
@@ -268,7 +285,7 @@
             foldr (uncurry M.insert . f) (envVtable env) (matchPat p t)
         }
       where
-        f (v, (_, als)) = (v, Consumable $ second (S.insert (AliasBound v)) als)
+        f (v, (_, als)) = (v, Consumable $ insertSelfAliases v als)
 
 bindingParam :: Pat ParamType -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
 bindingParam p m = do
@@ -281,8 +298,8 @@
             foldr (uncurry M.insert . f) (envVtable env) (patternMap p)
         }
     f (v, t)
-      | diet t == Consume = (v, Consumable $ t `setAliases` S.singleton (AliasBound v))
-      | otherwise = (v, Nonconsumable $ t `setAliases` S.singleton (AliasBound v))
+      | diet t == Consume = (v, Consumable $ selfAliasType v t)
+      | otherwise = (v, Nonconsumable $ selfAliasType v t)
 
 bindingIdent :: Diet -> Ident StructType -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
 bindingIdent d (Ident v (Info t) _) =
@@ -292,7 +309,7 @@
     d' = case d of
       Consume -> Consumable
       Observe -> Nonconsumable
-    t' = d' $ t `setAliases` S.singleton (AliasBound v)
+    t' = d' $ selfAliasType v t
 
 bindingParams :: [Pat ParamType] -> CheckM (a, TypeAliases) -> CheckM (a, TypeAliases)
 bindingParams params m =
@@ -335,7 +352,7 @@
           Just (Nonconsumable {}) -> True
           Just _ -> False
           Nothing -> True
-      checkIfConsumable (AliasBound v)
+      checkIfConsumable (AliasBound v _)
         | isBad v = do
             v' <- describeVar v
             addError loc mempty . withIndexLink "not-consumable" $
@@ -374,13 +391,15 @@
     isInstantiation vtable =
       any (`M.member` vtable) . fvVars . freeInType
 
-    selfAlias (Array als shape et) = Array (S.insert (AliasBound v) als) shape et
-    selfAlias (Scalar st) = Scalar $ selfAlias' st
-    selfAlias' (TypeVar als tn args) = TypeVar als tn args -- #1675 FIXME
-    selfAlias' (Record fs) = Record $ fmap selfAlias fs
-    selfAlias' (Sum fs) = Sum $ fmap (map selfAlias) fs
-    selfAlias' et@Arrow {} = et
-    selfAlias' et@Prim {} = et
+    selfAlias = onPath []
+      where
+        onPath fs (Array als shape et) = Array (S.insert (AliasBound v fs) als) shape et
+        onPath fs (Scalar st) = Scalar $ onPath' fs st
+        onPath' _ (TypeVar als tn args) = TypeVar als tn args -- #1675 FIXME
+        onPath' fs (Record ts) = Record $ M.mapWithKey (\f -> onPath (fs ++ [f])) ts
+        onPath' fs (Sum cs) = Sum $ fmap (map (onPath fs)) cs
+        onPath' _ et@Arrow {} = et
+        onPath' _ et@Prim {} = et
 
 -- Capture any newly consumed variables that occur during the provided action.
 contain :: CheckM a -> CheckM (a, Consumed)
@@ -494,10 +513,13 @@
       pure $ als <> seen
 
 consumeAsNeeded :: Loc -> ParamType -> TypeAliases -> CheckM ()
-consumeAsNeeded loc (Scalar (Record fs1)) (Scalar (Record fs2)) =
-  sequence_ $ M.elems $ M.intersectionWith (consumeAsNeeded loc) fs1 fs2
-consumeAsNeeded loc pt t =
-  when (diet pt == Consume) $ consumeAliases loc $ aliases t
+consumeAsNeeded loc pt t = consumeAliases loc $ consumeAliasesOf pt t
+  where
+    consumeAliasesOf (Scalar (Record fs1)) (Scalar (Record fs2)) =
+      mconcat $ M.elems $ M.intersectionWith consumeAliasesOf fs1 fs2
+    consumeAliasesOf p_t t_als
+      | diet p_t == Consume = aliases t_als
+      | otherwise = mempty
 
 checkArg :: [(Exp, TypeAliases)] -> ParamType -> Exp -> CheckM (Exp, TypeAliases)
 checkArg prev p_t e = do
@@ -560,6 +582,15 @@
   returnType closure_als rettype d arg_als
 applyArg t _ = error $ "applyArg: " <> show t
 
+applyLoopArg :: Aliases -> ParamType -> TypeAliases -> ResType -> TypeAliases
+applyLoopArg appres (Scalar (Record pfs)) (Scalar (Record afs)) (Scalar (Record rfs)) =
+  Scalar . Record $
+    M.mapWithKey
+      (\k p_t -> applyLoopArg appres p_t (afs M.! k) (rfs M.! k))
+      pfs
+applyLoopArg appres p_t arg_als rettype =
+  returnType appres rettype (diet p_t) arg_als
+
 boundFreeInExp :: Exp -> CheckM (M.Map VName TypeAliases)
 boundFreeInExp e = do
   vtable <- asks envVtable
@@ -710,12 +741,15 @@
   v <- VName "internal_loop_result" <$> incCounter
   modify $ \s -> s {stateNames = M.insert v (NameLoopRes (srclocOf loop_loc)) $ stateNames s}
 
-  let loopt =
-        funType [param'] (RetType [] $ paramToRes param_t)
-          `setAliases` S.singleton (AliasFree v)
+  let loop_als =
+        applyLoopArg
+          (S.singleton (AliasFree v []))
+          param_t
+          arg_als
+          (paramToRes param_t)
   pure
     ( (param', arg', form', body'),
-      applyArg loopt arg_als `combineAliases` body_als
+      loop_als `combineAliases` body_als
     )
 
 checkFuncall ::
@@ -728,7 +762,7 @@
 checkFuncall loc fname f_als arg_als = do
   v <- VName "internal_app_result" <$> incCounter
   modify $ \s -> s {stateNames = M.insert v (NameAppRes fname loc) $ stateNames s}
-  pure $ foldl applyArg (second (S.insert (AliasFree v)) f_als) arg_als
+  pure $ foldl applyArg (second (S.insert (AliasFree v [])) f_als) arg_als
 
 checkExp :: Exp -> CheckM (Exp, TypeAliases)
 -- First we have the complicated cases.
@@ -815,7 +849,7 @@
         pure (CasePat p body' caseloc, body_als)
 
 --
-checkExp (AppExp (LetFun fname (typarams, params, te, Info (RetType ext ret), funbody) letbody loc) appres) = do
+checkExp (AppExp (LetFun fname (typarams, params, retdecl, Info (RetType ext ret), funbody) letbody loc) appres) = do
   ((ret', funbody'), ftype) <- bindingParams params $ do
     -- Throw away the consumption - it can refer only to the parameters
     -- anyway.
@@ -823,13 +857,13 @@
     checkReturnAlias loc params ret funbody_als
     checkGlobalAliases loc params funbody_als
     free_bound <- boundFreeInExp funbody
-    let ret' = inferReturnUniqueness params ret funbody_als
+    let ret' = maybe (inferReturnUniqueness params ret funbody_als) (const ret) retdecl
         als = foldMap aliases (M.elems free_bound)
         ftype = funType params (RetType ext ret') `setAliases` als
     pure ((ret', funbody'), ftype)
   (letbody', letbody_als) <- bindingFun (fst fname) ftype $ checkExp letbody
   pure
-    ( AppExp (LetFun fname (typarams, params, te, Info (RetType ext ret'), funbody') letbody' loc) appres,
+    ( AppExp (LetFun fname (typarams, params, retdecl, Info (RetType ext ret'), funbody') letbody' loc) appres,
       letbody_als
     )
 
@@ -854,7 +888,7 @@
     checkReturnAlias loc params ret body_als
     checkGlobalAliases loc params body_als
     free_bound <- boundFreeInExp e
-    let ret' = inferReturnUniqueness params ret body_als
+    let ret' = maybe (inferReturnUniqueness params ret body_als) (const ret) te
         als = foldMap aliases (M.elems free_bound)
         ftype = funType params (RetType ext ret') `setAliases` als
     pure
diff --git a/src/Language/Futhark/TypeChecker/Names.hs b/src/Language/Futhark/TypeChecker/Names.hs
--- a/src/Language/Futhark/TypeChecker/Names.hs
+++ b/src/Language/Futhark/TypeChecker/Names.hs
@@ -119,7 +119,7 @@
   v' <- checkValName v loc
   case v' of
     QualName (q : _) _
-      | baseTag q <= maxIntrinsicTag -> do
+      | isIntrinsic q -> do
           me <- askImportName
           unless (isBuiltin (includeToFilePath me)) $
             warn loc "Using intrinsic functions directly can easily crash the compiler or result in wrong code generation."
diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs
--- a/src/Language/Futhark/TypeChecker/Terms.hs
+++ b/src/Language/Futhark/TypeChecker/Terms.hs
@@ -1241,7 +1241,7 @@
     check e@(AppExp (BinOp (QualName [] v, _) _ (x, _) _ loc) _)
       | baseName v == "==",
         Array {} <- typeOf x,
-        baseTag v <= maxIntrinsicTag = do
+        isIntrinsic v = do
           warn loc $
             textwrap
               "Comparing arrays with \"==\" is deprecated and will stop working in a future revision of the language."
diff --git a/src/Language/Futhark/TypeChecker/Terms/Monad.hs b/src/Language/Futhark/TypeChecker/Terms/Monad.hs
--- a/src/Language/Futhark/TypeChecker/Terms/Monad.hs
+++ b/src/Language/Futhark/TypeChecker/Terms/Monad.hs
@@ -385,7 +385,7 @@
 
 lookupQualNameEnv :: QualName VName -> TermTypeM TermScope
 lookupQualNameEnv (QualName [q] _)
-  | baseTag q <= maxIntrinsicTag = asks termScope -- Magical intrinsic module.
+  | isIntrinsic q = asks termScope -- Magical intrinsic module.
 lookupQualNameEnv qn@(QualName quals _) = do
   scope <- asks termScope
   descend scope quals