futhark 0.26.1 → 0.26.2

raw patch · 43 files changed

+847/−287 lines, 43 filesdep ~futhark-server

Dependency ranges changed: futhark-server

Files

• CHANGELOG.md +23/−0
• docs/server-protocol.rst +16/−0
• docs/usage.rst +2/−11
• futhark.cabal +2/−2
• prelude/ad.fut +38/−0
• rts/c/backends/c.h +1/−1
• rts/c/backends/hip.h +0/−6
• rts/c/backends/opencl.h +0/−6
• rts/c/cache.h +1/−1
• rts/c/server.h +200/−43
• rts/c/values.h +1/−1
• rts/javascript/server.js +3/−0
• rts/python/server.py +8/−4
• src-testing/Language/Futhark/TypeChecker/ConsumptionTests.hs +3/−3
• src/Futhark/AD/Fwd.hs +6/−0
• src/Futhark/AD/Rev/SOAC.hs +12/−0
• src/Futhark/CodeGen/Backends/GenericC/CLI.hs +28/−2
• src/Futhark/CodeGen/Backends/GenericC/Server.hs +107/−50
• src/Futhark/CodeGen/Backends/GenericPython.hs +23/−12
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +1/−1
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +157/−61
• src/Futhark/Doc/Generator.hs +5/−6
• src/Futhark/IR/SOACS.hs +1/−0
• src/Futhark/IR/SOACS/SOAC.hs +45/−0
• src/Futhark/IR/SOACS/Simplify.hs +5/−0
• src/Futhark/Internalise/Defunctionalise.hs +1/−1
• src/Futhark/Internalise/Exps.hs +12/−5
• src/Futhark/Internalise/Monomorphise.hs +6/−11
• src/Futhark/Internalise/TypesValues.hs +1/−1
• src/Futhark/Optimise/BlkRegTiling.hs +7/−1
• src/Futhark/Optimise/Fusion/GraphRep.hs +1/−0
• src/Futhark/Optimise/Fusion/Screma.hs +9/−2
• src/Futhark/Optimise/Simplify/Rules.hs +1/−0
• src/Futhark/Pass/AD.hs +24/−19
• src/Futhark/Pass/ExtractMulticore.hs +2/−0
• src/Futhark/Transform/FirstOrderTransform.hs +2/−0
• src/Language/Futhark/Interpreter.hs +4/−0
• src/Language/Futhark/Prop.hs +18/−0
• src/Language/Futhark/TypeChecker.hs +1/−1
• src/Language/Futhark/TypeChecker/Consumption.hs +67/−33
• src/Language/Futhark/TypeChecker/Names.hs +1/−1
• src/Language/Futhark/TypeChecker/Terms.hs +1/−1
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +1/−1

@@ -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

@@ -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 ~~~~~~~~~~~~~~~ 

@@ -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

@@ -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

@@ -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

@@ -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) {

@@ -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++) {

@@ -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++) {

@@ -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) {

@@ -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*);

@@ -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.

@@ -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;       }     }

@@ -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):

@@ -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" $

@@ -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 {} =

@@ -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 

@@ -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)

@@ -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           )

@@ -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+                  )                 ]             )           )

@@ -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)

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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)) =

@@ -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))

@@ -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]

@@ -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

@@ -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

@@ -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

@@ -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)     }

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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)

@@ -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

@@ -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."

@@ -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."

@@ -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