futhark 0.25.11 → 0.25.12

raw patch · 77 files changed

+1626/−1366 lines, 77 filesdep ~futhark-manifestPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependency ranges changed: futhark-manifest

API changes (from Hackage documentation)

Files

• CHANGELOG.md +35/−0
• docs/c-api.rst +67/−8
• docs/language-reference.rst +31/−23
• futhark.cabal +2/−3
• prelude/math.fut +6/−0
• rts/c/backends/c.h +1/−0
• rts/c/backends/cuda.h +13/−8
• rts/c/backends/hip.h +11/−6
• rts/c/backends/multicore.h +1/−0
• rts/c/backends/opencl.h +49/−22
• rts/c/context.h +6/−2
• rts/c/scalar.h +31/−1
• rts/c/scalar_f16.h +16/−0
• rts/python/scalar.py +2/−0
• src/Futhark/AD/Derivatives.hs +6/−0
• src/Futhark/CLI/Eval.hs +1/−0
• src/Futhark/CLI/Profile.hs +1/−0
• src/Futhark/CodeGen/Backends/GPU.hs +4/−4
• src/Futhark/CodeGen/Backends/GenericC.hs +1/−1
• src/Futhark/CodeGen/Backends/GenericC/CLI.hs +3/−3
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +12/−12
• src/Futhark/CodeGen/Backends/GenericC/Monad.hs +4/−4
• src/Futhark/CodeGen/Backends/GenericC/Pretty.hs +6/−0
• src/Futhark/CodeGen/Backends/GenericC/Server.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericC/Types.hs +217/−57
• src/Futhark/CodeGen/Backends/GenericPython.hs +14/−14
• src/Futhark/CodeGen/Backends/MulticoreISPC.hs +2/−2
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +1/−1
• src/Futhark/CodeGen/ImpCode.hs +7/−7
• src/Futhark/CodeGen/ImpGen.hs +11/−10
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +4/−4
• src/Futhark/CodeGen/ImpGen/GPU/Block.hs +2/−2
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +2/−2
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +1/−1
• src/Futhark/Compiler.hs +1/−1
• src/Futhark/IR/Mem.hs +154/−141
• src/Futhark/IR/Mem/IxFun.hs +0/−286
• src/Futhark/IR/Mem/LMAD.hs +108/−18
• src/Futhark/IR/Mem/Simplify.hs +8/−9
• src/Futhark/IR/Parse.hs +31/−20
• src/Futhark/IR/Pretty.hs +4/−0
• src/Futhark/IR/Prop/Types.hs +10/−0
• src/Futhark/IR/SOACS/Simplify.hs +4/−5
• src/Futhark/IR/Syntax/Core.hs +9/−0
• src/Futhark/IR/TypeCheck.hs +2/−0
• src/Futhark/Internalise/Entry.hs +59/−6
• src/Futhark/Internalise/Exps.hs +14/−4
• src/Futhark/Internalise/TypesValues.hs +19/−9
• src/Futhark/Optimise/ArrayShortCircuiting.hs +3/−3
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +16/−16
• src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs +4/−4
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +19/−20
• src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +10/−10
• src/Futhark/Optimise/BlkRegTiling.hs +2/−4
• src/Futhark/Optimise/DoubleBuffer.hs +109/−241
• src/Futhark/Optimise/EntryPointMem.hs +3/−3
• src/Futhark/Optimise/Simplify/Rules/Match.hs +14/−5
• src/Futhark/Optimise/TileLoops/Shared.hs +10/−10
• src/Futhark/Pass/ExpandAllocations.hs +240/−156
• src/Futhark/Pass/ExplicitAllocations.hs +73/−98
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +11/−13
• src/Futhark/Pass/ExplicitAllocations/SegOp.hs +9/−9
• src/Futhark/Profile.hs +5/−0
• src/Futhark/Util.hs +13/−5
• src/Futhark/Util/CMath.hs +14/−0
• src/Language/Futhark/Parser/Parser.y +19/−21
• src/Language/Futhark/Pretty.hs +5/−0
• src/Language/Futhark/Primitive.hs +4/−0
• src/Language/Futhark/TypeChecker.hs +1/−0
• src/Language/Futhark/TypeChecker/Consumption.hs +6/−5
• src/Language/Futhark/TypeChecker/Terms.hs +3/−3
• src/Language/Futhark/TypeChecker/Terms/Pat.hs +1/−1
• src/Language/Futhark/TypeChecker/Unify.hs +19/−19
• src/Language/Futhark/Warnings.hs +1/−0
• unittests/Futhark/IR/Mem/IxFunTests.hs +2/−2
• unittests/Futhark/IR/Mem/IxFunWrapper.hs +6/−8
• unittests/Futhark/Internalise/TypesValuesTests.hs +9/−12

@@ -5,6 +5,41 @@ 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.25.12]++### Added++* `f16.copysign`, `f32.copysign`, `f64.copysign`.++* Trailing commas are now allowed for all syntactical elements that+  involve comma-separation. (#2068)++* The C API now allows destruction and construction of sum types (with+  some caveats). (#2074)++* An overall reduction in memory copies, though simplifying the+  internal representation.++### Fixed++* C API would define distinct entry point types for Futhark types that+  differed only in naming of sizes (#2080).++* `==` and `!=` on sum types with array payloads. Constructing them is+  now a bit slower, though. (#2081)++* Somewhat obscure simplification error caused by neglecting to update+  metadata when removing dead scatter outputs.++* Compiler crash due to the type checker forgetting to respect the+  explicitly ascribed non-consuming diet of loop parameters (#2067).++* Size inference did incomplete level/scope checking, which could+  result in circular sizes, which usually manifested as the type+  checker going into an infinite loop (#2073).++* The OpenCL backend now more gracefully handles lack of platform.+ ## [0.25.11]  ### Added

@@ -304,9 +304,10 @@ a type abbreviation is the best way around this.  The API for opaque values is similar to that of arrays, and the same-rules for memory management apply.  You cannot construct them from-scratch, but must obtain them via entry points (or deserialisation,-see :c:func:`futhark_restore_opaque_foo`).+rules for memory management apply. You cannot construct them from+scratch (unless they correspond to records or tuples, see+:ref:`records`), but must obtain them via entry points (or+deserialisation, see :c:func:`futhark_restore_opaque_foo`).  .. c:struct:: futhark_opaque_foo @@ -355,13 +356,16 @@    Futhark program, and compiled with the same version of the Futhark    compiler. +.. _records:+ Records ~~~~~~~  A record is an opaque type (see above) that supports additional functions to *project* individual fields (read their values) and to-construct a value given values for the fields.  An opaque type is a-record if its definition is a record at the Futhark level.+construct a value given values for the fields. An opaque type is a+record if its definition is a record at the Futhark level. Note that a+tuple is simply a record with numeric fields.  The projection and construction functions are equivalent in functionality to writing entry points by hand, and so serve only to@@ -404,6 +408,61 @@    The resulting value *aliases* the record, but has its own lifetime,    and must eventually be freed. +.. _sums:++Sums+~~~~++A sum type is an opaque type (see above) that supports construction+and destruction functions. An opaque type is a sum type if its+definition is a sum type at the Futhark level.++Similarly to records (see :ref:`Records`), this functionality is+equivalent to writing entry points by hand, and have the same+properties regarding lifetimes.++A sum type consists of one or more variants. A value of this type is+always an instance of one of these variants. In the C API, these+variants are numbered from zero. The numbering is given by the order+in which they are represented in the manifest (see :ref:`manifest`),+which is also the order in which their associated functions are+defined in the header file.++For an opaque sum type ``t``, the following function is always+generated.++.. c:function:: int futhark_variant_opaque_t(struct futhark_context *ctx, const struct futhark_opaque_t *v);++   Return the identifying number of the variant of which this sum type+   is an instance (see above). Cannot fail.++For each variant ``foo``, construction and destruction functions are+defined. The following assume ``t`` is defined as ``type t = #foo+([]i32) bool``.++.. c:function:: int futhark_new_opaque_t_foo(struct futhark_context *ctx, struct futhark_opaque_contrived **out, const struct futhark_i32_1d *v0, const bool v1);++   Construct a value of type ``t`` that is an instance of the variant+   ``foo``. Arguments are provided in the same order as in the+   Futhark-level ``foo`` constructr.++   **Beware:** if ``t`` has size parameters that are only used for+   *other* variants than the one that is being instantiated, those+   size parameters will be set to 0. If this is a problem for your+   application, define your own entry point for constructing a value+   with the proper sizes.++.. c:function:: int futhark_destruct_opaque_contrived_foo(struct futhark_context *ctx, struct futhark_i32_1d **v0, bool *v1, const struct futhark_opaque_contrived *obj);++   Extract the payload of variant ``foo`` from the sum value. Despite+   the name, "destruction" does not free the sum type value. The+   extracted values alias the sum value, but has their own lifetime,+   and must eventually be freed.++   **Precondition:** ``t`` must be an instance of the ``foo`` variant,+   which can be determined with :c:func:`futhark_variant_opaque_t`.++ Entry points ------------ @@ -622,9 +681,9 @@ Manifest -------- -The C backends generate a machine-readable *manifest* in JSON format-that describes the API of the compiled Futhark program.  Specifically,-the manifest contains:+When compiling with ``--library``, the C backends generate a+machine-readable *manifest* in JSON format that describes the API of+the compiled Futhark program. Specifically, the manifest contains:  * A mapping from the name of each entry point to: 

@@ -26,6 +26,12 @@ tools.  Documentation comments are only allowed immediately before declarations. +Trailing commas+---------------++All syntactical elements that involve comma-separated sequencing+permit an optional trailing comma.+ Identifiers and Keywords ------------------------ @@ -142,17 +148,17 @@ :ref:`module-system`).  .. productionlist::-   tuple_type: "(" ")" | "(" `type` ("," `type`)+ ")"+   tuple_type: "(" ")" | "(" `type` ("," `type`)+ [","] ")"  A tuple value or type is written as a sequence of comma-separated-values or types enclosed in parentheses.  For example, ``(0, 1)`` is a-tuple value of type ``(i32,i32)``.  The elements of a tuple need not+values or types enclosed in parentheses. For example, ``(0, 1)`` is a+tuple value of type ``(i32,i32)``. The elements of a tuple need not have the same type -- the value ``(false, 1, 2.0)`` is of type-``(bool, i32, f64)``.  A tuple element can also be another tuple, as-in ``((1,2),(3,4))``, which is of type ``((i32,i32),(i32,i32))``.  A-tuple cannot have just one element, but empty tuples are permitted,-although they are not very useful.  Empty tuples are written ``()``-and are of type ``()``.+``(bool, i32, f64)``. A tuple element can also be another tuple, as in+``((1,2),(3,4))``, which is of type ``((i32,i32),(i32,i32))``. A tuple+cannot have just one element, but empty tuples are permitted, although+they are not very useful. Empty tuples are written ``()`` and are of+type ``()``.  .. productionlist::    array_type: "[" [`exp`] "]" `type`@@ -199,12 +205,13 @@ multiple constructors have large payloads.  .. productionlist::-   record_type: "{" "}" | "{" `fieldid` ":" `type` ("," `fieldid` ":" `type`)* "}"+   record_type: "{" "}" | "{" `fieldid` ":" `type` ("," `fieldid` ":" `type`)* [","] "}"  Records are mappings from field names to values, with the field names-known statically.  A tuple behaves in all respects like a record with-numeric field names starting from zero, and vice versa.  It is an-error for a record type to name the same field twice.+known statically. A tuple behaves in all respects like a record with+numeric field names starting from zero, and vice versa. It is an error+for a record type to name the same field twice. A trailing comma is+permitted.  .. productionlist::    type_application: `type` `type_arg` | "*" `type`@@ -458,18 +465,18 @@        : | `charlit`        : | "(" ")"        : | "(" `exp` ")" ("." `fieldid`)*-       : | "(" `exp` ("," `exp`)* ")"+       : | "(" `exp` ("," `exp`)+ [","] ")"        : | "{" "}"-       : | "{" `field` ("," `field`)* "}"-       : | `qualname` "[" `index` ("," `index`)* "]"-       : | "(" `exp` ")" "[" `index` ("," `index`)* "]"+       : | "{" `field` ("," `field`)* [","] "}"+       : | `qualname` `slice`+       : | "(" `exp` ")" `slice`        : | `quals` "." "(" `exp` ")"-       : | "[" `exp` ("," `exp`)* "]"+       : | "[" `exp` ("," `exp`)* [","] "]"        : | "(" `qualsymbol` ")"        : | "(" `exp` `qualsymbol` ")"        : | "(" `qualsymbol` `exp` ")"        : | "(" ( "." `field` )+ ")"-       : | "(" "." "[" `index` ("," `index`)* "]" ")"+       : | "(" "." `slice` ")"        : | "???"    exp:   `atom`       : | `exp` `qualsymbol` `exp`@@ -484,16 +491,17 @@       : | `exp` [ ".." `exp` ] "..>" `exp`       : | "if" `exp` "then" `exp` "else" `exp`       : | "let" `size`* `pat` "=" `exp` "in" `exp`-      : | "let" `name` "[" `index` ("," `index`)* "]" "=" `exp` "in" `exp`+      : | "let" `name` `slice` "=" `exp` "in" `exp`       : | "let" `name` `type_param`* `pat`+ [":" `type`] "=" `exp` "in" `exp`       : | "(" "\" `pat`+ [":" `type`] "->" `exp` ")"       : | "loop" `pat` ["=" `exp`] `loopform` "do" `exp`       : | "#[" `attr` "]" `exp`       : | "unsafe" `exp`       : | "assert" `atom` `atom`-      : | `exp` "with" "[" `index` ("," `index`)* "]" "=" `exp`+      : | `exp` "with" `slice` "=" `exp`       : | `exp` "with" `fieldid` ("." `fieldid`)* "=" `exp`       : | "match" `exp` ("case" `pat` "->" `exp`)++   slice: "[" `index` ("," `index`)* [","] "]"    field:   `fieldid` "=" `exp`         : | `name`    size : "[" `name` "]"@@ -502,9 +510,9 @@       : | "_"       : | "(" ")"       : | "(" `pat` ")"-      : | "(" `pat` ("," `pat`)+ ")"+      : | "(" `pat` ("," `pat`)+ [","] ")"       : | "{" "}"-      : | "{" `fieldid` ["=" `pat`] ("," `fieldid` ["=" `pat`])* "}"+      : | "{" `fieldid` ["=" `pat`] ("," `fieldid` ["=" `pat`])* [","] "}"       : | `constructor` `pat`*       : | `pat` ":" `type`       : | "#[" `attr` "]" `pat`@@ -1686,7 +1694,7 @@ .. productionlist::    attr:   `name`        : | `decimal`-       : | `name` "(" [`attr` ("," `attr`)*] ")"+       : | `name` "(" [`attr` ("," `attr`)* [","]] ")"  An expression, declaration, pattern, or module type spec can be prefixed with an attribute, written as ``#[attr]``.  This may affect

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.11+version:        0.25.12 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -254,7 +254,6 @@       Futhark.IR.MCMem       Futhark.IR.Mem       Futhark.IR.Mem.Interval-      Futhark.IR.Mem.IxFun       Futhark.IR.Mem.LMAD       Futhark.IR.Mem.Simplify       Futhark.IR.Parse@@ -453,7 +452,7 @@     , free >=5.1.10     , futhark-data >= 1.1.0.0     , futhark-server >= 1.2.2.1-    , futhark-manifest >= 1.2.0.1+    , futhark-manifest >= 1.3.0.0     , githash >=0.1.6.1     , half >= 0.3     , haskeline

@@ -257,6 +257,9 @@    -- | Multiplies floating-point value by 2 raised to an integer power.   val ldexp : t -> i32 -> t++  -- | Compose a floating-point value with the magnitude of `x` and the sign of `y`.+  val copysign : (x: t) -> (y: t) -> t }  -- | Boolean numbers.  When converting from a number to `bool`, 0 is@@ -971,6 +974,7 @@    def nextafter x y = intrinsics.nextafter64 (x,y)   def ldexp x y = intrinsics.ldexp64 (x,y)+  def copysign x y = intrinsics.copysign64 (x,y)    def to_bits (x: f64): u64 = u64m.i64 (intrinsics.to_bits64 x)   def from_bits (x: u64): f64 = intrinsics.from_bits64 (intrinsics.sign_i64 x)@@ -1087,6 +1091,7 @@    def nextafter x y = intrinsics.nextafter32 (x,y)   def ldexp x y = intrinsics.ldexp32 (x,y)+  def copysign x y = intrinsics.copysign32 (x,y)    def to_bits (x: f32): u32 = u32m.i32 (intrinsics.to_bits32 x)   def from_bits (x: u32): f32 = intrinsics.from_bits32 (intrinsics.sign_i32 x)@@ -1207,6 +1212,7 @@    def nextafter x y = intrinsics.nextafter16 (x,y)   def ldexp x y = intrinsics.ldexp16 (x,y)+  def copysign x y = intrinsics.copysign16 (x,y)    def to_bits (x: f16): u16 = u16m.i16 (intrinsics.to_bits16 x)   def from_bits (x: u16): f16 = intrinsics.from_bits16 (intrinsics.sign_i16 x)

@@ -43,6 +43,7 @@   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;   struct program* program;+  bool program_initialised; };  int backend_context_setup(struct futhark_context* ctx) {

@@ -273,6 +273,8 @@   struct event_list event_list;   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;+  struct program* program;+  bool program_initialised;   // Uniform fields above.    CUdeviceptr global_failure;@@ -284,7 +286,6 @@   long int total_runtime;   int64_t peak_mem_usage_device;   int64_t cur_mem_usage_device;-  struct program* program;    CUdevice dev;   CUcontext cu_ctx;@@ -819,6 +820,7 @@   ctx->total_runtime = 0;   ctx->peak_mem_usage_device = 0;   ctx->cur_mem_usage_device = 0;+  ctx->kernels = NULL;    CUDA_SUCCEED_FATAL(cuInit(0));   if (cuda_device_setup(ctx) != 0) {@@ -868,13 +870,16 @@ }  void backend_context_teardown(struct futhark_context* ctx) {-  free_builtin_kernels(ctx, ctx->kernels);-  cuMemFree(ctx->global_failure);-  cuMemFree(ctx->global_failure_args);-  CUDA_SUCCEED_FATAL(gpu_free_all(ctx));-  CUDA_SUCCEED_FATAL(cuStreamDestroy(ctx->stream));-  CUDA_SUCCEED_FATAL(cuModuleUnload(ctx->module));-  CUDA_SUCCEED_FATAL(cuCtxDestroy(ctx->cu_ctx));+  if (ctx->kernels != NULL) {+    free_builtin_kernels(ctx, ctx->kernels);+    cuMemFree(ctx->global_failure);+    cuMemFree(ctx->global_failure_args);+    CUDA_SUCCEED_FATAL(gpu_free_all(ctx));+    CUDA_SUCCEED_FATAL(cuStreamDestroy(ctx->stream));+    CUDA_SUCCEED_FATAL(cuModuleUnload(ctx->module));+    CUDA_SUCCEED_FATAL(cuCtxDestroy(ctx->cu_ctx));+  }+  free_list_destroy(&ctx->gpu_free_list); }  // GPU ABSTRACTION LAYER

@@ -248,6 +248,7 @@   struct event_list event_list;   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;+  bool program_initialised;   // Uniform fields above.    void* global_failure;@@ -676,6 +677,7 @@   ctx->total_runtime = 0;   ctx->peak_mem_usage_device = 0;   ctx->cur_mem_usage_device = 0;+  ctx->kernels = NULL;    HIP_SUCCEED_FATAL(hipInit(0));   if (hip_device_setup(ctx) != 0) {@@ -724,12 +726,15 @@ }  void backend_context_teardown(struct futhark_context* ctx) {-  free_builtin_kernels(ctx, ctx->kernels);-  hipFree(ctx->global_failure);-  hipFree(ctx->global_failure_args);-  HIP_SUCCEED_FATAL(gpu_free_all(ctx));-  HIP_SUCCEED_FATAL(hipStreamDestroy(ctx->stream));-  HIP_SUCCEED_FATAL(hipModuleUnload(ctx->module));+  if (ctx->kernels != NULL) {+    free_builtin_kernels(ctx, ctx->kernels);+    hipFree(ctx->global_failure);+    hipFree(ctx->global_failure_args);+    HIP_SUCCEED_FATAL(gpu_free_all(ctx));+    HIP_SUCCEED_FATAL(hipStreamDestroy(ctx->stream));+    HIP_SUCCEED_FATAL(hipModuleUnload(ctx->module));+  }+  free_list_destroy(&ctx->gpu_free_list); }  // GPU ABSTRACTION LAYER

@@ -50,6 +50,7 @@   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;   struct program* program;+  bool program_initialised;   // Uniform fields above.    lock_t event_list_lock;

@@ -513,8 +513,8 @@   int64_t peak_mem_usage_default;   int64_t cur_mem_usage_default;   struct program* program;--  // Common fields above.+  bool program_initialised;+  // Uniform fields above.    cl_mem global_failure;   cl_mem global_failure_args;@@ -546,7 +546,7 @@   struct builtin_kernels* kernels; }; -static cl_build_status build_gpu_program(cl_program program, cl_device_id device, const char* options) {+static cl_build_status build_gpu_program(cl_program program, cl_device_id device, const char* options, char** log) {   cl_int clBuildProgram_error = clBuildProgram(program, 1, &device, options, NULL, NULL);    // Avoid termination due to CL_BUILD_PROGRAM_FAILURE@@ -563,7 +563,7 @@                                              &build_status,                                              NULL)); -  if (build_status != CL_SUCCESS) {+  if (build_status != CL_BUILD_SUCCESS) {     char *build_log;     size_t ret_val_size;     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &ret_val_size));@@ -571,12 +571,10 @@     build_log = (char*) malloc(ret_val_size+1);     OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL)); -    // The spec technically does not say whether the build log is zero-terminated, so let's be careful.+    // The spec technically does not say whether the build log is+    // zero-terminated, so let's be careful.     build_log[ret_val_size] = '\0';--    fprintf(stderr, "Build log:\n%s\n", build_log);--    free(build_log);+    *log = build_log;   }    return build_status;@@ -977,10 +975,25 @@   if (ctx->cfg->logging) {     fprintf(stderr, "Building OpenCL program...\n");   }-  OPENCL_SUCCEED_FATAL(build_gpu_program(prog, device_option.device, compile_opts));-+  char* build_log;+  cl_build_status status =+    build_gpu_program(prog, device_option.device, compile_opts, &build_log);   free(compile_opts); +  if (status != CL_BUILD_SUCCESS) {+    ctx->error = msgprintf("Compilation of OpenCL program failed.\nBuild log:\n%s",+                           build_log);+    // We are giving up on initialising this OpenCL context. That also+    // means we need to free all the OpenCL bits we have managed to+    // allocate thus far, as futhark_context_free() will not touch+    // these unless initialisation was completely successful.+    (void)clReleaseProgram(prog);+    (void)clReleaseCommandQueue(ctx->queue);+    (void)clReleaseContext(ctx->ctx);+    free(build_log);+    return;+  }+   size_t binary_size = 0;   unsigned char *binary = NULL;   int store_in_cache = cache_fname != NULL && !loaded_from_cache;@@ -1011,7 +1024,8 @@   ctx->clprogram = prog; } -static struct opencl_device_option get_preferred_device(const struct futhark_context_config *cfg) {+static struct opencl_device_option get_preferred_device(struct futhark_context *ctx,+                                                        const struct futhark_context_config *cfg) {   struct opencl_device_option *devices;   size_t num_devices; @@ -1038,15 +1052,20 @@     }   } -  futhark_panic(1, "Could not find acceptable OpenCL device.\n");-  exit(1); // Never reached+  ctx->error = strdup("Could not find acceptable OpenCL device.\n");+  struct opencl_device_option device;+  return device; }  static void setup_opencl(struct futhark_context *ctx,                          const char *extra_build_opts[],                          const char* cache_fname) {-  struct opencl_device_option device_option = get_preferred_device(ctx->cfg);+  struct opencl_device_option device_option = get_preferred_device(ctx, ctx->cfg); +  if (ctx->error != NULL) {+    return;+  }+   if (ctx->cfg->logging) {     fprintf(stderr, "Using platform: %s\n", device_option.platform_name);     fprintf(stderr, "Using device: %s\n", device_option.device_name);@@ -1084,6 +1103,7 @@   ctx->total_runtime = 0;   ctx->peak_mem_usage_device = 0;   ctx->cur_mem_usage_device = 0;+  ctx->kernels = NULL;    if (ctx->cfg->queue_set) {     setup_opencl_with_command_queue(ctx, ctx->cfg->queue, (const char**)ctx->cfg->build_opts, ctx->cfg->cache_fname);@@ -1091,6 +1111,10 @@     setup_opencl(ctx, (const char**)ctx->cfg->build_opts, ctx->cfg->cache_fname);   } +  if (ctx->error != NULL) {+    return 1;+  }+   cl_int error;   cl_int no_error = -1;   ctx->global_failure =@@ -1116,13 +1140,16 @@ static int gpu_free_all(struct futhark_context *ctx);  void backend_context_teardown(struct futhark_context* ctx) {-  free_builtin_kernels(ctx, ctx->kernels);-  OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure));-  OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure_args));-  (void)gpu_free_all(ctx);-  (void)clReleaseProgram(ctx->clprogram);-  (void)clReleaseCommandQueue(ctx->queue);-  (void)clReleaseContext(ctx->ctx);+  if (ctx->kernels != NULL) {+    free_builtin_kernels(ctx, ctx->kernels);+    OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure));+    OPENCL_SUCCEED_FATAL(clReleaseMemObject(ctx->global_failure_args));+    (void)gpu_free_all(ctx);+    (void)clReleaseProgram(ctx->clprogram);+    (void)clReleaseCommandQueue(ctx->queue);+    (void)clReleaseContext(ctx->ctx);+  }+  free_list_destroy(&ctx->gpu_free_list); }  cl_command_queue futhark_context_get_command_queue(struct futhark_context* ctx) {

@@ -116,6 +116,7 @@   assert(!cfg->in_use);   ctx->cfg = cfg;   ctx->cfg->in_use = 1;+  ctx->program_initialised = false;   create_lock(&ctx->error_lock);   create_lock(&ctx->lock);   free_list_init(&ctx->free_list);@@ -134,6 +135,7 @@   if (backend_context_setup(ctx) == 0) {     setup_program(ctx);     init_constants(ctx);+    ctx->program_initialised = true;     (void)futhark_context_clear_caches(ctx);     (void)futhark_context_sync(ctx);   }@@ -141,8 +143,10 @@ }  void futhark_context_free(struct futhark_context* ctx) {-  free_constants(ctx);-  teardown_program(ctx);+  if (ctx->program_initialised) {+    free_constants(ctx);+    teardown_program(ctx);+  }   backend_context_teardown(ctx);   free_all_in_free_list(ctx);   free_list_destroy(&ctx->free_list);

@@ -17,6 +17,9 @@ // Double-precision definitions are only included if the preprocessor // macro FUTHARK_F64_ENABLED is set. +SCALAR_FUN_ATTR int32_t futrts_to_bits32(float x);+SCALAR_FUN_ATTR float futrts_from_bits32(int32_t x);+ SCALAR_FUN_ATTR uint8_t add8(uint8_t x, uint8_t y) {   return x + y; }@@ -1890,6 +1893,10 @@   return ldexp(x, y); } +SCALAR_FUN_ATTR float futrts_copysign32(float x, float y) {+  return copysign(x, y);+}+ SCALAR_FUN_ATTR float futrts_mad32(float a, float b, float c) {   return mad(a, b, c); }@@ -2103,6 +2110,12 @@   return x * pow((double)2.0, (double)y); } +SCALAR_FUN_ATTR float futrts_copysign32(float x, float y) {+  int32_t xb = futrts_to_bits32(x);+  int32_t yb = futrts_to_bits32(y);+  return futrts_from_bits32((xb & ~(1<<31)) | (yb & (1<<31)));+}+ SCALAR_FUN_ATTR float futrts_mad32(float a, float b, float c) {   return a * b + c; }@@ -2241,6 +2254,10 @@   return ldexpf(x, y); } +SCALAR_FUN_ATTR float futrts_copysign32(float x, float y) {+  return copysignf(x, y);+}+ SCALAR_FUN_ATTR float futrts_mad32(float a, float b, float c) {   return a * b + c; }@@ -2286,6 +2303,9 @@  #ifdef FUTHARK_F64_ENABLED +SCALAR_FUN_ATTR double futrts_from_bits64(int64_t x);+SCALAR_FUN_ATTR int64_t futrts_to_bits64(double x);+ #if ISPC SCALAR_FUN_ATTR bool futrts_isinf64(float x) {   return !isnan(x) && isnan(x - x);@@ -2652,10 +2672,16 @@   return v0 + (v1 - v0) * t; } -SCALAR_FUN_ATTR float futrts_ldexp64(double x, int32_t y) {+SCALAR_FUN_ATTR double futrts_ldexp64(double x, int32_t y) {   return x * pow((double)2.0, (double)y); } +SCALAR_FUN_ATTR double futrts_copysign64(double x, double y) {+  int64_t xb = futrts_to_bits64(x);+  int64_t yb = futrts_to_bits64(y);+  return futrts_from_bits64((xb & ~(((int64_t)1)<<63)) | (yb & (((int64_t)1)<<63)));+}+ SCALAR_FUN_ATTR double futrts_mad64(double a, double b, double c) {   return a * b + c; }@@ -2988,6 +3014,10 @@  SCALAR_FUN_ATTR double futrts_ldexp64(double x, int32_t y) {   return ldexp(x, y);+}++SCALAR_FUN_ATTR float futrts_copysign64(double x, double y) {+  return copysign(x, y); }  SCALAR_FUN_ATTR double futrts_mad64(double a, double b, double c) {

@@ -334,6 +334,10 @@   return ldexp(x, y); } +SCALAR_FUN_ATTR f16 futrts_copysign16(f16 x, f16 y) {+  return copysign(x, y);+}+ SCALAR_FUN_ATTR f16 futrts_mad16(f16 a, f16 b, f16 c) {   return mad(a, b, c); }@@ -495,6 +499,10 @@   return futrts_ldexp32((float)x, y); } +SCALAR_FUN_ATTR f16 futrts_copysign16(f16 x, f16 y) {+  return futrts_copysign32((float)x, y);+}+ SCALAR_FUN_ATTR f16 futrts_mad16(f16 a, f16 b, f16 c) {   return a * b + c; }@@ -633,6 +641,10 @@   return futrts_ldexp32((float)x, y); } +SCALAR_FUN_ATTR f16 futrts_copysign16(f16 x, f16 y) {+  return futrts_copysign32((float)x, y);+}+ SCALAR_FUN_ATTR f16 futrts_mad16(f16 a, f16 b, f16 c) {   return a * b + c; }@@ -837,6 +849,10 @@  SCALAR_FUN_ATTR f16 futrts_ldexp16(f16 x, int32_t y) {   return futrts_ldexp32(x, y);+}++SCALAR_FUN_ATTR f16 futrts_copysign16(f16 x, f16 y) {+  return futrts_copysign32((float)x, y); }  SCALAR_FUN_ATTR f16 futrts_mad16(f16 a, f16 b, f16 c) {

@@ -1026,4 +1026,6 @@     return a * b + c  +futhark_copysign16 = futhark_copysign32 = futhark_copysign64 = np.copysign+ # End of scalar.py.

@@ -366,6 +366,12 @@   Just [untyped $ negate $ (2 / sqrt pi) * exp (negate (isF32 z * isF32 z))] pdBuiltin "erfc64" [z] =   Just [untyped $ negate $ (2 / sqrt pi) * exp (negate (isF64 z * isF64 z))]+pdBuiltin "copysign16" [_x, y] =+  Just [untyped $ 1 * isF16 (UnOpExp (FSignum Float16) y), fConst Float16 0]+pdBuiltin "copysign32" [_x, y] =+  Just [untyped $ 1 * isF32 (UnOpExp (FSignum Float32) y), fConst Float32 0]+pdBuiltin "copysign64" [_x, y] =+  Just [untyped $ 1 * isF64 (UnOpExp (FSignum Float64) y), fConst Float64 0] -- More problematic derivatives follow below. pdBuiltin "umul_hi8" [x, y] = Just [y, x] pdBuiltin "umul_hi16" [x, y] = Just [y, x]

@@ -1,3 +1,4 @@+-- | @futhark eval@ module Futhark.CLI.Eval (main) where  import Control.Exception

@@ -1,3 +1,4 @@+-- | @futhark profile@ module Futhark.CLI.Profile (main) where  import Control.Exception (catch)

@@ -155,7 +155,7 @@ callKernel (LaunchKernel safety kernel_name shared_memory args num_tblocks tblock_size) =   genLaunchKernel safety kernel_name shared_memory args num_tblocks tblock_size -copygpu2gpu :: GC.DoLMADCopy op s+copygpu2gpu :: GC.DoCopy op s copygpu2gpu _ t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do   let fname = "lmad_copy_gpu2gpu_" <> show (primByteSize t :: Int) <> "b"       r = length shape@@ -175,7 +175,7 @@          }      |] -copyhost2gpu :: GC.DoLMADCopy op s+copyhost2gpu :: GC.DoCopy op s copyhost2gpu sync t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do   let r = length shape       dststride_inits = [[C.cinit|$exp:e|] | Count e <- dststride]@@ -199,7 +199,7 @@       GC.CopyBarrier -> [C.cexp|true|]       GC.CopyNoBarrier -> [C.cexp|false|] -copygpu2host :: GC.DoLMADCopy op s+copygpu2host :: GC.DoCopy op s copygpu2host sync t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do   let r = length shape       dststride_inits = [[C.cinit|$exp:e|] | Count e <- dststride]@@ -223,7 +223,7 @@       GC.CopyBarrier -> [C.cexp|true|]       GC.CopyNoBarrier -> [C.cexp|false|] -gpuCopies :: M.Map (Space, Space) (GC.DoLMADCopy op s)+gpuCopies :: M.Map (Space, Space) (GC.DoCopy op s) gpuCopies =   M.fromList     [ ((Space "device", Space "device"), copygpu2gpu),

@@ -62,7 +62,7 @@               err = FUTHARK_PROGRAM_ERROR;               goto cleanup;|] -lmadcopyCPU :: DoLMADCopy op s+lmadcopyCPU :: DoCopy op s lmadcopyCPU _ t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do   let fname :: String       (fname, ty) =

@@ -190,7 +190,7 @@             [C.cstm|;|],             [C.cexp|$id:dest|]           )-    Just (TypeOpaque desc _ _) ->+    Just (TypeOpaque desc _ _ _) ->       ( [C.citems|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:(T.unpack desc));|],         [C.cstm|;|],         [C.cstm|;|],@@ -256,7 +256,7 @@             [C.cexp|$id:result|],             [C.cstm|assert($id:(arrayFree ops)(ctx, $id:result) == 0);|]           )-        Just (TypeOpaque t ops _) ->+        Just (TypeOpaque t ops _ _) ->           ( [C.citem|typename $id:t $id:result;|],             [C.cexp|$id:result|],             [C.cstm|assert($id:(opaqueFree ops)(ctx, $id:result) == 0);|]@@ -269,7 +269,7 @@     Nothing ->       let info = tname <> "_info"        in [C.cstm|write_scalar(stdout, binary_output, &$id:info, &$exp:e);|]-    Just (TypeOpaque desc _ _) ->+    Just (TypeOpaque desc _ _ _) ->       [C.cstm|{          fprintf(stderr, "Values of type \"%s\" have no external representation.\n", $string:(T.unpack desc));          retval = 1;

@@ -7,8 +7,8 @@     compileCode,     compileDest,     compileArg,-    compileLMADCopy,-    compileLMADCopyWith,+    compileCopy,+    compileCopyWith,     errorMsgString,     linearCode,   )@@ -326,11 +326,11 @@       [C.cstm|if ($exp:cond') { $items:tbranch' } else $stm:x|]     _ ->       [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|]-compileCode (LMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) = do+compileCode (Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) = do   cp <- asks $ M.lookup (dstspace, srcspace) . opsCopies . envOperations   case cp of     Nothing ->-      compileLMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)+      compileCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)     Just cp' -> do       shape' <- traverse (traverse (compileExp . untyped)) shape       dst' <- rawMem dst@@ -396,9 +396,9 @@       <*> mapM compileArg args   stms $ mconcat unpack_dest --- | Compile an 'LMADCopy' using sequential nested loops, but+-- | Compile an 'Copy' using sequential nested loops, but -- parameterised over how to do the reads and writes.-compileLMADCopyWith ::+compileCopyWith ::   [Count Elements (TExp Int64)] ->   (C.Exp -> C.Exp -> CompilerM op s ()) ->   ( Count Elements (TExp Int64),@@ -409,7 +409,7 @@     [Count Elements (TExp Int64)]   ) ->   CompilerM op s ()-compileLMADCopyWith shape doWrite dst_lmad doRead src_lmad = do+compileCopyWith shape doWrite dst_lmad doRead src_lmad = do   let (dstoffset, dststrides) = dst_lmad       (srcoffset, srcstrides) = src_lmad   shape' <- mapM (compileExp . untyped . unCount) shape@@ -432,10 +432,10 @@       [C.citems|for (typename int64_t $id:i = 0; $id:i < $exp:n; $id:i++)                   { $items:(loops ins body) }|] --- | Compile an 'LMADCopy' using sequential nested loops and--- 'Read'/'Write' of individual scalars.  This always works, but can+-- | Compile an 'Copy' using sequential nested loops and+-- t'Read'/t'Write' of individual scalars.  This always works, but can -- be pretty slow if those reads and writes are costly.-compileLMADCopy ::+compileCopy ::   PrimType ->   [Count Elements (TExp Int64)] ->   (VName, Space) ->@@ -447,9 +447,9 @@     [Count Elements (TExp Int64)]   ) ->   CompilerM op s ()-compileLMADCopy t shape (dst, dstspace) dst_lmad (src, srcspace) src_lmad = do+compileCopy t shape (dst, dstspace) dst_lmad (src, srcspace) src_lmad = do   src' <- rawMem src   dst' <- rawMem dst   let doWrite dst_i = generateWrite dst' dst_i t dstspace Nonvolatile       doRead src_i = generateRead src' src_i t srcspace Nonvolatile-  compileLMADCopyWith shape doWrite dst_lmad doRead src_lmad+  compileCopyWith shape doWrite dst_lmad doRead src_lmad

@@ -18,7 +18,7 @@     Deallocate,     CopyBarrier (..),     Copy,-    DoLMADCopy,+    DoCopy,      -- * Monadic compiler interface     CompilerM,@@ -201,9 +201,9 @@   C.Exp ->   CompilerM op s () --- | Perform an 'LMADCopy'.  It is expected that these functions are+-- | Perform an 'Copy'.  It is expected that these functions are -- each specialised on which spaces they operate on, so that is not part of their arguments.-type DoLMADCopy op s =+type DoCopy op s =   CopyBarrier ->   PrimType ->   [Count Elements C.Exp] ->@@ -231,7 +231,7 @@     opsError :: ErrorCompiler op s,     opsCall :: CallCompiler op s,     -- | @(dst,src)@-space mapping to copy functions.-    opsCopies :: M.Map (Space, Space) (DoLMADCopy op s),+    opsCopies :: M.Map (Space, Space) (DoCopy op s),     -- | If true, use reference counting.  Otherwise, bare     -- pointers.     opsFatMemory :: Bool,

@@ -18,20 +18,26 @@ import Text.PrettyPrint.Mainland qualified as MPP import Text.PrettyPrint.Mainland.Class qualified as MPP +-- | Prettyprint a C expression. expText :: C.Exp -> T.Text expText = T.pack . MPP.pretty 8000 . MPP.ppr +-- | Prettyprint a list of C definitions. definitionsText :: [C.Definition] -> T.Text definitionsText = T.unlines . map (T.pack . MPP.pretty 8000 . MPP.ppr) +-- | Prettyprint a single C type. typeText :: C.Type -> T.Text typeText = T.pack . MPP.pretty 8000 . MPP.ppr +-- | Prettyprint a single identifier. idText :: C.Id -> T.Text idText = T.pack . MPP.pretty 8000 . MPP.ppr +-- | Prettyprint a single function. funcText :: C.Func -> T.Text funcText = T.pack . MPP.pretty 8000 . MPP.ppr +-- | Prettyprint a list of functions. funcsText :: [C.Func] -> T.Text funcsText = T.unlines . map funcText

@@ -119,7 +119,7 @@ cType manifest tname =   case M.lookup tname $ manifestTypes manifest of     Just (TypeArray ctype _ _ _) -> [C.cty|typename $id:(T.unpack ctype)|]-    Just (TypeOpaque ctype _ _) -> [C.cty|typename $id:(T.unpack ctype)|]+    Just (TypeOpaque ctype _ _ _) -> [C.cty|typename $id:(T.unpack ctype)|]     Nothing -> uncurry primAPIType $ scalarToPrim tname  -- First component is forward declaration so we don't have to worry@@ -156,7 +156,7 @@                 .aux = &$id:aux_name               };|]       )-typeBoilerplate manifest (tname, TypeOpaque c_type_name ops record) =+typeBoilerplate manifest (tname, TypeOpaque c_type_name ops record _sumops) =   let type_name = typeStructName tname       aux_name = type_name <> "_aux"       (record_edecls, record_init) = recordDefs type_name record

@@ -204,6 +204,7 @@  opaquePayload :: OpaqueTypes -> OpaqueType -> [ValueType] opaquePayload _ (OpaqueType ts) = ts+opaquePayload _ (OpaqueSum ts _) = ts opaquePayload types (OpaqueRecord fs) = concatMap f fs   where     f (_, TypeOpaque s) = opaquePayload types $ lookupOpaqueType s types@@ -226,56 +227,62 @@     typeLength (TypeOpaque desc) =       length $ opaquePayload types $ lookupOpaqueType desc types -opaqueProjectFunctions ::+projectField ::+  Operations op s ->+  EntryPointType ->+  [(Int, ValueType)] ->+  CompilerM op s (C.Type, [C.BlockItem])+projectField _ (TypeTransparent (ValueType sign (Rank 0) pt)) [(i, _)] = do+  pure+    ( primAPIType sign pt,+      [C.citems|v = obj->$id:(tupleField i);|]+    )+projectField ops (TypeTransparent vt) [(i, _)] = do+  ct <- valueTypeToCType Public vt+  pure+    ( [C.cty|$ty:ct *|],+      criticalSection+        ops+        [C.citems|v = malloc(sizeof($ty:ct));+                  memcpy(v, obj->$id:(tupleField i), sizeof($ty:ct));+                  (void)(*(v->mem.references))++;|]+    )+projectField _ (TypeTransparent _) rep =+  error $ "projectField: invalid representation of transparent type: " ++ show rep+projectField ops (TypeOpaque f_desc) components = do+  ct <- opaqueToCType f_desc+  let setField j (i, ValueType _ (Rank r) _) =+        if r == 0+          then [C.citems|v->$id:(tupleField j) = obj->$id:(tupleField i);|]+          else+            [C.citems|v->$id:(tupleField j) = malloc(sizeof(*v->$id:(tupleField j)));+                      *v->$id:(tupleField j) = *obj->$id:(tupleField i);+                      (void)(*(v->$id:(tupleField j)->mem.references))++;|]+  pure+    ( [C.cty|$ty:ct *|],+      criticalSection+        ops+        [C.citems|v = malloc(sizeof($ty:ct));+                  $items:(concat (zipWith setField [0..] components))|]+    )++recordProjectFunctions ::   OpaqueTypes ->   Name ->   [(Name, EntryPointType)] ->   [ValueType] ->   CompilerM op s [Manifest.RecordField]-opaqueProjectFunctions types desc fs vds = do+recordProjectFunctions types desc fs vds = do   opaque_type <- opaqueToCType desc   ctx_ty <- contextType   ops <- asks envOperations-  let mkProject (TypeTransparent (ValueType sign (Rank 0) pt)) [(i, _)] = do-        pure-          ( primAPIType sign pt,-            [C.citems|v = obj->$id:(tupleField i);|]-          )-      mkProject (TypeTransparent vt) [(i, _)] = do-        ct <- valueTypeToCType Public vt-        pure-          ( [C.cty|$ty:ct *|],-            criticalSection-              ops-              [C.citems|v = malloc(sizeof($ty:ct));-                        memcpy(v, obj->$id:(tupleField i), sizeof($ty:ct));-                        (void)(*(v->mem.references))++;|]-          )-      mkProject (TypeTransparent _) rep =-        error $ "mkProject: invalid representation of transparent type: " ++ show rep-      mkProject (TypeOpaque f_desc) components = do-        ct <- opaqueToCType f_desc-        let setField j (i, ValueType _ (Rank r) _) =-              if r == 0-                then [C.citems|v->$id:(tupleField j) = obj->$id:(tupleField i);|]-                else-                  [C.citems|v->$id:(tupleField j) = malloc(sizeof(*v->$id:(tupleField j)));-                            *v->$id:(tupleField j) = *obj->$id:(tupleField i);-                            (void)(*(v->$id:(tupleField j)->mem.references))++;|]-        pure-          ( [C.cty|$ty:ct *|],-            criticalSection-              ops-              [C.citems|v = malloc(sizeof($ty:ct));-                        $items:(concat (zipWith setField [0..] components))|]-          )   let onField ((f, et), elems) = do         let f' =               if isValidCName $ opaqueName desc <> "_" <> nameToText f                 then nameToText f                 else zEncodeText (nameToText f)         project <- publicName $ "project_" <> opaqueName desc <> "_" <> f'-        (et_ty, project_items) <- mkProject et elems+        (et_ty, project_items) <- projectField ops et elems         headerDecl           (OpaqueDecl desc)           [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj);|]@@ -292,13 +299,22 @@   mapM onField . zip fs . recordFieldPayloads types (map snd fs) $     zip [0 ..] vds -opaqueNewFunctions ::+setFieldField :: (C.ToExp a) => Int -> a -> ValueType -> C.Stm+setFieldField i e (ValueType _ (Rank r) _)+  | r == 0 =+      [C.cstm|v->$id:(tupleField i) = $exp:e;|]+  | otherwise =+      [C.cstm|{v->$id:(tupleField i) = malloc(sizeof(*$exp:e));+               *v->$id:(tupleField i) = *$exp:e;+               (void)(*(v->$id:(tupleField i)->mem.references))++;}|]++recordNewFunctions ::   OpaqueTypes ->   Name ->   [(Name, EntryPointType)] ->   [ValueType] ->   CompilerM op s Manifest.CFuncName-opaqueNewFunctions types desc fs vds = do+recordNewFunctions types desc fs vds = do   opaque_type <- opaqueToCType desc   ctx_ty <- contextType   ops <- asks envOperations@@ -319,7 +335,7 @@                 $ty:opaque_type* v = malloc(sizeof($ty:opaque_type));                 $items:(criticalSection ops new_stms)                 *out = v;-                return 0;+                return FUTHARK_SUCCESS;               }|]   pure new   where@@ -359,33 +375,173 @@               )             ) -    setFieldField i e (ValueType _ (Rank r) _)-      | r == 0 =-          [C.cstm|v->$id:(tupleField i) = $exp:e;|]-      | otherwise =-          [C.cstm|{v->$id:(tupleField i) = malloc(sizeof(*$exp:e));-                   *v->$id:(tupleField i) = *$exp:e;-                   (void)(*(v->$id:(tupleField i)->mem.references))++;}|]+sumVariants ::+  Name ->+  [(Name, [(EntryPointType, [Int])])] ->+  [ValueType] ->+  CompilerM op s [Manifest.SumVariant]+sumVariants desc variants vds = do+  opaque_ty <- opaqueToCType desc+  ctx_ty <- contextType+  ops <- asks envOperations +  let onVariant i (name, payload) = do+        construct <- publicName $ "new_" <> opaqueName desc <> "_" <> nameToText name+        destruct <- publicName $ "destruct_" <> opaqueName desc <> "_" <> nameToText name++        constructFunction ops ctx_ty opaque_ty i construct payload+        destructFunction ops ctx_ty opaque_ty i destruct payload++        pure $+          Manifest.SumVariant+            { Manifest.sumVariantName = nameToText name,+              Manifest.sumVariantPayload = map (entryTypeName . fst) payload,+              Manifest.sumVariantConstruct = construct,+              Manifest.sumVariantDestruct = destruct+            }++  zipWithM onVariant [0 :: Int ..] variants+  where+    constructFunction ops ctx_ty opaque_ty i fname payload = do+      (params, new_stms) <- unzip <$> zipWithM constructPayload [0 ..] payload++      let used = concatMap snd payload+      set_unused_stms <-+        mapM setUnused $ filter ((`notElem` used) . fst) (zip [0 ..] vds)++      headerDecl+        (OpaqueDecl desc)+        [C.cedecl|int $id:fname($ty:ctx_ty *ctx,+                                $ty:opaque_ty **out,+                                $params:params);|]++      libDecl+        [C.cedecl|int $id:fname($ty:ctx_ty *ctx,+                                $ty:opaque_ty **out,+                                $params:params) {+                    (void)ctx;+                    $ty:opaque_ty* v = malloc(sizeof($ty:opaque_ty));+                    v->$id:(tupleField 0) = $int:i;+                    { $items:(criticalSection ops new_stms) }+                    // Set other fields+                    { $items:set_unused_stms }+                    *out = v;+                    return FUTHARK_SUCCESS;+                  }|]++    -- We must initialise some of the fields that are unused in this+    -- variant; specifically the ones corresponding to arrays. This+    -- has the unfortunate effect that all arrays in the nonused+    -- constructor are set to have size 0.+    setUnused (_, ValueType _ (Rank 0) _) =+      pure [C.citem|{}|]+    setUnused (i, ValueType signed (Rank rank) pt) = do+      new_array <- publicName $ "new_" <> arrayName pt signed rank+      let dims = replicate rank [C.cexp|0|]+      pure [C.citem|v->$id:(tupleField i) = $id:new_array(ctx, NULL, $args:dims);|]++    constructPayload j (et, is) = do+      let param_name = "v" <> show (j :: Int)+      case et of+        TypeTransparent (ValueType sign (Rank 0) pt) -> do+          let ct = primAPIType sign pt+              i = head is+          pure+            ( [C.cparam|const $ty:ct $id:param_name|],+              [C.citem|v->$id:(tupleField i) = $id:param_name;|]+            )+        TypeTransparent vt -> do+          ct <- valueTypeToCType Public vt+          let i = head is+          pure+            ( [C.cparam|const $ty:ct* $id:param_name|],+              [C.citem|{v->$id:(tupleField i) = malloc(sizeof($ty:ct));+                        memcpy(v->$id:(tupleField i), $id:param_name, sizeof(const $ty:ct));+                        (void)(*(v->$id:(tupleField i)->mem.references))++;}|]+            )+        TypeOpaque f_desc -> do+          ct <- opaqueToCType f_desc+          let param_fields = do+                i <- [0 ..]+                pure [C.cexp|$id:param_name->$id:(tupleField i)|]+              vts = map (vds !!) is+          pure+            ( [C.cparam|const $ty:ct* $id:param_name|],+              [C.citem|{$stms:(zipWith3 setFieldField is param_fields vts)}|]+            )++    destructFunction ops ctx_ty opaque_ty i fname payload = do+      (params, destruct_stms) <- unzip <$> zipWithM (destructPayload ops) [0 ..] payload+      headerDecl+        (OpaqueDecl desc)+        [C.cedecl|int $id:fname($ty:ctx_ty *ctx,+                                $params:params,+                                const $ty:opaque_ty *obj);|]++      libDecl+        [C.cedecl|int $id:fname($ty:ctx_ty *ctx,+                                $params:params,+                                const $ty:opaque_ty *obj) {+                    (void)ctx;+                    assert(obj->$id:(tupleField 0) == $int:i);+                    $stms:destruct_stms+                    return FUTHARK_SUCCESS;+                  }|]++    destructPayload ops j (et, is) = do+      let param_name = "v" <> show (j :: Int)+      (ct, project_items) <- projectField ops et $ zip is $ map (vds !!) is+      pure+        ( [C.cparam|$ty:ct* $id:param_name|],+          [C.cstm|{$ty:ct v;+                   $items:project_items+                   *$id:param_name = v;+                  }|]+        )++sumVariantFunction :: Name -> CompilerM op s Manifest.CFuncName+sumVariantFunction desc = do+  opaque_ty <- opaqueToCType desc+  ctx_ty <- contextType+  variant <- publicName $ "variant_" <> opaqueName desc+  headerDecl+    (OpaqueDecl desc)+    [C.cedecl|int $id:variant($ty:ctx_ty *ctx, const $ty:opaque_ty* v);|]+  -- This depends on the assumption that the first value always+  -- encodes the variant.+  libDecl+    [C.cedecl|int $id:variant($ty:ctx_ty *ctx, const $ty:opaque_ty* v) {+                (void)ctx;+                return v->$id:(tupleField 0);+              }|]+  pure variant+ processOpaqueRecord ::   OpaqueTypes ->   Name ->   OpaqueType ->   [ValueType] ->-  CompilerM op s (Maybe Manifest.RecordOps)-processOpaqueRecord _ _ (OpaqueType _) _ = pure Nothing+  CompilerM op s (Maybe Manifest.RecordOps, Maybe Manifest.SumOps)+processOpaqueRecord _ _ (OpaqueType _) _ =+  pure (Nothing, Nothing)+processOpaqueRecord _types desc (OpaqueSum _ cs) vds =+  (Nothing,) . Just+    <$> ( Manifest.SumOps+            <$> sumVariants desc cs vds+            <*> sumVariantFunction desc+        ) processOpaqueRecord types desc (OpaqueRecord fs) vds =-  Just+  (,Nothing) . Just     <$> ( Manifest.RecordOps-            <$> opaqueProjectFunctions types desc fs vds-            <*> opaqueNewFunctions types desc fs vds+            <$> recordProjectFunctions types desc fs vds+            <*> recordNewFunctions types desc fs vds         )  opaqueLibraryFunctions ::   OpaqueTypes ->   Name ->   OpaqueType ->-  CompilerM op s (Manifest.OpaqueOps, Maybe Manifest.RecordOps)+  CompilerM op s (Manifest.OpaqueOps, Maybe Manifest.RecordOps, Maybe Manifest.SumOps) opaqueLibraryFunctions types desc ot = do   name <- publicName $ opaqueName desc   free_opaque <- publicName $ "free_" <> opaqueName desc@@ -484,7 +640,7 @@     (OpaqueDecl desc)     [C.cedecl|$ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx, const void *p);|] -  record <- processOpaqueRecord types desc ot vds+  (record, sumops) <- processOpaqueRecord types desc ot vds    -- We do not need to enclose most bodies in a critical section,   -- because when we operate on the components of the opaque, we are@@ -511,6 +667,7 @@            $ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx,                                               const void *p) {+            (void)ctx;             int err = 0;             const unsigned char *src = p;             $ty:opaque_type* obj = malloc(sizeof($ty:opaque_type));@@ -531,7 +688,8 @@           Manifest.opaqueStore = store_opaque,           Manifest.opaqueRestore = restore_opaque         },-      record+      record,+      sumops     )  generateArray ::@@ -564,9 +722,9 @@   name <- publicName $ opaqueName desc   members <- zipWithM field (opaquePayload types ot) [(0 :: Int) ..]   libDecl [C.cedecl|struct $id:name { $sdecls:members };|]-  (ops, record) <- opaqueLibraryFunctions types desc ot+  (ops, record, sumops) <- opaqueLibraryFunctions types desc ot   let opaque_type = [C.cty|struct $id:name*|]-  pure (nameToText desc, Manifest.TypeOpaque (typeText opaque_type) ops record)+  pure (nameToText desc, Manifest.TypeOpaque (typeText opaque_type) ops record sumops)   where     field vt@(ValueType _ (Rank r) _) i = do       ct <- valueTypeToCType Private vt@@ -583,10 +741,12 @@   pure $ M.fromList $ catMaybes array_ts <> opaque_ts   where     -- Ensure that array types will be generated before the opaque-    -- records that allow projection of them.  This is because the+    -- types that allow projection of them.  This is because the     -- projection functions somewhat uglily directly poke around in     -- the innards to increment reference counts.     findNecessaryArrays (OpaqueType _) =       pure ()+    findNecessaryArrays (OpaqueSum _ variants) =+      mapM_ (mapM_ (entryPointTypeToCType Public . fst) . snd) variants     findNecessaryArrays (OpaqueRecord fs) =       mapM_ (entryPointTypeToCType Public . snd) fs

@@ -19,7 +19,7 @@     fromStorage,     toStorage,     Operations (..),-    DoLMADCopy,+    DoCopy,     defaultOperations,     unpackDim,     CompilerM (..),@@ -102,10 +102,10 @@   PrimType ->   CompilerM op s () --- | Perform an 'Imp.LMADCopy'.  It is expected that these functions+-- | Perform an 'Imp.Copy'.  It is expected that these functions -- are each specialised on which spaces they operate on, so that is -- not part of their arguments.-type DoLMADCopy op s =+type DoCopy op s =   PrimType ->   [Count Elements PyExp] ->   PyExp ->@@ -142,7 +142,7 @@     opsReadScalar :: ReadScalar op s,     opsAllocate :: Allocate op s,     -- | @(dst,src)@-space mapping to copy functions.-    opsCopies :: M.Map (Space, Space) (DoLMADCopy op s),+    opsCopies :: M.Map (Space, Space) (DoCopy op s),     opsCompiler :: OpCompiler op s,     opsEntryOutput :: EntryOutput op s,     opsEntryInput :: EntryInput op s@@ -1190,9 +1190,9 @@ generateWrite dst iexp _ DefaultSpace elemexp =   stm $ Exp $ simpleCall "writeScalarArray" [dst, iexp, elemexp] --- | Compile an 'LMADCopy' using sequential nested loops, but+-- | Compile an 'Copy' using sequential nested loops, but -- parameterised over how to do the reads and writes.-compileLMADCopyWith ::+compileCopyWith ::   [Count Elements (TExp Int64)] ->   (PyExp -> PyExp -> CompilerM op s ()) ->   ( Count Elements (TExp Int64),@@ -1203,7 +1203,7 @@     [Count Elements (TExp Int64)]   ) ->   CompilerM op s ()-compileLMADCopyWith shape doWrite dst_lmad doRead src_lmad = do+compileCopyWith shape doWrite dst_lmad doRead src_lmad = do   let (dstoffset, dststrides) = dst_lmad       (srcoffset, srcstrides) = src_lmad   shape' <- mapM (compileExp . untyped . unCount) shape@@ -1225,10 +1225,10 @@     loops ((i, n) : ins) body =       [For (compileName i) (simpleCall "range" [n]) $ loops ins body] --- | Compile an 'LMADCopy' using sequential nested loops and+-- | Compile an 'Copy' using sequential nested loops and -- 'Imp.Read'/'Imp.Write' of individual scalars.  This always works, -- but can be pretty slow if those reads and writes are costly.-compileLMADCopy ::+compileCopy ::   PrimType ->   [Count Elements (TExp Int64)] ->   (VName, Space) ->@@ -1240,12 +1240,12 @@     [Count Elements (TExp Int64)]   ) ->   CompilerM op s ()-compileLMADCopy t shape (dst, dstspace) dst_lmad (src, srcspace) src_lmad = do+compileCopy t shape (dst, dstspace) dst_lmad (src, srcspace) src_lmad = do   src' <- compileVar src   dst' <- compileVar dst   let doWrite dst_i = generateWrite dst' dst_i t dstspace       doRead src_i = generateRead src' src_i t srcspace-  compileLMADCopyWith shape doWrite dst_lmad doRead src_lmad+  compileCopyWith shape doWrite dst_lmad doRead src_lmad  compileCode :: Imp.Code op -> CompilerM op s () compileCode Imp.DebugPrint {} =@@ -1351,11 +1351,11 @@   stm =<< Assign <$> compileVar name <*> pure allocate' compileCode (Imp.Free name _) =   stm =<< Assign <$> compileVar name <*> pure None-compileCode (Imp.LMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) = do+compileCode (Imp.Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) = do   cp <- asks $ M.lookup (dstspace, srcspace) . opsCopies . envOperations   case cp of     Nothing ->-      compileLMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)+      compileCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)     Just cp' -> do       shape' <- traverse (traverse (compileExp . untyped)) shape       dst' <- compileVar dst@@ -1377,7 +1377,7 @@   stm . Assign x' =<< generateRead src' iexp' pt space compileCode Imp.Skip = pure () -lmadcopyCPU :: DoLMADCopy op s+lmadcopyCPU :: DoCopy op s lmadcopyCPU t shape dst (dstoffset, dststride) src (srcoffset, srcstride) =   stm . Exp . simpleCall "lmad_copy" $     [ Var (compilePrimType t),

@@ -563,7 +563,7 @@         <$> compileExp (untyped iexp)         <*> getMemType src restype   GC.stm [C.cstm|$id:x = $exp:e;|]-compileCode (LMADCopy t shape (dst, DefaultSpace) dst_lmad (src, DefaultSpace) src_lmad) = do+compileCode (Copy t shape (dst, DefaultSpace) dst_lmad (src, DefaultSpace) src_lmad) = do   dst' <- GC.rawMem dst   src' <- GC.rawMem src   let doWrite dst_i ve = do@@ -575,7 +575,7 @@         GC.stm [C.cstm|$exp:deref = $exp:(toStorage t ve);|]       doRead src_i =         fromStorage t . GC.derefPointer src' src_i <$> getMemType src t-  GC.compileLMADCopyWith shape doWrite dst_lmad doRead src_lmad+  GC.compileCopyWith shape doWrite dst_lmad doRead src_lmad compileCode (Free name space) = do   cached <- isJust <$> GC.cacheMem name   unless cached $ unRefMem name space

@@ -390,7 +390,7 @@ finishIfSynchronous =   stm $ If (Var "synchronous") [Exp $ simpleCall "sync" [Var "self"]] [] -copygpu2gpu :: DoLMADCopy op s+copygpu2gpu :: DoCopy op s copygpu2gpu t shape dst (dstoffset, dststride) src (srcoffset, srcstride) = do   stm . Exp . simpleCall "lmad_copy_gpu2gpu" $     [ Var "self",

@@ -28,7 +28,7 @@ -- ImpCode does not have arrays. 'DeclareArray' is for declaring -- constant array literals, not arrays in general.  Instead, ImpCode -- deals only with memory.  Array operations present in core IR--- programs are turned into 'Write', v'Read', and 'LMADCopy'+-- programs are turned into 'Write', v'Read', and 'Copy' -- operations that use flat indexes and offsets based on the index -- function of the original array. --@@ -280,8 +280,8 @@     -- all memory blocks will be freed with this statement.     -- Backends are free to ignore it entirely.     Free VName Space-  | -- | @LMADcopy pt dest dest_lmad src src_lmad shape@-    LMADCopy+  | -- | @Copy pt shape dest dest_lmad src src_lmad@.+    Copy       PrimType       [Count Elements (TExp Int64)]       (VName, Space)@@ -609,7 +609,7 @@     pretty dest <+> "<-" <+> pretty from <+> "@" <> pretty space   pretty (Assert e msg _) =     "assert" <> parens (commasep [pretty msg, pretty e])-  pretty (LMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) =+  pretty (Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) =     ("lmadcopy_" <> pretty (length shape) <> "d_" <> pretty t)       <> (parens . align)         ( foldMap (brackets . pretty) shape@@ -707,8 +707,8 @@     pure $ Allocate name size s   traverse _ (Free name space) =     pure $ Free name space-  traverse _ (LMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) =-    pure $ LMADCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)+  traverse _ (Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) =+    pure $ Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)   traverse _ (Write name i bt val space vol) =     pure $ Write name i bt val space vol   traverse _ (Read x name i bt space vol) =@@ -781,7 +781,7 @@     freeIn' name <> freeIn' size <> freeIn' space   freeIn' (Free name _) =     freeIn' name-  freeIn' (LMADCopy _ shape (dst, _) (dstoffset, dststrides) (src, _) (srcoffset, srcstrides)) =+  freeIn' (Copy _ shape (dst, _) (dstoffset, dststrides) (src, _) (srcoffset, srcstrides)) =     freeIn' shape <> freeIn' dst <> freeIn' dstoffset <> freeIn' dststrides <> freeIn' src <> freeIn' srcoffset <> freeIn' srcstrides   freeIn' (SetMem x y _) =     freeIn' x <> freeIn' y

@@ -140,7 +140,6 @@ import Futhark.CodeGen.ImpCode qualified as Imp import Futhark.Construct hiding (ToExp (..)) import Futhark.IR.Mem-import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SOACS (SOACS) import Futhark.Util@@ -486,6 +485,7 @@   case lookupOpaqueType desc types of     OpaqueType vts -> map valueTypeSign vts     OpaqueRecord fs -> foldMap (entryPointSignedness types . snd) fs+    OpaqueSum vts _ -> map valueTypeSign vts  -- | How many value parameters are accepted by this entry point?  This -- is used to determine which of the function parameters correspond to@@ -497,6 +497,7 @@   case lookupOpaqueType desc types of     OpaqueType vts -> length vts     OpaqueRecord fs -> sum $ map (entryPointSize types . snd) fs+    OpaqueSum vts _ -> length vts  compileInParam ::   (Mem rep inner) =>@@ -508,7 +509,7 @@   MemMem space ->     pure $ Left $ Imp.MemParam name space   MemArray bt shape _ (ArrayIn mem lmad) ->-    pure $ Right $ ArrayDecl name bt $ MemLoc mem (shapeDims shape) $ IxFun.ixfunLMAD lmad+    pure $ Right $ ArrayDecl name bt $ MemLoc mem (shapeDims shape) lmad   MemAcc {} ->     error "Functions may not have accumulator parameters."   where@@ -1103,7 +1104,7 @@ memBoundToVarEntry e (MemAcc acc ispace ts _) =   AccVar e (acc, ispace, ts) memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem lmad)) =-  let location = MemLoc mem (shapeDims shape) $ IxFun.ixfunLMAD lmad+  let location = MemLoc mem (shapeDims shape) lmad    in ArrayVar         e         ArrayEntry@@ -1405,7 +1406,7 @@   srcspace <- entryMemSpace <$> lookupMemory srcmem   dstspace <- entryMemSpace <$> lookupMemory dstmem   emit $-    Imp.LMADCopy+    Imp.Copy       t       (elements <$> LMAD.shape dstlmad)       (dstmem, dstspace)@@ -1582,8 +1583,8 @@         case dest_entry of           ScalarVar _ _ ->             ScalarDestination dest-          ArrayVar _ (ArrayEntry (MemLoc mem shape ixfun) _) ->-            ArrayDestination $ Just $ MemLoc mem shape ixfun+          ArrayVar _ (ArrayEntry (MemLoc mem shape lmad) _) ->+            ArrayDestination $ Just $ MemLoc mem shape lmad           MemVar _ _ ->             MemoryDestination dest           AccVar {} ->@@ -1708,9 +1709,9 @@   pure name'  sArray :: String -> PrimType -> ShapeBase SubExp -> VName -> LMAD -> ImpM rep r op VName-sArray name bt shape mem ixfun = do+sArray name bt shape mem lmad = do   name' <- newVName name-  dArray name' bt shape mem ixfun+  dArray name' bt shape mem lmad   pure name'  -- | Declare an array in row-major order in the given memory block.@@ -1726,9 +1727,9 @@ sAllocArrayPerm name pt shape space perm = do   let permuted_dims = rearrangeShape perm $ shapeDims shape   mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space-  let iota_ixfun = LMAD.iota 0 $ map (isInt64 . primExpFromSubExp int64) permuted_dims+  let iota_lmad = LMAD.iota 0 $ map (isInt64 . primExpFromSubExp int64) permuted_dims   sArray name pt shape mem $-    LMAD.permute iota_ixfun $+    LMAD.permute iota_lmad $       rearrangeInverse perm  -- | Uses linear/iota index function.

@@ -1314,11 +1314,11 @@  replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) replicateIsFill arr v = do-  ArrayEntry (MemLoc arr_mem arr_shape arr_ixfun) _ <- lookupArray arr+  ArrayEntry (MemLoc arr_mem arr_shape arr_lmad) _ <- lookupArray arr   v_t <- subExpType v   case v_t of     Prim v_t'-      | LMAD.isDirect arr_ixfun -> pure $+      | LMAD.isDirect arr_lmad -> pure $           Just $ do             fname <- replicateForType v_t'             emit $@@ -1417,8 +1417,8 @@   IntType ->   CallKernelGen () sIota arr n x s et = do-  ArrayEntry (MemLoc arr_mem _ arr_ixfun) _ <- lookupArray arr-  if LMAD.isDirect arr_ixfun+  ArrayEntry (MemLoc arr_mem _ arr_lmad) _ <- lookupArray arr+  if LMAD.isDirect arr_lmad     then do       fname <- iotaForType et       emit $

@@ -45,7 +45,7 @@  sliceArray :: Imp.TExp Int64 -> TV Int64 -> VName -> ImpM rep r op VName sliceArray start size arr = do-  MemLoc mem _ ixfun <- entryArrayLoc <$> lookupArray arr+  MemLoc mem _ lmad <- entryArrayLoc <$> lookupArray arr   arr_t <- lookupType arr   let slice =         fullSliceNum@@ -56,7 +56,7 @@     (elemType arr_t)     (arrayShape arr_t `setOuterDim` Var (tvVar size))     mem-    $ LMAD.slice ixfun slice+    $ LMAD.slice lmad slice  -- | @applyLambda lam dests args@ emits code that: --

@@ -44,9 +44,9 @@ -- -- An optimization specfically targeted at non-segmented and large-segments -- segmented reductions with non-commutative is made: The stage one main loop is--- essentially stripmined by a factor `chunk`, inserting collective copies via+-- essentially stripmined by a factor *chunk*, inserting collective copies via -- local memory of each reduction parameter going into the intra-block (partial)--- reductions. This saves a factor `chunk` number of intra-block reductions at+-- reductions. This saves a factor *chunk* number of intra-block reductions at -- the cost of some overhead in collective copies. module Futhark.CodeGen.ImpGen.GPU.SegRed   ( compileSegRed,

@@ -810,7 +810,7 @@ typesInCode (DeclareArray _ t _) = S.singleton t typesInCode (Allocate _ (Count (TPrimExp e)) _) = typesInExp e typesInCode Free {} = mempty-typesInCode (LMADCopy _ shape _ (Count (TPrimExp dstoffset), dststrides) _ (Count (TPrimExp srcoffset), srcstrides)) =+typesInCode (Copy _ shape _ (Count (TPrimExp dstoffset), dststrides) _ (Count (TPrimExp srcoffset), srcstrides)) =   foldMap (typesInExp . untyped . unCount) shape     <> typesInExp dstoffset     <> foldMap (typesInExp . untyped . unCount) dststrides

@@ -148,7 +148,7 @@     onError (ProgWarning loc msg) =       annotate (color Yellow) $ "Warning at " <> pretty (locText (srclocOf loc)) <> ":" </> unAnnotate msg --- | Throw an exception formatted with 'pprProgErrors' if there's+-- | Throw an exception formatted with 'prettyProgErrors' if there's -- an error. throwOnProgError ::   (MonadError CompilerError m) =>

@@ -62,10 +62,9 @@     MemBound,     MemBind (..),     MemReturn (..),-    IxFun,-    ExtIxFun,     LMAD,-    isStaticIxFun,+    ExtLMAD,+    isStaticLMAD,     ExpReturns,     BodyReturns,     FunReturns,@@ -81,7 +80,7 @@     subExpMemInfo,     lookupArraySummary,     lookupMemSpace,-    existentialiseIxFun,+    existentialiseLMAD,      -- * Type checking parts     matchBranchReturnType,@@ -123,7 +122,7 @@     removePatAliases,     removeScopeAliases,   )-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.Pretty import Futhark.IR.Prop import Futhark.IR.Prop.Aliases@@ -250,14 +249,11 @@   indexOp vtable k (Inner op) is = ST.indexOp vtable k op is   indexOp _ _ _ _ = Nothing --- | The index function representation used for memory annotations.-type IxFun = IxFun.IxFun (TPrimExp Int64 VName)- -- | The LMAD representation used for memory annotations.-type LMAD = IxFun.LMAD (TPrimExp Int64 VName)+type LMAD = LMAD.LMAD (TPrimExp Int64 VName)  -- | An index function that may contain existential variables.-type ExtIxFun = IxFun.IxFun (TPrimExp Int64 (Ext VName))+type ExtLMAD = LMAD.LMAD (TPrimExp Int64 (Ext VName))  -- | A summary of the memory information for every let-bound -- identifier, function parameter, and return value.  Parameterisered@@ -297,10 +293,16 @@ instance (FixExt ret) => FixExt (MemInfo ExtSize u ret) where   fixExt _ _ (MemPrim pt) = MemPrim pt   fixExt _ _ (MemMem space) = MemMem space+  fixExt _ _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u   fixExt i se (MemArray pt shape u ret) =     MemArray pt (fixExt i se shape) u (fixExt i se ret)-  fixExt _ _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u +  mapExt _ (MemPrim pt) = MemPrim pt+  mapExt _ (MemMem space) = MemMem space+  mapExt _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u+  mapExt f (MemArray pt shape u ret) =+    MemArray pt (mapExt f shape) u (mapExt f ret)+ instance Typed (MemInfo SubExp Uniqueness ret) where   typeOf = fromDecl . declTypeOf @@ -343,20 +345,20 @@ instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where   rename = substituteRename -simplifyIxFun ::+simplifyLMAD ::   (Engine.SimplifiableRep rep) =>-  IxFun ->-  Engine.SimpleM rep IxFun-simplifyIxFun = traverse $ fmap isInt64 . simplifyPrimExp . untyped+  LMAD ->+  Engine.SimpleM rep LMAD+simplifyLMAD = traverse $ fmap isInt64 . simplifyPrimExp . untyped -simplifyExtIxFun ::+simplifyExtLMAD ::   (Engine.SimplifiableRep rep) =>-  ExtIxFun ->-  Engine.SimpleM rep ExtIxFun-simplifyExtIxFun = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped+  ExtLMAD ->+  Engine.SimpleM rep ExtLMAD+simplifyExtLMAD = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped -isStaticIxFun :: ExtIxFun -> Maybe IxFun-isStaticIxFun = traverse $ traverse inst+isStaticLMAD :: ExtLMAD -> Maybe LMAD+isStaticLMAD = traverse $ traverse inst   where     inst Ext {} = Nothing     inst (Free x) = Just x@@ -395,7 +397,7 @@ data MemBind   = -- | Located in this memory block with this index     -- function.-    ArrayIn VName IxFun+    ArrayIn VName LMAD   deriving (Show)  instance Eq MemBind where@@ -408,25 +410,25 @@   rename = substituteRename  instance Substitute MemBind where-  substituteNames substs (ArrayIn ident ixfun) =-    ArrayIn (substituteNames substs ident) (substituteNames substs ixfun)+  substituteNames substs (ArrayIn ident lmad) =+    ArrayIn (substituteNames substs ident) (substituteNames substs lmad)  instance PP.Pretty MemBind where-  pretty (ArrayIn mem ixfun) =-    PP.pretty mem <+> "->" PP.</> PP.pretty ixfun+  pretty (ArrayIn mem lmad) =+    PP.pretty mem <+> "->" PP.</> PP.pretty lmad  instance FreeIn MemBind where-  freeIn' (ArrayIn mem ixfun) = freeIn' mem <> freeIn' ixfun+  freeIn' (ArrayIn mem lmad) = freeIn' mem <> freeIn' lmad  -- | A description of the memory properties of an array being returned -- by an operation. data MemReturn   = -- | The array is located in a memory block that is     -- already in scope.-    ReturnsInBlock VName ExtIxFun+    ReturnsInBlock VName ExtLMAD   | -- | The operation returns a new (existential) memory     -- block.-    ReturnsNewBlock Space Int ExtIxFun+    ReturnsNewBlock Space Int ExtLMAD   deriving (Show)  instance Eq MemReturn where@@ -439,33 +441,41 @@   rename = substituteRename  instance Substitute MemReturn where-  substituteNames substs (ReturnsInBlock ident ixfun) =-    ReturnsInBlock (substituteNames substs ident) (substituteNames substs ixfun)-  substituteNames substs (ReturnsNewBlock space i ixfun) =-    ReturnsNewBlock space i (substituteNames substs ixfun)+  substituteNames substs (ReturnsInBlock ident lmad) =+    ReturnsInBlock (substituteNames substs ident) (substituteNames substs lmad)+  substituteNames substs (ReturnsNewBlock space i lmad) =+    ReturnsNewBlock space i (substituteNames substs lmad)  instance FixExt MemReturn where-  fixExt i (Var v) (ReturnsNewBlock _ j ixfun)+  fixExt i (Var v) (ReturnsNewBlock _ j lmad)     | j == i =         ReturnsInBlock v $-          fixExtIxFun+          fixExtLMAD             i             (primExpFromSubExp int64 (Var v))-            ixfun-  fixExt i se (ReturnsNewBlock space j ixfun) =+            lmad+  fixExt i se (ReturnsNewBlock space j lmad) =     ReturnsNewBlock       space       j'-      (fixExtIxFun i (primExpFromSubExp int64 se) ixfun)+      (fixExtLMAD i (primExpFromSubExp int64 se) lmad)     where       j'         | i < j = j - 1         | otherwise = j-  fixExt i se (ReturnsInBlock mem ixfun) =-    ReturnsInBlock mem (fixExtIxFun i (primExpFromSubExp int64 se) ixfun)+  fixExt i se (ReturnsInBlock mem lmad) =+    ReturnsInBlock mem (fixExtLMAD i (primExpFromSubExp int64 se) lmad) -fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun-fixExtIxFun i e = fmap $ isInt64 . replaceInPrimExp update . untyped+  mapExt f (ReturnsNewBlock space i lmad) =+    ReturnsNewBlock space (f i) lmad+  mapExt f (ReturnsInBlock mem lmad) =+    ReturnsInBlock mem (fmap (fmap f') lmad)+    where+      f' (Ext i) = Ext $ f i+      f' v = v++fixExtLMAD :: Int -> PrimExp VName -> ExtLMAD -> ExtLMAD+fixExtLMAD i e = fmap $ isInt64 . replaceInPrimExp update . untyped   where     update (Ext j) t       | j > i = LeafExp (Ext $ j - 1) t@@ -476,30 +486,30 @@ leafExp :: Int -> TPrimExp Int64 (Ext a) leafExp i = isInt64 $ LeafExp (Ext i) int64 -existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun-existentialiseIxFun ctx = IxFun.substituteInIxFun ctx' . fmap (fmap Free)+existentialiseLMAD :: [VName] -> LMAD -> ExtLMAD+existentialiseLMAD ctx = LMAD.substitute ctx' . fmap (fmap Free)   where     ctx' = M.map leafExp $ M.fromList $ zip (map Free ctx) [0 ..]  instance PP.Pretty MemReturn where-  pretty (ReturnsInBlock v ixfun) =-    pretty v <+> "->" PP.</> PP.pretty ixfun-  pretty (ReturnsNewBlock space i ixfun) =-    "?" <> pretty i <> PP.pretty space <+> "->" PP.</> PP.pretty ixfun+  pretty (ReturnsInBlock v lmad) =+    pretty v <+> "->" PP.</> PP.pretty lmad+  pretty (ReturnsNewBlock space i lmad) =+    "?" <> pretty i <> PP.pretty space <+> "->" PP.</> PP.pretty lmad  instance FreeIn MemReturn where-  freeIn' (ReturnsInBlock v ixfun) = freeIn' v <> freeIn' ixfun-  freeIn' (ReturnsNewBlock space _ ixfun) = freeIn' space <> freeIn' ixfun+  freeIn' (ReturnsInBlock v lmad) = freeIn' v <> freeIn' lmad+  freeIn' (ReturnsNewBlock space _ lmad) = freeIn' space <> freeIn' lmad  instance Engine.Simplifiable MemReturn where-  simplify (ReturnsNewBlock space i ixfun) =-    ReturnsNewBlock space i <$> simplifyExtIxFun ixfun-  simplify (ReturnsInBlock v ixfun) =-    ReturnsInBlock <$> Engine.simplify v <*> simplifyExtIxFun ixfun+  simplify (ReturnsNewBlock space i lmad) =+    ReturnsNewBlock space i <$> simplifyExtLMAD lmad+  simplify (ReturnsInBlock v lmad) =+    ReturnsInBlock <$> Engine.simplify v <*> simplifyExtLMAD lmad  instance Engine.Simplifiable MemBind where-  simplify (ArrayIn mem ixfun) =-    ArrayIn <$> Engine.simplify mem <*> simplifyIxFun ixfun+  simplify (ArrayIn mem lmad) =+    ArrayIn <$> Engine.simplify mem <*> simplifyLMAD lmad  instance Engine.Simplifiable [FunReturns] where   simplify = mapM Engine.simplify@@ -553,11 +563,11 @@ bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns  varInfoToExpReturns :: MemInfo SubExp NoUniqueness MemBind -> ExpReturns-varInfoToExpReturns (MemArray et shape u (ArrayIn mem ixfun)) =+varInfoToExpReturns (MemArray et shape u (ArrayIn mem lmad)) =   MemArray et (fmap Free shape) u $     Just $       ReturnsInBlock mem $-        existentialiseIxFun [] ixfun+        existentialiseLMAD [] lmad varInfoToExpReturns (MemPrim pt) = MemPrim pt varInfoToExpReturns (MemAcc acc ispace ts u) = MemAcc acc ispace ts u varInfoToExpReturns (MemMem space) = MemMem space@@ -589,15 +599,15 @@         MemPrim _ -> pure ()         MemMem {} -> pure ()         MemAcc {} -> pure ()-        MemArray _ _ _ (ArrayIn _ ixfun)-          | IxFun.isDirect ixfun ->+        MemArray _ _ _ (ArrayIn _ lmad)+          | LMAD.isDirect lmad ->               pure ()           | otherwise ->               TC.bad . TC.TypeError $                 "Array "                   <> prettyText v                   <> " returned by function, but has nontrivial index function:\n"-                  <> prettyText ixfun+                  <> prettyText lmad  matchLoopResultMem ::   (Mem rep inner, TC.Checkable rep) =>@@ -625,15 +635,15 @@       MemMem space     toRet (MemAcc acc ispace ts u) =       MemAcc acc ispace ts u-    toRet (MemArray pt shape u (ArrayIn mem ixfun))+    toRet (MemArray pt shape u (ArrayIn mem lmad))       | Just i <- mem `elemIndex` param_names,         Param _ _ (MemMem space) : _ <- drop i params =-          MemArray pt shape' u $ ReturnsNewBlock space i ixfun'+          MemArray pt shape' u $ ReturnsNewBlock space i lmad'       | otherwise =-          MemArray pt shape' u $ ReturnsInBlock mem ixfun'+          MemArray pt shape' u $ ReturnsInBlock mem lmad'       where         shape' = fmap toExtSE shape-        ixfun' = existentialiseIxFun param_names ixfun+        lmad' = existentialiseLMAD param_names lmad  matchBranchReturnType ::   (Mem rep inner, TC.Checkable rep) =>@@ -680,8 +690,8 @@   [MemInfo SubExp NoUniqueness MemBind] ->   TC.TypeM rep () matchReturnType rettype res ts = do-  let existentialiseIxFun0 :: IxFun -> ExtIxFun-      existentialiseIxFun0 = fmap $ fmap Free+  let existentialiseLMAD0 :: LMAD -> ExtLMAD+      existentialiseLMAD0 = fmap $ fmap Free        fetchCtx i = case maybeNth i $ zip res ts of         Nothing ->@@ -714,27 +724,27 @@         unless (x == y) . throwError . T.unwords $           ["Expected ext dim", prettyText i, "=>", prettyText x, "but got", prettyText y] -      checkMemReturn (ReturnsInBlock x_mem x_ixfun) (ArrayIn y_mem y_ixfun)+      checkMemReturn (ReturnsInBlock x_mem x_lmad) (ArrayIn y_mem y_lmad)         | x_mem == y_mem =-            unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $+            unless (LMAD.closeEnough x_lmad $ existentialiseLMAD0 y_lmad) $               throwError . T.unwords $                 [ "Index function unification failed (ReturnsInBlock)",-                  "\nixfun of body result: ",-                  prettyText y_ixfun,-                  "\nixfun of return type: ",-                  prettyText x_ixfun+                  "\nlmad of body result: ",+                  prettyText y_lmad,+                  "\nlmad of return type: ",+                  prettyText x_lmad                 ]       checkMemReturn-        (ReturnsNewBlock x_space x_ext x_ixfun)-        (ArrayIn y_mem y_ixfun) = do+        (ReturnsNewBlock x_space x_ext x_lmad)+        (ArrayIn y_mem y_lmad) = do           (x_mem, x_mem_type) <- fetchCtx x_ext-          unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $+          unless (LMAD.closeEnough x_lmad $ existentialiseLMAD0 y_lmad) $             throwError . docText $               "Index function unification failed (ReturnsNewBlock)"-                </> "Ixfun of body result:"-                </> indent 2 (pretty y_ixfun)-                </> "Ixfun of return type:"-                </> indent 2 (pretty x_ixfun)+                </> "Lmad of body result:"+                </> indent 2 (pretty y_lmad)+                </> "Lmad of return type:"+                </> indent 2 (pretty x_lmad)           case x_mem_type of             MemMem y_space ->               unless (x_space == y_space) . throwError . T.unwords $@@ -763,6 +773,13 @@             </> indent 2 (ppTupleLines' $ map pretty ts)             </> pretty s +  unless (length rettype == length ts) $+    TC.bad . TC.TypeError . docText $+      "Return type"+        </> indent 2 (ppTupleLines' $ map pretty rettype)+        </> "does not have same number of elements as results"+        </> indent 2 (ppTupleLines' $ map pretty ts)+   either bad pure =<< runExceptT (zipWithM_ checkReturn rettype ts)  matchPatToExp ::@@ -802,22 +819,22 @@       x_pt == y_pt         && x_shape == y_shape         && case (x_ret, y_ret) of-          (ReturnsInBlock _ x_ixfun, Just (ReturnsInBlock _ y_ixfun)) ->-            let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun-                y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun-             in IxFun.closeEnough x_ixfun' y_ixfun'-          ( ReturnsInBlock _ x_ixfun,-            Just (ReturnsNewBlock _ _ y_ixfun)+          (ReturnsInBlock _ x_lmad, Just (ReturnsInBlock _ y_lmad)) ->+            let x_lmad' = LMAD.substitute ctxids x_lmad+                y_lmad' = LMAD.substitute ctxexts y_lmad+             in LMAD.closeEnough x_lmad' y_lmad'+          ( ReturnsInBlock _ x_lmad,+            Just (ReturnsNewBlock _ _ y_lmad)             ) ->-              let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun-                  y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun-               in IxFun.closeEnough x_ixfun' y_ixfun'-          ( ReturnsNewBlock _ x_i x_ixfun,-            Just (ReturnsNewBlock _ y_i y_ixfun)+              let x_lmad' = LMAD.substitute ctxids x_lmad+                  y_lmad' = LMAD.substitute ctxexts y_lmad+               in LMAD.closeEnough x_lmad' y_lmad'+          ( ReturnsNewBlock _ x_i x_lmad,+            Just (ReturnsNewBlock _ y_i y_lmad)             ) ->-              let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun-                  y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun-               in x_i == y_i && IxFun.closeEnough x_ixfun' y_ixfun'+              let x_lmad' = LMAD.substitute ctxids x_lmad+                  y_lmad' = LMAD.substitute ctxexts y_lmad+               in x_i == y_i && LMAD.closeEnough x_lmad' y_lmad'           (_, Nothing) -> True           _ -> False     matches _ _ _ _ = False@@ -859,12 +876,12 @@ lookupArraySummary ::   (Mem rep inner, HasScope rep m, Monad m) =>   VName ->-  m (VName, IxFun.IxFun (TPrimExp Int64 VName))+  m (VName, LMAD.LMAD (TPrimExp Int64 VName)) lookupArraySummary name = do   summary <- lookupMemInfo name   case summary of-    MemArray _ _ _ (ArrayIn mem ixfun) ->-      pure (mem, ixfun)+    MemArray _ _ _ (ArrayIn mem lmad) ->+      pure (mem, lmad)     _ ->       error . T.unpack $         "Expected "@@ -898,7 +915,7 @@ checkMemInfo _ (MemMem _) = pure () checkMemInfo _ (MemAcc acc ispace ts u) =   TC.checkType $ Acc acc ispace ts u-checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do+checkMemInfo name (MemArray _ shape _ (ArrayIn v lmad)) = do   t <- lookupType v   case t of     Mem {} ->@@ -912,13 +929,13 @@             <> prettyText t             <> "." -  TC.context ("in index function " <> prettyText ixfun) $ do-    traverse_ (TC.requirePrimExp int64 . untyped) ixfun-    unless (IxFun.shape ixfun == map pe64 (shapeDims shape)) $+  TC.context ("in index function " <> prettyText lmad) $ do+    traverse_ (TC.requirePrimExp int64 . untyped) lmad+    unless (LMAD.shape lmad == map pe64 (shapeDims shape)) $       TC.bad $         TC.TypeError $           "Shape of index function ("-            <> prettyText (IxFun.shape ixfun)+            <> prettyText (LMAD.shape lmad)             <> ") does not match shape of array "             <> prettyText name             <> " ("@@ -942,13 +959,13 @@         case patElemDec pe of           MemPrim pt -> MemPrim pt           MemMem space -> MemMem space-          MemArray pt shape u (ArrayIn mem ixfun) ->+          MemArray pt shape u (ArrayIn mem lmad) ->             MemArray pt (Shape $ map ext $ shapeDims shape) u $               case find ((== mem) . patElemName . snd) $ zip [0 ..] ctx of                 Just (i, PatElem _ (MemMem space)) ->                   ReturnsNewBlock space i $-                    existentialiseIxFun (map patElemName ctx) ixfun-                _ -> ReturnsInBlock mem $ existentialiseIxFun [] ixfun+                    existentialiseLMAD (map patElemName ctx) lmad+                _ -> ReturnsInBlock mem $ existentialiseLMAD [] lmad           MemAcc acc ispace ts u -> MemAcc acc ispace ts u       ) @@ -963,13 +980,9 @@     addDec t@(Array bt shape u)       | existential t = do           i <- get <* modify (+ 1)-          pure $-            MemArray bt shape u $-              Just $-                ReturnsNewBlock DefaultSpace i $-                  IxFun.iota $-                    map convert $-                      shapeDims shape+          pure . MemArray bt shape u . Just $+            ReturnsNewBlock DefaultSpace i $+              LMAD.iota 0 (map convert $ shapeDims shape)       | otherwise =           pure $ MemArray bt shape u Nothing     addDec (Acc acc ispace ts u) =@@ -980,12 +993,12 @@ arrayVarReturns ::   (HasScope rep m, Monad m, Mem rep inner) =>   VName ->-  m (PrimType, Shape, VName, IxFun)+  m (PrimType, Shape, VName, LMAD) arrayVarReturns v = do   summary <- lookupMemInfo v   case summary of-    MemArray et shape _ (ArrayIn mem ixfun) ->-      pure (et, Shape $ shapeDims shape, mem, ixfun)+    MemArray et shape _ (ArrayIn mem lmad) ->+      pure (et, Shape $ shapeDims shape, mem, lmad)     _ ->       error . T.unpack $ "arrayVarReturns: " <> prettyText v <> " is not an array." @@ -998,12 +1011,12 @@   case summary of     MemPrim bt ->       pure $ MemPrim bt-    MemArray et shape _ (ArrayIn mem ixfun) ->+    MemArray et shape _ (ArrayIn mem lmad) ->       pure $         MemArray et (fmap Free shape) NoUniqueness $           Just $             ReturnsInBlock mem $-              existentialiseIxFun [] ixfun+              existentialiseLMAD [] lmad     MemMem space ->       pure $ MemMem space     MemAcc acc ispace ts u ->@@ -1029,29 +1042,29 @@ expReturns (BasicOp (Opaque _ (Var v))) =   Just . pure <$> varReturns v expReturns (BasicOp (Reshape k newshape v)) = do-  (et, _, mem, ixfun) <- arrayVarReturns v-  case reshaper k ixfun $ map pe64 $ shapeDims newshape of-    Just ixfun' ->+  (et, _, mem, lmad) <- arrayVarReturns v+  case reshaper k lmad $ map pe64 $ shapeDims newshape of+    Just lmad' ->       pure . Just $         [ MemArray et (fmap Free newshape) NoUniqueness . Just $-            ReturnsInBlock mem (existentialiseIxFun [] ixfun')+            ReturnsInBlock mem (existentialiseLMAD [] lmad')         ]     Nothing -> pure Nothing   where-    reshaper ReshapeArbitrary ixfun =-      IxFun.reshape ixfun-    reshaper ReshapeCoerce ixfun =-      Just . IxFun.coerce ixfun+    reshaper ReshapeArbitrary lmad =+      LMAD.reshape lmad+    reshaper ReshapeCoerce lmad =+      Just . LMAD.coerce lmad expReturns (BasicOp (Rearrange perm v)) = do-  (et, Shape dims, mem, ixfun) <- arrayVarReturns v-  let ixfun' = IxFun.permute ixfun perm+  (et, Shape dims, mem, lmad) <- arrayVarReturns v+  let lmad' = LMAD.permute lmad perm       dims' = rearrangeShape perm dims   pure $     Just       [ MemArray et (Shape $ map Free dims') NoUniqueness $           Just $             ReturnsInBlock mem $-              existentialiseIxFun [] ixfun'+              existentialiseLMAD [] lmad'       ] expReturns (BasicOp (Index v slice)) = do   Just . pure . varInfoToExpReturns <$> sliceInfo v slice@@ -1070,15 +1083,15 @@     typeWithDec t p =       case (t, paramDec p) of         ( Array pt shape u,-          MemArray _ _ _ (ArrayIn mem ixfun)+          MemArray _ _ _ (ArrayIn mem lmad)           )             | Just (i, mem_p) <- isLoopVar mem,               Mem space <- paramType mem_p ->-                pure $ MemArray pt shape u $ Just $ ReturnsNewBlock space i ixfun'+                pure $ MemArray pt shape u $ Just $ ReturnsNewBlock space i lmad'             | otherwise ->-                pure $ MemArray pt shape u $ Just $ ReturnsInBlock mem ixfun'+                pure $ MemArray pt shape u $ Just $ ReturnsInBlock mem lmad'             where-              ixfun' = existentialiseIxFun (map paramName mergevars) ixfun+              lmad' = existentialiseLMAD (map paramName mergevars) lmad         (Array {}, _) ->           error "expReturns: Array return type but not array merge variable."         (Acc acc ispace ts u, _) ->@@ -1114,13 +1127,13 @@   Slice SubExp ->   m (MemInfo SubExp NoUniqueness MemBind) sliceInfo v slice = do-  (et, _, mem, ixfun) <- arrayVarReturns v+  (et, _, mem, lmad) <- arrayVarReturns v   case sliceDims slice of     [] -> pure $ MemPrim et     dims ->       pure $         MemArray et (Shape dims) NoUniqueness . ArrayIn mem $-          IxFun.slice ixfun (fmap pe64 slice)+          LMAD.slice lmad (fmap pe64 slice)  flatSliceInfo ::   (Monad m, HasScope rep m, Mem rep inner) =>@@ -1128,10 +1141,10 @@   FlatSlice SubExp ->   m (MemInfo SubExp NoUniqueness MemBind) flatSliceInfo v slice@(FlatSlice offset idxs) = do-  (et, _, mem, ixfun) <- arrayVarReturns v+  (et, _, mem, lmad) <- arrayVarReturns v   map (fmap pe64) idxs     & FlatSlice (pe64 offset)-    & IxFun.flatSlice ixfun+    & LMAD.flatSlice lmad     & MemArray et (Shape (flatSliceDims slice)) NoUniqueness . ArrayIn mem     & pure @@ -1182,11 +1195,11 @@     correctDim (Ext i) = Ext i     correctDim (Free se) = Free $ substSubExp se -    correctSummary (ReturnsNewBlock space i ixfun) =-      ReturnsNewBlock space i ixfun-    correctSummary (ReturnsInBlock mem ixfun) =-      -- FIXME: we should also do a replacement in ixfun here.-      ReturnsInBlock mem' ixfun+    correctSummary (ReturnsNewBlock space i lmad) =+      ReturnsNewBlock space i lmad+    correctSummary (ReturnsInBlock mem lmad) =+      -- FIXME: we should also do a replacement in lmad here.+      ReturnsInBlock mem' lmad       where         mem' = case M.lookup mem parammap of           Just (Var v, _) -> v

@@ -1,286 +0,0 @@-{-# OPTIONS_GHC -fno-warn-redundant-constraints #-}---- | This module contains a representation for the index function based on--- linear-memory accessor descriptors; see Zhu, Hoeflinger and David work.-module Futhark.IR.Mem.IxFun-  ( IxFun (..),-    Shape,-    LMAD (..),-    LMADDim (..),-    index,-    mkExistential,-    iota,-    permute,-    reshape,-    coerce,-    slice,-    flatSlice,-    expand,-    shape,-    rank,-    isDirect,-    substituteInIxFun,-    substituteInLMAD,-    existentialize,-    existentialized,-    closeEnough,-    disjoint,-    disjoint2,-    disjoint3,-  )-where--import Control.Category-import Control.Monad-import Control.Monad.State-import Data.Map.Strict qualified as M-import Data.Traversable-import Futhark.Analysis.PrimExp-import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.Mem.LMAD hiding-  ( equivalent,-    flatSlice,-    index,-    iota,-    isDirect,-    mkExistential,-    permute,-    rank,-    reshape,-    shape,-    slice,-  )-import Futhark.IR.Mem.LMAD qualified as LMAD-import Futhark.IR.Prop-import Futhark.IR.Syntax-  ( FlatSlice (..),-    Slice (..),-    unitSlice,-  )-import Futhark.IR.Syntax.Core (Ext (..))-import Futhark.Transform.Rename-import Futhark.Transform.Substitute-import Futhark.Util.IntegralExp-import Futhark.Util.Pretty-import Prelude hiding (gcd, id, mod, (.))---- | An index function is a mapping from a multidimensional array--- index space (the domain) to a one-dimensional memory index space.--- Essentially, it explains where the element at position @[i,j,p]@ of--- some array is stored inside the flat one-dimensional array that--- constitutes its memory.  For example, we can use this to--- distinguish row-major and column-major representations.------ An index function is represented as an LMAD.-data IxFun num = IxFun-  { ixfunLMAD :: LMAD num,-    -- | the shape of the support array, i.e., the original array-    --   that birthed (is the start point) of this index function.-    base :: Shape num-  }-  deriving (Show, Eq)--instance (Pretty num) => Pretty (IxFun num) where-  pretty (IxFun lmad oshp) =-    braces . semistack $-      [ "base:" <+> brackets (commasep $ map pretty oshp),-        "LMAD:" <+> pretty lmad-      ]--instance (Substitute num) => Substitute (IxFun num) where-  substituteNames substs = fmap $ substituteNames substs--instance (Substitute num) => Rename (IxFun num) where-  rename = substituteRename--instance (FreeIn num) => FreeIn (IxFun num) where-  freeIn' = foldMap freeIn'--instance Functor IxFun where-  fmap = fmapDefault--instance Foldable IxFun where-  foldMap = foldMapDefault---- It is important that the traversal order here is the same as in--- mkExistential.-instance Traversable IxFun where-  traverse f (IxFun lmad oshp) =-    IxFun <$> traverse f lmad <*> traverse f oshp---- | Substitute a name with a PrimExp in an index function.-substituteInIxFun ::-  (Ord a) =>-  M.Map a (TPrimExp t a) ->-  IxFun (TPrimExp t a) ->-  IxFun (TPrimExp t a)-substituteInIxFun tab (IxFun lmad oshp) =-  IxFun-    (substituteInLMAD tab lmad)-    (map (TPrimExp . substituteInPrimExp tab' . untyped) oshp)-  where-    tab' = fmap untyped tab---- | Is this is a row-major array?-isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool-isDirect (IxFun (LMAD offset dims) oshp) =-  let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp))-   in length oshp == length dims-        && offset == 0-        && all-          (\(LMADDim s n, d, se) -> s == se && n == d)-          (zip3 dims oshp strides_expected)---- | The index space of the index function.  This is the same as the--- shape of arrays that the index function supports.-shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num-shape = LMAD.shape . ixfunLMAD---- | Compute the flat memory index for a complete set @inds@ of array indices--- and a certain element size @elem_size@.-index ::-  (IntegralExp num, Eq num) =>-  IxFun num ->-  Indices num ->-  num-index = LMAD.index . ixfunLMAD---- | iota with offset.-iotaOffset :: (IntegralExp num) => num -> Shape num -> IxFun num-iotaOffset o ns = IxFun (LMAD.iota o ns) ns---- | iota.-iota :: (IntegralExp num) => Shape num -> IxFun num-iota = iotaOffset 0---- | Create a single-LMAD index function that is existential in--- everything except shape, with the provided shape.-mkExistential :: Int -> Shape (Ext a) -> Int -> IxFun (Ext a)-mkExistential basis_rank lmad_shape start =-  IxFun (LMAD.mkExistential lmad_shape start) basis-  where-    basis = take basis_rank $ map Ext [start + 1 + length lmad_shape ..]---- | Permute dimensions.-permute ::-  (IntegralExp num) =>-  IxFun num ->-  Permutation ->-  IxFun num-permute (IxFun lmad oshp) perm_new =-  IxFun (LMAD.permute lmad perm_new) oshp---- | Slice an index function.-slice ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Slice num ->-  IxFun num-slice ixfun@(IxFun lmad@(LMAD _ _) oshp) (Slice is)-  -- Avoid identity slicing.-  | is == map (unitSlice 0) (shape ixfun) = ixfun-  | otherwise =-      IxFun (LMAD.slice lmad (Slice is)) oshp---- | Flat-slice an index function.-flatSlice ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  FlatSlice num ->-  IxFun num-flatSlice (IxFun lmad oshp) s = IxFun (LMAD.flatSlice lmad s) oshp---- | Reshape an index function.------ There are four conditions that all must hold for the result of a reshape--- operation to remain in the one-LMAD domain:------   (1) the permutation of the underlying LMAD must leave unchanged---       the LMAD dimensions that were *not* reshape coercions.---   (2) the repetition of dimensions of the underlying LMAD must---       refer only to the coerced-dimensions of the reshape operation.------ If any of these conditions do not hold, then the reshape operation--- will conservatively add a new LMAD to the list, leading to a--- representation that provides less opportunities for further--- analysis-reshape ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Shape num ->-  Maybe (IxFun num)-reshape (IxFun lmad _) new_shape =-  IxFun <$> LMAD.reshape lmad new_shape <*> pure new_shape---- | Coerce an index function to look like it has a new shape.--- Dynamically the shape must be the same.-coerce ::-  (Eq num, IntegralExp num) =>-  IxFun num ->-  Shape num ->-  IxFun num-coerce (IxFun lmad _) new_shape =-  IxFun (onLMAD lmad) new_shape-  where-    onLMAD (LMAD offset dims) = LMAD offset $ zipWith onDim dims new_shape-    onDim ld d = ld {ldShape = d}---- | The number of dimensions in the domain of the input function.-rank :: (IntegralExp num) => IxFun num -> Int-rank (IxFun (LMAD _ sss) _) = length sss---- | Conceptually expand index function to be a particular slice of--- another by adjusting the offset and strides.  Used for memory--- expansion.-expand ::-  (Eq num, IntegralExp num) => num -> num -> IxFun num -> Maybe (IxFun num)-expand o p (IxFun lmad base) =-  let onDim ld = ld {LMAD.ldStride = p * LMAD.ldStride ld}-      lmad' =-        LMAD-          (o + p * LMAD.offset lmad)-          (map onDim (LMAD.dims lmad))-   in Just $ IxFun lmad' base---- | Turn all the leaves of the index function into 'Ext's, except for---  the shape, which where the leaves are simply made 'Free'.-existentialize ::-  Int ->-  IxFun (TPrimExp Int64 a) ->-  IxFun (TPrimExp Int64 (Ext a))-existentialize start (IxFun lmad base) = evalState (IxFun <$> lmad' <*> base') start-  where-    mkExt = do-      i <- get-      put $ i + 1-      pure $ TPrimExp $ LeafExp (Ext i) int64-    lmad' = LMAD <$> mkExt <*> mapM onDim (dims lmad)-    base' = traverse (const mkExt) base-    onDim ld = LMADDim <$> mkExt <*> pure (fmap Free (ldShape ld))---- | Retrieve those elements that 'existentialize' changes. That is,--- everything except the shape (and in the same order as--- 'existentialise' existentialises them).-existentialized :: IxFun a -> [a]-existentialized (IxFun (LMAD offset dims) base) =-  offset : concatMap onDim dims <> base-  where-    onDim (LMADDim ldstride _) = [ldstride]---- | When comparing index functions as part of the type check in KernelsMem,--- we may run into problems caused by the simplifier. As index functions can be--- generalized over if-then-else expressions, the simplifier might hoist some of--- the code from inside the if-then-else (computing the offset of an array, for--- instance), but now the type checker cannot verify that the generalized index--- function is valid, because some of the existentials are computed somewhere--- else. To Work around this, we've had to relax the KernelsMem type-checker--- a bit, specifically, we've introduced this function to verify whether two--- index functions are "close enough" that we can assume that they match. We use--- this instead of `ixfun1 == ixfun2` and hope that it's good enough.-closeEnough :: IxFun num -> IxFun num -> Bool-closeEnough ixf1 ixf2 =-  (length (base ixf1) == length (base ixf2))-    && closeEnoughLMADs (ixfunLMAD ixf1) (ixfunLMAD ixf2)-  where-    closeEnoughLMADs lmad1 lmad2 =-      length (LMAD.dims lmad1) == length (LMAD.dims lmad2)

@@ -4,7 +4,8 @@ -- This module is designed to be used as a qualified import, as the -- exported names are quite generic. module Futhark.IR.Mem.LMAD-  ( Shape,+  ( -- * Core+    Shape,     Indices,     LMAD (..),     LMADDim (..),@@ -13,23 +14,31 @@     slice,     flatSlice,     reshape,+    coerce,     permute,     shape,-    rank,-    substituteInLMAD,+    substitute,+    iota,+    equivalent,+    range,++    -- * Exotic+    expand,+    isDirect,     disjoint,     disjoint2,     disjoint3,     dynamicEqualsLMAD,-    iota,     mkExistential,-    equivalent,-    isDirect,+    closeEnough,+    existentialize,+    existentialized,   ) where  import Control.Category import Control.Monad+import Control.Monad.State import Data.Function (on, (&)) import Data.List (elemIndex, partition, sortBy) import Data.Map.Strict qualified as M@@ -91,7 +100,7 @@ -- However, we expect that the common case is when the index function is one -- LMAD -- we call this the "nice" representation. ----- Finally, the list of LMADs is kept in an @IxFun@ together with the shape of+-- Finally, the list of LMADs is kept in an @LMAD@ together with the shape of -- the original array, and a bit to indicate whether the index function is -- contiguous, i.e., if we instantiate all the points of the current index -- function, do we get a contiguous memory interval?@@ -201,9 +210,20 @@     helper s0 (FlatDimIndex n s) = LMADDim (s0 * s) n flatSlice (LMAD offset []) _ = LMAD offset [] --- | Handle the case where a reshape operation can stay inside a--- single LMAD.  See "Futhark.IR.Mem.IxFun.reshape" for--- conditions.+-- | Reshape an LMAD.+--+-- There are four conditions that all must hold for the result of a reshape+-- operation to remain in the one-LMAD domain:+--+--   (1) the permutation of the underlying LMAD must leave unchanged+--       the LMAD dimensions that were *not* reshape coercions.+--   (2) the repetition of dimensions of the underlying LMAD must+--       refer only to the coerced-dimensions of the reshape operation.+--+-- If any of these conditions do not hold, then the reshape operation+-- will conservatively add a new LMAD to the list, leading to a+-- representation that provides less opportunities for further+-- analysis reshape ::   (Eq num, IntegralExp num) => LMAD num -> Shape num -> Maybe (LMAD num) --@@ -232,13 +252,22 @@   Just $ iotaStrided off base_stride newshape {-# NOINLINE reshape #-} +-- | Coerce an index function to look like it has a new shape.+-- Dynamically the shape must be the same.+coerce :: LMAD num -> Shape num -> LMAD num+coerce (LMAD offset dims) new_shape =+  LMAD offset $ zipWith onDim dims new_shape+  where+    onDim ld d = ld {ldShape = d}+{-# NOINLINE coerce #-}+ -- | Substitute a name with a PrimExp in an LMAD.-substituteInLMAD ::+substitute ::   (Ord a) =>   M.Map a (TPrimExp t a) ->   LMAD (TPrimExp t a) ->   LMAD (TPrimExp t a)-substituteInLMAD tab (LMAD offset dims) =+substitute tab (LMAD offset dims) =   LMAD (sub offset) $ map (\(LMADDim s n) -> LMADDim (sub s) (sub n)) dims   where     tab' = fmap untyped tab@@ -248,10 +277,6 @@ shape :: LMAD num -> Shape num shape = map ldShape . dims --- | Rank of an LMAD.-rank :: LMAD num -> Int-rank = length . shape- iotaStrided ::   (IntegralExp num) =>   -- | Offset@@ -533,8 +558,7 @@ -- Equivalence in this case is matching in offsets and strides. equivalent :: (Eq num) => LMAD num -> LMAD num -> Bool equivalent lmad1 lmad2 =-  length (dims lmad1) == length (dims lmad2)-    && offset lmad1 == offset lmad2+  offset lmad1 == offset lmad2     && map ldStride (dims lmad1) == map ldStride (dims lmad2) {-# NOINLINE equivalent #-} @@ -542,3 +566,69 @@ isDirect :: (Eq num, IntegralExp num) => LMAD num -> Bool isDirect lmad = lmad == iota 0 (map ldShape $ dims lmad) {-# NOINLINE isDirect #-}++-- | The largest possible linear address reachable by this LMAD, not+-- counting the offset. If you add one to this number (and multiply it+-- with the element size), you get the amount of bytes you need to+-- allocate for an array with this LMAD (assuming zero offset).+range :: (Pretty num) => LMAD (TPrimExp Int64 num) -> TPrimExp Int64 num+range lmad =+  -- The idea is that the largest possible offset must be the sum of+  -- the maximum offsets reachable in each dimension, which must be at+  -- either the minimum or maximum index.+  sum (map dimRange $ dims lmad)+  where+    dimRange LMADDim {ldStride, ldShape} =+      0 `sMax64` ((0 `sMax64` (ldShape - 1)) * ldStride)+{-# NOINLINE range #-}++-- | When comparing LMADs as part of the type check in GPUMem, we+-- may run into problems caused by the simplifier. As index functions+-- can be generalized over if-then-else expressions, the simplifier+-- might hoist some of the code from inside the if-then-else+-- (computing the offset of an array, for instance), but now the type+-- checker cannot verify that the generalized index function is valid,+-- because some of the existentials are computed somewhere else. To+-- Work around this, we've had to relax the KernelsMem type-checker a+-- bit, specifically, we've introduced this function to verify whether+-- two index functions are "close enough" that we can assume that they+-- match. We use this instead of `lmad1 == lmad2` and hope that it's+-- good enough.+closeEnough :: LMAD num -> LMAD num -> Bool+closeEnough lmad1 lmad2 =+  length (dims lmad1) == length (dims lmad2)+{-# NOINLINE closeEnough #-}++-- | Turn all the leaves of the LMAD into 'Ext's, except for+--  the shape, which where the leaves are simply made 'Free'.+existentialize ::+  Int ->+  LMAD (TPrimExp Int64 a) ->+  LMAD (TPrimExp Int64 (Ext a))+existentialize start lmad = evalState lmad' start+  where+    mkExt = do+      i <- get+      put $ i + 1+      pure $ TPrimExp $ LeafExp (Ext i) int64+    lmad' = LMAD <$> mkExt <*> mapM onDim (dims lmad)+    onDim ld = LMADDim <$> mkExt <*> pure (fmap Free (ldShape ld))++-- | Retrieve those elements that 'existentialize' changes. That is,+-- everything except the shape (and in the same order as+-- 'existentialise' existentialises them).+existentialized :: LMAD a -> [a]+existentialized (LMAD offset dims) =+  offset : concatMap onDim dims+  where+    onDim (LMADDim ldstride _) = [ldstride]++-- | Conceptually expand LMAD to be a particular slice of+-- another by adjusting the offset and strides.  Used for memory+-- expansion.+expand ::+  (IntegralExp num) => num -> num -> LMAD num -> LMAD num+expand o p lmad =+  LMAD (o + p * offset lmad) (map onDim (dims lmad))+  where+    onDim ld = ld {ldStride = p * ldStride ld}

@@ -15,7 +15,7 @@ import Futhark.Analysis.UsageTable qualified as UT import Futhark.Construct import Futhark.IR.Mem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.Prop.Aliases (AliasedOp) import Futhark.Optimise.Simplify qualified as Simplify import Futhark.Optimise.Simplify.Engine qualified as Engine@@ -145,10 +145,10 @@             zipWithM updateResult (patElems pat) res           updateResult pat_elem (SubExpRes cs (Var v))             | Just mem <- lookup (patElemName pat_elem) arr_to_mem,-              (_, MemArray pt shape u (ArrayIn _ ixfun)) <- patElemDec pat_elem = do+              (_, MemArray pt shape u (ArrayIn _ lmad)) <- patElemDec pat_elem = do                 v_copy <- newVName $ baseString v <> "_nonext_copy"                 let v_pat =-                      Pat [PatElem v_copy $ MemArray pt shape u $ ArrayIn mem ixfun]+                      Pat [PatElem v_copy $ MemArray pt shape u $ ArrayIn mem lmad]                 addStm $ mkWiseStm v_pat (defAux ()) $ BasicOp $ Replicate mempty $ Var v                 pure $ SubExpRes cs $ Var v_copy             | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem =@@ -167,7 +167,7 @@     inContext = (`elem` patNames pat)      hasConcretisableMemory fixable pat_elem-      | (_, MemArray pt shape _ (ArrayIn mem ixfun)) <- patElemDec pat_elem,+      | (_, MemArray pt shape _ (ArrayIn mem lmad)) <- patElemDec pat_elem,         Just (j, Mem space) <-           fmap patElemType             <$> find@@ -176,12 +176,11 @@         Just cases_ses <- mapM (maybeNth j . bodyResult . caseBody) cases,         Just defbody_se <- maybeNth j $ bodyResult defbody,         mem `onlyUsedIn` patElemName pat_elem,-        length (IxFun.base ixfun) == shapeRank shape, -- See #1325         all knownSize (shapeDims shape),-        not $ freeIn ixfun `namesIntersect` namesFromList (patNames pat),-        any (defbody_se /=) cases_ses =-          let mem_size =-                untyped $ product $ primByteSize pt : map sExt64 (IxFun.base ixfun)+        not $ freeIn lmad `namesIntersect` namesFromList (patNames pat),+        any (defbody_se /=) cases_ses,+        LMAD.offset lmad == 0 =+          let mem_size = untyped $ primByteSize pt * (1 + LMAD.range lmad)            in (pat_elem, mem_size, mem, space) : fixable       | otherwise =           fixable

@@ -32,7 +32,7 @@ import Futhark.IR.MC.Op qualified as MC import Futhark.IR.MCMem (MCMem) import Futhark.IR.Mem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SOACS (SOACS) import Futhark.IR.SOACS.SOAC qualified as SOAC import Futhark.IR.SegOp qualified as SegOp@@ -641,11 +641,27 @@ pOpaqueType =   (,)     <$> (keyword "type" *> (nameFromText <$> pStringLiteral) <* pEqual)-    <*> choice [pRecord, pOpaque]+    <*> choice [pRecord, pSum, pOpaque]   where     pFieldName = choice [pName, nameFromString . show <$> pInt]     pField = (,) <$> pFieldName <* pColon <*> pEntryPointType     pRecord = keyword "record" $> OpaqueRecord <*> braces (many pField)++    pConstructor = "#" *> pName+    pPayload =+      parens $+        (,)+          <$> (pEntryPointType <* pComma)+          <*> brackets (pInt `sepBy` pComma)+    pVariant = (,) <$> pConstructor <*> many pPayload+    pSum =+      keyword "sum"+        *> braces+          ( OpaqueSum+              <$> brackets (pValueType `sepBy` pComma)+              <*> many pVariant+          )+     pOpaque = keyword "opaque" $> OpaqueType <*> braces (many pValueType)  pOpaqueTypes :: Parser OpaqueTypes@@ -961,19 +977,14 @@   where     pMCSegOp = pSegOp pr (void $ lexeme "()") -pIxFunBase :: Parser a -> Parser (IxFun.IxFun a)-pIxFunBase pNum =-  braces $ do-    base <- pLab "base" $ brackets (pNum `sepBy` pComma) <* pSemi-    lmad <- pLab "LMAD" pLMAD-    pure $ IxFun.IxFun lmad base+pLMADBase :: Parser a -> Parser (LMAD.LMAD a)+pLMADBase pNum = braces $ do+  offset <- pLab "offset" pNum <* pSemi+  strides <- pLab "strides" $ brackets (pNum `sepBy` pComma) <* pSemi+  shape <- pLab "shape" $ brackets (pNum `sepBy` pComma)+  pure $ LMAD.LMAD offset $ zipWith LMAD.LMADDim strides shape   where     pLab s m = keyword s *> pColon *> m-    pLMAD = braces $ do-      offset <- pLab "offset" pNum <* pSemi-      strides <- pLab "strides" $ brackets (pNum `sepBy` pComma) <* pSemi-      shape <- pLab "shape" $ brackets (pNum `sepBy` pComma)-      pure $ IxFun.LMAD offset $ zipWith IxFun.LMADDim strides shape  pPrimExpLeaf :: Parser VName pPrimExpLeaf = pVName@@ -981,11 +992,11 @@ pExtPrimExpLeaf :: Parser (Ext VName) pExtPrimExpLeaf = pExt pVName -pIxFun :: Parser IxFun-pIxFun = pIxFunBase $ isInt64 <$> pPrimExp int64 pPrimExpLeaf+pLMAD :: Parser LMAD+pLMAD = pLMADBase $ isInt64 <$> pPrimExp int64 pPrimExpLeaf -pExtIxFun :: Parser ExtIxFun-pExtIxFun = pIxFunBase $ isInt64 <$> pPrimExp int64 pExtPrimExpLeaf+pExtLMAD :: Parser ExtLMAD+pExtLMAD = pLMADBase $ isInt64 <$> pPrimExp int64 pExtPrimExpLeaf  pMemInfo :: Parser d -> Parser u -> Parser ret -> Parser (MemInfo d u ret) pMemInfo pd pu pret =@@ -1023,16 +1034,16 @@       ]  pMemBind :: Parser MemBind-pMemBind = ArrayIn <$> pVName <* lexeme "->" <*> pIxFun+pMemBind = ArrayIn <$> pVName <* lexeme "->" <*> pLMAD  pMemReturn :: Parser MemReturn pMemReturn =   choice-    [ ReturnsInBlock <$> pVName <* lexeme "->" <*> pExtIxFun,+    [ ReturnsInBlock <$> pVName <* lexeme "->" <*> pExtLMAD,       do         i <- "?" *> pInt         space <- choice [pSpace, pure DefaultSpace] <* lexeme "->"-        ReturnsNewBlock space i <$> pExtIxFun+        ReturnsNewBlock space i <$> pExtLMAD     ]  pRetTypeMem :: Parser RetTypeMem

@@ -404,6 +404,10 @@     "record" <+> nestedBlock "{" "}" (stack $ map p fs)     where       p (f, et) = pretty f <> ":" <+> pretty et+  pretty (OpaqueSum ts cs) =+    "sum" <+> nestedBlock "{" "}" (stack $ pretty ts : map p cs)+    where+      p (c, ets) = hsep $ "#" <> pretty c : map pretty ets  instance Pretty OpaqueTypes where   pretty (OpaqueTypes ts) = "types" <+> nestedBlock "{" "}" (stack $ map p ts)

@@ -620,14 +620,20 @@   -- value.   fixExt :: Int -> SubExp -> t -> t +  -- | Map a function onto any existential.+  mapExt :: (Int -> Int) -> t -> t+ instance (FixExt shape, ArrayShape shape) => FixExt (TypeBase shape u) where   fixExt i se = modifyArrayShape $ fixExt i se+  mapExt f = modifyArrayShape $ mapExt f  instance (FixExt d) => FixExt (ShapeBase d) where   fixExt i se = fmap $ fixExt i se+  mapExt f = fmap $ mapExt f  instance (FixExt a) => FixExt [a] where   fixExt i se = fmap $ fixExt i se+  mapExt f = fmap $ mapExt f  instance FixExt ExtSize where   fixExt i se (Ext j)@@ -636,5 +642,9 @@     | otherwise = Ext j   fixExt _ _ (Free x) = Free x +  mapExt f (Ext i) = Ext $ f i+  mapExt _ (Free x) = Free x+ instance FixExt () where   fixExt _ _ () = ()+  mapExt _ () = ()

@@ -576,15 +576,14 @@         unzip6 $ filter isUsed $ zip6 (patElems pat) i_ses v_ses i_ts v_ts dests       fun' =         fun-          { lambdaBody = (lambdaBody fun) {bodyResult = concat i_ses' ++ v_ses'},+          { lambdaBody =+              mkBody (bodyStms (lambdaBody fun)) (concat i_ses' ++ v_ses'),             lambdaReturnType = concat i_ts' ++ v_ts'           }    in if pat /= Pat pat'         then-          Simplify . auxing aux $-            letBind (Pat pat') $-              Op $-                Scatter w arrs fun' dests'+          Simplify . auxing aux . letBind (Pat pat') $+            Op (Scatter w arrs fun' dests')         else Skip removeDeadWrite _ _ _ _ = Skip 

@@ -599,6 +599,15 @@   | -- | Note that the field ordering here denote the actual     -- representation - make sure it is preserved.     OpaqueRecord [(Name, EntryPointType)]+  | -- | Constructor ordering also denotes representation, in that the+    -- index of the constructor is the identifying number.+    --+    -- The total values used to represent a sum values is the+    -- 'ValueType' list. The 'Int's associated with each+    -- 'EntryPointType' are the indexes of the values used to+    -- represent that constructor payload. This is necessary because+    -- we deduplicate payloads across constructors.+    OpaqueSum [ValueType] [(Name, [(EntryPointType, [Int])])]   deriving (Eq, Ord, Show)  -- | Names of opaque types and their representation.

@@ -558,6 +558,8 @@       descend (name : known) ts     check known (OpaqueRecord fs) =       mapM_ (checkEntryPointType known . snd) fs+    check known (OpaqueSum _ cs) =+      mapM_ (mapM_ (checkEntryPointType known . fst) . snd) cs     check _ (OpaqueType _) =       pure ()     checkEntryPointType known (TypeOpaque s) =

@@ -8,10 +8,11 @@  import Control.Monad import Control.Monad.State-import Data.List (find)+import Data.List (find, intersperse) import Data.Map qualified as M import Futhark.IR qualified as I-import Futhark.Internalise.TypesValues (internalisedTypeSize)+import Futhark.Internalise.TypesValues (internaliseSumTypeRep, internalisedTypeSize)+import Futhark.Util (chunks) import Futhark.Util.Pretty (prettyTextOneLine) import Language.Futhark qualified as E hiding (TypeArg) import Language.Futhark.Core (Name, Uniqueness (..), VName, nameFromText)@@ -49,15 +50,29 @@ rootType :: E.TypeExp E.Info VName -> E.TypeExp E.Info VName rootType (E.TEApply te E.TypeArgExpSize {} _) = rootType te rootType (E.TEUnique te _) = rootType te+rootType (E.TEDim _ te _) = rootType te rootType te = te  typeExpOpaqueName :: E.TypeExp E.Info VName -> Name-typeExpOpaqueName = f . rootType+typeExpOpaqueName = nameFromText . f   where-    f (E.TEArray _ te _) =+    f = g . rootType+    g (E.TEArray _ te _) =       let (d, te') = withoutDims te-       in nameFromText (mconcat (replicate (1 + d) "[]")) <> typeExpOpaqueName te'-    f te = nameFromText $ prettyTextOneLine te+       in mconcat (replicate (1 + d) "[]") <> f te'+    g (E.TETuple tes _) =+      "(" <> mconcat (intersperse ", " (map f tes)) <> ")"+    g (E.TERecord tes _) =+      "{" <> mconcat (intersperse ", " (map onField tes)) <> "}"+      where+        onField (k, te) = E.nameToText k <> ":" <> f te+    g (E.TESum cs _) =+      mconcat (intersperse " | " (map onConstr cs))+      where+        onConstr (k, tes) =+          E.nameToText k <> ":" <> mconcat (intersperse " " (map f tes))+    g (E.TEParens te _) = "(" <> f te <> ")"+    g te = prettyTextOneLine te  type GenOpaque = State I.OpaqueTypes @@ -105,6 +120,38 @@   where     opaqueField e_t i_ts = snd <$> entryPointType types e_t i_ts +isSum :: VisibleTypes -> E.TypeExp E.Info VName -> Maybe (M.Map Name [E.TypeExp E.Info VName])+isSum _ (E.TESum cs _) = Just $ M.fromList cs+isSum types (E.TEVar v _) = isSum types =<< findType (E.qualLeaf v) types+isSum _ _ = Nothing++sumConstrs ::+  VisibleTypes ->+  M.Map Name [E.StructType] ->+  Maybe (E.TypeExp E.Info VName) ->+  [(Name, [E.EntryType])]+sumConstrs types cs t =+  case isSum types . rootType =<< t of+    Just e_cs ->+      zipWith f (E.sortConstrs cs) (E.sortConstrs e_cs)+      where+        f (k, c_ts) (_, e_c_ts) = (k, zipWith E.EntryType c_ts $ map Just e_c_ts)+    Nothing ->+      map (fmap (map (`E.EntryType` Nothing))) $ E.sortConstrs cs++opaqueSum ::+  VisibleTypes ->+  [(Name, ([E.EntryType], [Int]))] ->+  [I.TypeBase I.Rank Uniqueness] ->+  GenOpaque [(Name, [(I.EntryPointType, [Int])])]+opaqueSum types cs ts = mapM (traverse f) cs+  where+    f (ets, is) = do+      let ns = map (internalisedTypeSize . E.entryType) ets+          is' = chunks ns is+      ets' <- map snd <$> zipWithM (entryPointType types) ets (map (map (ts !!)) is')+      pure $ zip ets' $ map (map (+ 1)) is' -- Adjust for tag.+ entryPointType ::   VisibleTypes ->   E.EntryType ->@@ -129,6 +176,12 @@           | not $ null fs ->               let fs' = recordFields types fs $ E.entryAscribed t                in addType desc . I.OpaqueRecord =<< opaqueRecord types fs' ts+        E.Scalar (E.Sum cs) -> do+          let (_, places) = internaliseSumTypeRep cs+              cs' = sumConstrs types cs $ E.entryAscribed t+              cs'' = zip (map fst cs') (zip (map snd cs') (map snd places))+          addType desc . I.OpaqueSum (map valueType ts)+            =<< opaqueSum types cs'' (drop 1 ts)         _ -> addType desc $ I.OpaqueType $ map valueType ts       pure (u, I.TypeOpaque desc)   where

@@ -23,7 +23,7 @@ import Futhark.Internalise.Monad as I import Futhark.Internalise.TypesValues import Futhark.Transform.Rename as I-import Futhark.Util (splitAt3)+import Futhark.Util (lookupWithIndex, splitAt3) import Futhark.Util.Pretty (align, docText, pretty) import Language.Futhark as E hiding (TypeArg) import Language.Futhark.TypeChecker.Types qualified as E@@ -796,7 +796,7 @@   let noExt _ = pure $ intConst Int64 0   ts' <- instantiateShapes noExt $ map fromDecl ts -  case M.lookup c constr_map of+  case lookupWithIndex c constr_map of     Just (i, js) ->       (intConst Int8 (toInteger i) :) <$> clauses 0 ts' (zip js es')     Nothing ->@@ -806,7 +806,17 @@       | Just e <- j `lookup` js_to_es =           (e :) <$> clauses (j + 1) ts js_to_es       | otherwise = do-          blank <- letSubExp "zero" =<< eBlank t+          blank <-+            -- Cannot use eBlank here for arrays, because when doing+            -- equality comparisons on sum types, we end up looking at+            -- the array elements. (#2081) This is a bit of an edge+            -- case, but arrays in sum types are known to be+            -- inefficient.+            letSubExp "zero"+              =<< case t of+                I.Array {} ->+                  pure $ BasicOp $ Replicate (I.arrayShape t) $ I.Constant $ blankPrimValue $ elemType t+                _ -> eBlank t           (blank :) <$> clauses (j + 1) ts js_to_es     clauses _ [] _ =       pure []@@ -889,7 +899,7 @@       pure ([Just $ internalisePatLit l t], [se], ses)     compares (E.PatConstr c (Info (E.Scalar (E.Sum fs))) pats _) (_ : ses) = do       (payload_ts, m) <- internaliseSumType $ M.map (map toStruct) fs-      case M.lookup c m of+      case lookupWithIndex c m of         Just (tag, payload_is) -> do           let (payload_ses, ses') = splitAt (length payload_ts) ses           (cmps, pertinent, _) <-

@@ -9,6 +9,7 @@     internaliseLoopParamType,     internalisePrimType,     internalisedTypeSize,+    internaliseSumTypeRep,     internaliseSumType,     Tree, @@ -24,6 +25,7 @@ import Control.Monad import Control.Monad.Free (Free (..)) import Control.Monad.State+import Data.Bifunctor import Data.Bitraversable (bitraverse) import Data.Foldable (toList) import Data.List (delete, find, foldl')@@ -285,15 +287,15 @@ internaliseConstructors ::   M.Map Name [Tree (I.TypeBase ExtShape Uniqueness)] ->   ( [Tree (I.TypeBase ExtShape Uniqueness)],-    M.Map Name (Int, [Int])+    [(Name, [Int])]   ) internaliseConstructors cs =-  foldl' onConstructor mempty $ zip (E.sortConstrs cs) [0 ..]+  L.mapAccumL onConstructor mempty $ E.sortConstrs cs   where-    onConstructor (ts, mapping) ((c, c_ts), i) =+    onConstructor ts (c, c_ts) =       let (_, js, new_ts) =             foldl' f (withOffsets (map (fmap fromDecl) ts), mempty, mempty) c_ts-       in (ts ++ new_ts, M.insert c (i, js) mapping)+       in (ts ++ new_ts, (c, js))       where         size = sum . map length         f (ts', js, new_ts) t@@ -308,16 +310,24 @@                 new_ts ++ [t]               ) +internaliseSumTypeRep ::+  M.Map Name [E.StructType] ->+  ( [I.TypeBase ExtShape Uniqueness],+    [(Name, [Int])]+  )+internaliseSumTypeRep cs =+  first (foldMap toList) . runInternaliseTypeM $+    internaliseConstructors+      <$> traverse (fmap concat . mapM (internaliseTypeM mempty . E.toRes E.Nonunique)) cs+ internaliseSumType ::   M.Map Name [E.StructType] ->   InternaliseM     ( [I.TypeBase ExtShape Uniqueness],-      M.Map Name (Int, [Int])+      [(Name, [Int])]     )-internaliseSumType cs =-  bitraverse (mapM mkAccCerts . foldMap toList) pure . runInternaliseTypeM $-    internaliseConstructors-      <$> traverse (fmap concat . mapM (internaliseTypeM mempty . E.toRes E.Nonunique)) cs+internaliseSumType =+  bitraverse (mapM mkAccCerts) pure . internaliseSumTypeRep  -- | How many core language values are needed to represent one source -- language value of the given type?

@@ -20,7 +20,7 @@ import Futhark.IR.Aliases import Futhark.IR.GPUMem import Futhark.IR.MCMem-import Futhark.IR.Mem.IxFun (substituteInIxFun)+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SeqMem import Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing import Futhark.Optimise.ArrayShortCircuiting.DataStructs@@ -193,7 +193,7 @@     coaltab <- asks envCoalesceTab     if any (M.member vname . vartab) coaltab       then-        existentialiseIxFun (map patElemName pat_elems) ixf+        existentialiseLMAD (map patElemName pat_elems) ixf           & ReturnsInBlock mem           & MemArray pt shp u           & pure@@ -221,7 +221,7 @@   case M.lookup vname $ foldMap vartab coaltab of     Just (Coalesced _ (MemBlock pt shp mem ixf) subs) ->       ixf-        & fixPoint (substituteInIxFun subs)+        & fixPoint (LMAD.substitute subs)         & ArrayIn mem         & MemArray pt shp u         & f vname

@@ -26,7 +26,7 @@ import Futhark.IR.Aliases import Futhark.IR.GPUMem as GPU import Futhark.IR.MCMem as MC-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SeqMem import Futhark.MonadFreshNames import Futhark.Optimise.ArrayShortCircuiting.DataStructs@@ -405,8 +405,8 @@                   maybe                     Undeterminable                     ( ixfunToAccessSummary-                        . IxFun.slice ixf-                        . fullSlice (IxFun.shape ixf)+                        . LMAD.slice ixf+                        . fullSlice (LMAD.shape ixf)                     )                     $ threadSlice space res                 Nothing -> mempty@@ -580,7 +580,7 @@       unSegSpace space         & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)         & Slice-    resultSlice ixf = IxFun.slice ixf $ fullSlice (IxFun.shape ixf) thread_slice+    resultSlice ixf = LMAD.slice ixf $ fullSlice (LMAD.shape ixf) thread_slice makeSegMapCoals _ _ td_env _ _ x (_, _, WriteReturns _ return_name _) =   case getScopeMemInfo return_name $ scope td_env of     Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem@@ -1256,8 +1256,8 @@                                       ((M.insert mb info' a_acc, inhb), s_acc)                         _ -> (failed, s_acc) -- fail! -ixfunToAccessSummary :: IxFun.IxFun (TPrimExp Int64 VName) -> AccessSummary-ixfunToAccessSummary (IxFun.IxFun lmad _) = Set $ S.singleton lmad+ixfunToAccessSummary :: LMAD.LMAD (TPrimExp Int64 VName) -> AccessSummary+ixfunToAccessSummary = Set . S.singleton  -- | Check safety conditions 2 and 5 and update new substitutions: -- called on the pat-elements of loop and if-then-else expressions.@@ -1406,13 +1406,13 @@ -- | Information about a particular short-circuit point type SSPointInfo =   ( CoalescedKind,-    IxFun -> IxFun,+    LMAD -> LMAD,     VName,     VName,-    IxFun,+    LMAD,     VName,     VName,-    IxFun,+    LMAD,     PrimType,     Shape,     Certs@@ -1444,7 +1444,7 @@ --  3. The array being indexed is last-used in that statement, is free in the --  'SegMap', is unique or has been recently allocated (specifically, it should --  not be a non-unique argument to the enclosing function), has elements with---  the same bit-size as the pattern elements, and has the exact same 'IxFun' as+--  the same bit-size as the pattern elements, and has the exact same 'LMAD' as --  the pattern of the 'SegMap' statement. -- -- There can be multiple candidate arrays, but the current implementation will@@ -1541,10 +1541,10 @@     b `nameIn` last_uses =       pure $ Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)]   where-    updateIndFunSlice :: IxFun -> Slice SubExp -> IxFun+    updateIndFunSlice :: LMAD -> Slice SubExp -> LMAD     updateIndFunSlice ind_fun slc_x =       let slc_x' = map (fmap pe64) $ unSlice slc_x-       in IxFun.slice ind_fun $ Slice slc_x'+       in LMAD.slice ind_fun $ Slice slc_x' genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (FlatUpdate x slice_x b)))   | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat,     Just last_uses <- M.lookup x' lutab,@@ -1553,9 +1553,9 @@     b `nameIn` last_uses =       pure $ Just [(InPlaceCoal, (`updateIndFunSlice` slice_x), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)]   where-    updateIndFunSlice :: IxFun -> FlatSlice SubExp -> IxFun+    updateIndFunSlice :: LMAD -> FlatSlice SubExp -> LMAD     updateIndFunSlice ind_fun (FlatSlice offset dims) =-      IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims+      LMAD.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims  -- CASE b) @let x = concat(a, b^{lu})@ genCoalStmtInfo lutab td_env scopetab (Let pat aux (BasicOp (Concat concat_dim (b0 :| bs) _)))@@ -1578,7 +1578,7 @@                       map (unitSlice zero . pe64) (take concat_dim dims)                         <> [unitSlice offs (pe64 d)]                         <> map (unitSlice zero . pe64) (drop (concat_dim + 1) dims)-               in ( acc ++ [(ConcatCoal, (`IxFun.slice` slc), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)],+               in ( acc ++ [(ConcatCoal, (`LMAD.slice` slc), x, m_x, ind_x, b, m_b, ind_b, tpb, shpb, stmAuxCerts aux)],                     offs',                     True                   )@@ -1655,7 +1655,7 @@       -- by definition of if-stmt, r and b have the same basic type, shape and       -- index function, hence, for example, do not need to rebase       -- We will check whether it is translatable at the definition point of r.-      let ind_r = IxFun.substituteInIxFun exist_subs ind_b+      let ind_r = LMAD.substitute exist_subs ind_b           subst_r = M.union exist_subs subst_b           mem_info = Coalesced knd (MemBlock btp shp (dstmem etry) ind_r) subst_r        in if m_r == m_b -- already unified, just add binding for @r@

@@ -40,7 +40,7 @@ import Futhark.IR.Aliases import Futhark.IR.GPUMem as GPU import Futhark.IR.MCMem as MC-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SeqMem import Futhark.Util.Pretty hiding (line, sep, (</>)) import Prelude@@ -50,7 +50,7 @@ type ScopeTab rep = Scope (Aliases rep)  -- | An LMAD specialized to TPrimExps (a typed primexp)-type LmadRef = IxFun.LMAD (TPrimExp Int64 VName)+type LmadRef = LMAD.LMAD (TPrimExp Int64 VName)  -- | Summary of all memory accesses at a given point in the code data AccessSummary@@ -110,7 +110,7 @@   { primType :: PrimType,     shape :: Shape,     memName :: VName,-    ixfun :: IxFun+    ixfun :: LMAD   }  -- | Free variable substitutions@@ -132,7 +132,7 @@   { -- | destination memory block     dstmem :: VName,     -- | index function of the destination (used for rebasing)-    dstind :: IxFun,+    dstind :: LMAD,     -- | aliased destination memory blocks can appear     --   due to repeated (optimistic) coalescing.     alsmem :: Names,

@@ -24,7 +24,7 @@ import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Aliases import Futhark.IR.Mem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.MonadFreshNames import Futhark.Optimise.ArrayShortCircuiting.DataStructs import Futhark.Optimise.ArrayShortCircuiting.TopdownAnalysis@@ -69,7 +69,7 @@ translateAccessSummary scope0 scals0 (Set slmads)   | Just subs <- freeVarSubstitutions scope0 scals0 slmads =       slmads-        & S.map (IxFun.substituteInLMAD subs)+        & S.map (LMAD.substitute subs)         & Set translateAccessSummary _ _ _ = Undeterminable @@ -81,12 +81,12 @@   Stm (Aliases rep) ->   -- | A pair of written and written+read memory locations, along with their   -- associated array and the index function used-  Maybe ([(VName, VName, IxFun)], [(VName, VName, IxFun)])+  Maybe ([(VName, VName, LMAD)], [(VName, VName, LMAD)]) getUseSumFromStm td_env coal_tab (Let _ _ (BasicOp (Index arr (Slice slc))))   | Just (MemBlock _ shp _ _) <- getScopeMemInfo arr (scope td_env),     length slc == length (shapeDims shp) && all isFix slc = do       (mem_b, mem_arr, ixfn_arr) <- getDirAliasedIxfn td_env coal_tab arr-      let new_ixfn = IxFun.slice ixfn_arr $ Slice $ map (fmap pe64) slc+      let new_ixfn = LMAD.slice ixfn_arr $ Slice $ map (fmap pe64) slc       pure ([], [(mem_b, mem_arr, new_ixfn)])   where     isFix DimFix {} = True@@ -105,7 +105,7 @@ --   been added in the active coalesced table. getUseSumFromStm td_env coal_tab (Let (Pat [x']) _ (BasicOp (Update _ _x (Slice slc) a_se))) = do   (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x')-  let x_ixfn_slc = IxFun.slice x_ixfn $ Slice $ map (fmap pe64) slc+  let x_ixfn_slc = LMAD.slice x_ixfn $ Slice $ map (fmap pe64) slc       r1 = (m_b, m_x, x_ixfn_slc)   case a_se of     Constant _ -> Just ([r1], [r1])@@ -136,7 +136,7 @@ getUseSumFromStm td_env coal_tab (Let (Pat [x]) _ (BasicOp (FlatUpdate _ (FlatSlice offset slc) v)))   | Just (m_b, m_x, x_ixfn) <- getDirAliasedIxfn td_env coal_tab (patElemName x) = do       let x_ixfn_slc =-            IxFun.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc+            LMAD.flatSlice x_ixfn $ FlatSlice (pe64 offset) $ map (fmap pe64) slc       let r1 = (m_b, m_x, x_ixfn_slc)       case getDirAliasedIxfn td_env coal_tab v of         Nothing -> Just ([r1], [r1])@@ -251,9 +251,8 @@         <> acc         <> fromMaybe mempty (M.lookup m (m_alias td_env))     mbLmad indfun-      | Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun,-        (IxFun.IxFun lmad _) <- IxFun.substituteInIxFun subs indfun =-          Just lmad+      | Just subs <- freeVarSubstitutions (scope td_env) (scals bu_env) indfun =+          Just $ LMAD.substitute subs indfun     mbLmad _ = Nothing     addLmads wrts uses etry =       etry {memrefs = MemRefs uses wrts <> memrefs etry}@@ -274,9 +273,9 @@               ( \i ->                   all                     ( \j ->-                        IxFun.disjoint less_thans (nonNegatives td_env) i j-                          || IxFun.disjoint2 () () less_thans (nonNegatives td_env) i j-                          || IxFun.disjoint3 (typeOf <$> scope td_env) asserts less_thans non_negs i j+                        LMAD.disjoint less_thans (nonNegatives td_env) i j+                          || LMAD.disjoint2 () () less_thans (nonNegatives td_env) i j+                          || LMAD.disjoint3 (typeOf <$> scope td_env) asserts less_thans non_negs i j                     )                     js               )@@ -285,8 +284,8 @@   where     less_thans = map (fmap $ fixPoint $ substituteInPrimExp $ scalarTable td_env) $ knownLessThan td_env     asserts = map (fixPoint (substituteInPrimExp $ scalarTable td_env) . primExpFromSubExp Bool) $ td_asserts td_env-    is = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList is0-    js = map (fixPoint (IxFun.substituteInLMAD $ TPrimExp <$> scalarTable td_env)) $ S.toList js0+    is = map (fixPoint (LMAD.substitute $ TPrimExp <$> scalarTable td_env)) $ S.toList is0+    js = map (fixPoint (LMAD.substitute $ TPrimExp <$> scalarTable td_env)) $ S.toList js0 noMemOverlap _ _ _ = False  -- | Computes the total aggregated access summary for a loop by expanding the@@ -322,7 +321,7 @@ aggSummaryLoopTotal _ _ scalars_loop (Just (iterator_var, (lower_bound, upper_bound))) (Set lmads) =   concatMapM     ( aggSummaryOne iterator_var lower_bound upper_bound-        . fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)+        . fixPoint (LMAD.substitute $ fmap TPrimExp scalars_loop)     )     (S.toList lmads) aggSummaryLoopTotal _ _ _ _ _ = pure Undeterminable@@ -355,7 +354,7 @@         iterator_var         (isInt64 (LeafExp iterator_var $ IntType Int64) + 1)         (upper_bound - typedLeafExp iterator_var - 1)-        . fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars_loop)+        . fixPoint (LMAD.substitute $ fmap TPrimExp scalars_loop)     )     (S.toList lmads) @@ -409,7 +408,7 @@       )     ]   where-    lmads = map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $ S.toList lmads0+    lmads = map (fixPoint (LMAD.substitute $ fmap TPrimExp scalars)) $ S.toList lmads0     helper (x, y) = concatMapM (aggSummaryOne gtid x y) lmads  -- | Computes to total access summary over a multi-dimensional map.@@ -435,7 +434,7 @@   where     lmads =       S.fromList $-        map (fixPoint (IxFun.substituteInLMAD $ fmap TPrimExp scalars)) $+        map (fixPoint (LMAD.substitute $ fmap TPrimExp scalars)) $           S.toList lmads0  -- | Helper function that aggregates the accesses of single LMAD according to a@@ -448,7 +447,7 @@ -- The function returns 'Underterminable' if the iterator is free in the output -- LMAD or the dimensions of the input LMAD . aggSummaryOne :: (MonadFreshNames m) => VName -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> LmadRef -> m AccessSummary-aggSummaryOne iterator_var lower_bound spn lmad@(IxFun.LMAD offset0 dims0)+aggSummaryOne iterator_var lower_bound spn lmad@(LMAD.LMAD offset0 dims0)   | iterator_var `nameIn` freeIn dims0 = pure Undeterminable   | iterator_var `notNameIn` freeIn offset0 = pure $ Set $ S.singleton lmad   | otherwise = do@@ -458,7 +457,7 @@           new_stride = TPrimExp $ constFoldPrimExp $ simplify $ untyped $ offsetp1 - offset           new_offset = replaceIteratorWith lower_bound offset0           new_lmad =-            IxFun.LMAD new_offset $ IxFun.LMADDim new_stride spn : dims0+            LMAD.LMAD new_offset $ LMAD.LMADDim new_stride spn : dims0       if new_var `nameIn` freeIn new_lmad         then pure Undeterminable         else pure $ Set $ S.singleton new_lmad

@@ -24,13 +24,13 @@ import Futhark.IR.Aliases import Futhark.IR.GPUMem as GPU import Futhark.IR.MCMem as MC-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Optimise.ArrayShortCircuiting.DataStructs -type DirAlias = IxFun -> Maybe IxFun+type DirAlias = LMAD -> Maybe LMAD -- ^ A direct aliasing transformation -type InvAlias = Maybe (IxFun -> IxFun)+type InvAlias = Maybe (LMAD -> LMAD) -- ^ An inverse aliasing transformation  type VarAliasTab = M.Map VName (VName, DirAlias, InvAlias)@@ -72,18 +72,18 @@ getDirAliasFromExp (BasicOp (SubExp (Var x))) = Just (x, Just) getDirAliasFromExp (BasicOp (Opaque _ (Var x))) = Just (x, Just) getDirAliasFromExp (BasicOp (Reshape ReshapeCoerce shp x)) =-  Just (x, Just . (`IxFun.coerce` shapeDims (fmap pe64 shp)))+  Just (x, Just . (`LMAD.coerce` shapeDims (fmap pe64 shp))) getDirAliasFromExp (BasicOp (Reshape ReshapeArbitrary shp x)) =-  Just (x, (`IxFun.reshape` shapeDims (fmap pe64 shp)))+  Just (x, (`LMAD.reshape` shapeDims (fmap pe64 shp))) getDirAliasFromExp (BasicOp (Rearrange _ _)) =   Nothing getDirAliasFromExp (BasicOp (Index x slc)) =-  Just (x, Just . (`IxFun.slice` (Slice $ map (fmap pe64) $ unSlice slc)))+  Just (x, Just . (`LMAD.slice` (Slice $ map (fmap pe64) $ unSlice slc))) getDirAliasFromExp (BasicOp (Update _ x _ _elm)) = Just (x, Just) getDirAliasFromExp (BasicOp (FlatIndex x (FlatSlice offset idxs))) =   Just     ( x,-      Just . (`IxFun.flatSlice` FlatSlice (pe64 offset) (map (fmap pe64) idxs))+      Just . (`LMAD.flatSlice` FlatSlice (pe64 offset) (map (fmap pe64) idxs))     ) getDirAliasFromExp (BasicOp (FlatUpdate x _ _)) = Just (x, Just) getDirAliasFromExp _ = Nothing@@ -110,7 +110,7 @@ getInvAliasFromExp (BasicOp (Opaque _ (Var _))) = Just id getInvAliasFromExp (BasicOp Update {}) = Just id getInvAliasFromExp (BasicOp (Rearrange perm _)) =-  Just (`IxFun.permute` rearrangeInverse perm)+  Just (`LMAD.permute` rearrangeInverse perm) getInvAliasFromExp _ = Nothing  class TopDownHelper inner where@@ -225,7 +225,7 @@ -- | Get direct aliased index function.  Returns a triple of current memory -- block to be coalesced, the destination memory block and the index function of -- the access in the space of the destination block.-getDirAliasedIxfn :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)+getDirAliasedIxfn :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD) getDirAliasedIxfn td_env coals_tab x =   case getScopeMemInfo x (scope td_env) of     Just (MemBlock _ _ m_x orig_ixfun) ->@@ -241,7 +241,7 @@  -- | Like 'getDirAliasedIxfn', but this version returns 'Nothing' if the value -- is not currently subject to coalescing.-getDirAliasedIxfn' :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, IxFun)+getDirAliasedIxfn' :: (HasMemBlock (Aliases rep)) => TopdownEnv rep -> CoalsTab -> VName -> Maybe (VName, VName, LMAD) getDirAliasedIxfn' td_env coals_tab x =   case getScopeMemInfo x (scope td_env) of     Just (MemBlock _ _ m_x _) ->

@@ -24,7 +24,6 @@ import Data.Maybe import Data.Sequence qualified as Seq import Futhark.IR.GPU-import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.MonadFreshNames import Futhark.Optimise.TileLoops.Shared@@ -140,14 +139,13 @@       isInnerCoal (_, ixfn_env) slc_X (Let pat _ (BasicOp (Index x _)))         | [slc_X'] <- patNames pat,           slc_X == slc_X',-          Just ixf_fn <- M.lookup x ixfn_env,-          (IxFun.IxFun lmad _) <- ixf_fn =+          Just lmad <- M.lookup x ixfn_env =             innerHasStride1 lmad       isInnerCoal _ _ _ =         error "kkLoopBody.isInnerCoal: not an error, but I would like to know why!"       innerHasStride1 lmad =         let lmad_dims = LMAD.dims lmad-            stride = IxFun.ldStride $ last lmad_dims+            stride = LMAD.ldStride $ last lmad_dims          in stride == pe64 (intConst Int64 1)       --       mkRedomapOneTileBody acc_merge asss bsss fits_ij = do

@@ -71,56 +71,20 @@ -- per iteration (and an initial one, elided above). module Futhark.Optimise.DoubleBuffer (doubleBufferGPU, doubleBufferMC) where -import Control.Monad import Control.Monad.Reader import Control.Monad.State-import Control.Monad.Writer import Data.Bifunctor-import Data.List (find)+import Data.List qualified as L import Data.Map.Strict qualified as M import Data.Maybe import Futhark.Construct import Futhark.IR.GPUMem as GPU import Futhark.IR.MCMem as MC-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Pass-import Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp) import Futhark.Pass.ExplicitAllocations.GPU () import Futhark.Transform.Substitute-import Futhark.Util (mapAccumLM, maybeHead)---- | The double buffering pass definition.-doubleBuffer :: (Mem rep inner) => String -> String -> OptimiseOp rep -> Pass rep rep-doubleBuffer name desc onOp =-  Pass-    { passName = name,-      passDescription = desc,-      passFunction = intraproceduralTransformation optimise-    }-  where-    optimise scope stms = modifyNameSource $ \src ->-      let m =-            runDoubleBufferM $ localScope scope $ optimiseStms $ stmsToList stms-       in runState (runReaderT m env) src--    env = Env mempty doNotTouchLoop onOp-    doNotTouchLoop pat merge body = pure (mempty, pat, merge, body)---- | The pass for GPU kernels.-doubleBufferGPU :: Pass GPUMem GPUMem-doubleBufferGPU =-  doubleBuffer-    "Double buffer GPU"-    "Double buffer memory in sequential loops (GPU rep)."-    optimiseGPUOp---- | The pass for multicore-doubleBufferMC :: Pass MCMem MCMem-doubleBufferMC =-  doubleBuffer-    "Double buffer MC"-    "Double buffer memory in sequential loops (MC rep)."-    optimiseMCOp+import Futhark.Util (mapAccumLM)  type OptimiseLoop rep =   Pat (LetDec rep) ->@@ -250,56 +214,107 @@     invariant Constant {} = True     invariant (Var v) = v `notNameIn` bound -optimiseLoop :: (Constraints rep inner) => OptimiseLoop rep-optimiseLoop pat merge body = do-  (outer_stms_1, pat', merge', body') <--    optimiseLoopBySwitching pat merge body-  (outer_stms_2, pat'', merge'', body'') <--    inScopeOf outer_stms_1 $ optimiseLoopByCopying pat' merge' body'-  pure (outer_stms_1 <> outer_stms_2, pat'', merge'', body'')- isArrayIn :: VName -> Param FParamMem -> Bool isArrayIn x (Param _ _ (MemArray _ _ _ (ArrayIn y _))) = x == y isArrayIn _ _ = False -optimiseLoopBySwitching :: (Constraints rep inner) => OptimiseLoop rep-optimiseLoopBySwitching (Pat pes) merge (Body _ body_stms body_res) = do+doubleBufferSpace :: Space -> Bool+doubleBufferSpace ScalarSpace {} = False+doubleBufferSpace _ = True++optimiseLoop :: (Constraints rep inner) => OptimiseLoop rep+optimiseLoop (Pat pes) merge body@(Body _ body_stms body_res) = do   ((pat', merge', body'), outer_stms) <- runBuilder $ do-    ((buffered, body_stms'), (pes', merge', body_res')) <--      second unzip3 <$> mapAccumLM check (mempty, body_stms) (zip3 pes merge body_res)-    merge'' <- mapM (maybeCopyInitial buffered) $ mconcat merge'-    pure (Pat $ mconcat pes', merge'', Body () body_stms' $ mconcat body_res')+    ((param_changes, body_stms'), (pes', merge', body_res')) <-+      second unzip3 <$> mapAccumLM check (id, body_stms) (zip3 pes merge body_res)+    pure+      ( Pat $ mconcat pes',+        map param_changes $ mconcat merge',+        Body () body_stms' $ mconcat body_res'+      )   pure (outer_stms, pat', merge', body')   where-    merge_bound = namesFromList $ map (paramName . fst) merge+    bound_in_loop =+      namesFromList (map (paramName . fst) merge) <> boundInBody body -    check (buffered, body_stms') (pe, (param, arg), res)+    findLmadOfArray v = listToMaybe . mapMaybe onStm $ stmsToList body_stms+      where+        onStm = listToMaybe . mapMaybe onPatElem . patElems . stmPat+        onPatElem (PatElem pe_v (MemArray _ _ _ (ArrayIn _ lmad)))+          | v == pe_v,+            not $ bound_in_loop `namesIntersect` freeIn lmad =+              Just lmad+        onPatElem _ = Nothing++    changeParam p_needle new (p, p_initial) =+      if p == p_needle then new else (p, p_initial)++    check (param_changes, body_stms') (pe, (param, arg), res)       | Mem space <- paramType param,+        doubleBufferSpace space,         Var arg_v <- arg,         -- XXX: what happens if there are multiple arrays in the same         -- memory block?-        [arr_param] <- filter (isArrayIn (paramName param)) $ map fst merge,-        MemArray pt _ _ (ArrayIn _ ixfun) <- paramDec arr_param,-        not $ merge_bound `namesIntersect` freeIn (IxFun.base ixfun),-        Var res_v <- resSubExp res,-        Just (res_v_alloc, body_stms'') <- extractAllocOf merge_bound res_v body_stms' = do+        [((arr_param, Var arr_param_initial), Var arr_v)] <-+          filter+            (isArrayIn (paramName param) . fst . fst)+            (zip merge $ map resSubExp body_res),+        MemArray pt shape _ (ArrayIn _ param_lmad) <- paramDec arr_param,+        Var arr_mem_out <- resSubExp res,+        Just arr_lmad <- findLmadOfArray arr_v,+        Just (arr_mem_out_alloc, body_stms'') <-+          extractAllocOf bound_in_loop arr_mem_out body_stms' = do+          -- Put the allocations outside the loop.           num_bytes <--            letSubExp "num_bytes" =<< toExp (product $ primByteSize pt : IxFun.base ixfun)+            letSubExp "num_bytes" =<< toExp (primByteSize pt * (1 + LMAD.range arr_lmad))           arr_mem_in <-             letExp (baseString arg_v <> "_in") $ Op $ Alloc num_bytes space+          addStm arr_mem_out_alloc++          -- Construct additional pattern element and parameter for+          -- the memory block that is not used afterwards.           pe_unused <-             PatElem               <$> newVName (baseString (patElemName pe) <> "_unused")               <*> pure (MemMem space)           param_out <-             newParam (baseString (paramName param) <> "_out") (MemMem space)-          addStm res_v_alloc++          -- Copy the initial array value to the input memory, with+          -- the same index function as the result.+          arr_v_copy <- newVName $ baseString arr_v <> "_db_copy"+          let arr_initial_info =+                MemArray pt shape NoUniqueness $ ArrayIn arr_mem_in arr_lmad+              arr_initial_pe =+                PatElem arr_v_copy arr_initial_info+          addStm . Let (Pat [arr_initial_pe]) (defAux ()) . BasicOp $+            Replicate mempty (Var arr_param_initial)+          -- AS a trick we must make the array parameter Unique to+          -- avoid unfortunate hoisting (see #1533) because we are+          -- invalidating the underlying memory.+          let arr_param' =+                Param mempty (paramName arr_param) $+                  MemArray pt shape Unique (ArrayIn (paramName param) param_lmad)++          -- We must also update the initial values of the parameters+          -- used in the index function of this array parameter, such+          -- that they match the result.+          let mkUpdate lmad_v =+                case L.find ((== lmad_v) . paramName . fst . fst) $+                  zip merge body_res of+                  Nothing -> id+                  Just ((p, _), p_res) -> changeParam p (p, resSubExp p_res)+              updateLmadParam =+                foldl (.) id $ map mkUpdate $ namesToList $ freeIn param_lmad+           pure-            ( ( M.insert (paramName param) arr_mem_in buffered,-                substituteNames (M.singleton res_v (paramName param_out)) body_stms''+            ( ( updateLmadParam+                  . changeParam arr_param (arr_param', Var arr_v_copy)+                  . param_changes,+                substituteNames (M.singleton arr_mem_out (paramName param_out)) body_stms''               ),               ( [pe, pe_unused],-                [(param, Var arr_mem_in), (param_out, resSubExp res)],+                [(param, Var arr_mem_in), (param_out, Var arr_mem_out)],                 [ res {resSubExp = Var $ paramName param_out},                   subExpRes $ Var $ paramName param                 ]@@ -307,186 +322,39 @@             )       | otherwise =           pure-            ( (buffered, body_stms'),+            ( (param_changes, body_stms'),               ([pe], [(param, arg)], [res])             ) -    maybeCopyInitial buffered (param@(Param _ _ (MemArray _ _ _ (ArrayIn mem _))), Var arg)-      | Just mem' <- mem `M.lookup` buffered = do-          arg_info <- lookupMemInfo arg-          case arg_info of-            MemArray pt shape u (ArrayIn _ arg_ixfun) -> do-              arg_copy <- newVName (baseString arg <> "_dbcopy")-              letBind (Pat [PatElem arg_copy $ MemArray pt shape u $ ArrayIn mem' arg_ixfun]) $-                BasicOp (Replicate mempty $ Var arg)-              -- We need to make this parameter unique to avoid invalid-              -- hoisting (see #1533), because we are invalidating the-              -- underlying memory.-              pure (fmap mkUnique param, Var arg_copy)-            _ -> pure (fmap mkUnique param, Var arg)-    maybeCopyInitial _ (param, arg) = pure (param, arg)--    mkUnique (MemArray bt shape _ ret) = MemArray bt shape Unique ret-    mkUnique x = x--optimiseLoopByCopying :: (Constraints rep inner) => OptimiseLoop rep-optimiseLoopByCopying pat merge body = do-  -- We start out by figuring out which of the merge variables should-  -- be double-buffered.-  buffered <--    doubleBufferLoopParams-      (zip (map fst merge) (bodyResult body))-      (boundInBody body)-  -- Then create the allocations of the buffers and copies of the-  -- initial values.-  (merge', allocs) <- allocStms merge buffered-  -- Modify the loop body to copy buffered result arrays.-  let body' = doubleBufferResult (map fst merge) buffered body-  pure (stmsFromList allocs, pat, merge', body')---- | The booleans indicate whether we should also play with the--- initial merge values.-data DoubleBuffer-  = BufferAlloc VName (PrimExp VName) Space Bool-  | -- | First name is the memory block to copy to,-    -- second is the name of the array copy.-    BufferCopy VName IxFun VName Bool-  | NoBuffer-  deriving (Show)--doubleBufferLoopParams ::-  (MonadFreshNames m) =>-  [(Param FParamMem, SubExpRes)] ->-  Names ->-  m [DoubleBuffer]-doubleBufferLoopParams ctx_and_res bound_in_loop =-  evalStateT (mapM buffer ctx_and_res) M.empty-  where-    params = map fst ctx_and_res-    loopVariant v =-      v-        `nameIn` bound_in_loop-        || v-          `elem` map (paramName . fst) ctx_and_res--    loopInvariantSize (Constant v) =-      Just (Constant v, True)-    loopInvariantSize (Var v) =-      case find ((== v) . paramName . fst) ctx_and_res of-        Just (_, SubExpRes _ (Constant val)) ->-          Just (Constant val, False)-        Just (_, SubExpRes _ (Var v'))-          | not $ loopVariant v' ->-              Just (Var v', False)-        Just _ ->-          Nothing-        Nothing ->-          Just (Var v, True)--    sizeForMem mem = maybeHead $ mapMaybe (arrayInMem . paramDec) params-      where-        arrayInMem (MemArray pt shape _ (ArrayIn arraymem ixfun))-          | IxFun.isDirect ixfun,-            Just (dims, b) <--              mapAndUnzipM loopInvariantSize $ shapeDims shape,-            mem == arraymem =-              Just-                ( arraySizeInBytesExp $-                    Array pt (Shape dims) NoUniqueness,-                  or b-                )-        arrayInMem _ = Nothing--    buffer (fparam, res) = case paramType fparam of-      Mem space-        | Just (size, b) <- sizeForMem $ paramName fparam,-          Var res_v <- resSubExp res,-          res_v `nameIn` bound_in_loop -> do-            -- Let us double buffer this!-            bufname <- lift $ newVName "double_buffer_mem"-            modify $ M.insert (paramName fparam) (bufname, b)-            pure $ BufferAlloc bufname size space b-      Array {}-        | MemArray _ _ _ (ArrayIn mem ixfun) <- paramDec fparam -> do-            buffered <- gets $ M.lookup mem-            case buffered of-              Just (bufname, b) -> do-                copyname <- lift $ newVName "double_buffer_array"-                pure $ BufferCopy bufname ixfun copyname b-              Nothing ->-                pure NoBuffer-      _ -> pure NoBuffer--allocStms ::-  (Constraints rep inner) =>-  [(FParam rep, SubExp)] ->-  [DoubleBuffer] ->-  DoubleBufferM rep ([(FParam rep, SubExp)], [Stm rep])-allocStms merge = runWriterT . zipWithM allocation merge-  where-    allocation m@(Param attrs pname _, _) (BufferAlloc name size space b) = do-      stms <- lift $-        runBuilder_ $ do-          size' <- toSubExp "double_buffer_size" size-          letBindNames [name] $ Op $ Alloc size' space-      tell $ stmsToList stms-      if b-        then pure (Param attrs pname $ MemMem space, Var name)-        else pure m-    allocation (f, Var v) (BufferCopy mem _ _ b) | b = do-      v_copy <- lift $ newVName $ baseString v ++ "_double_buffer_copy"-      (_v_mem, v_ixfun) <- lift $ lookupArraySummary v-      let bt = elemType $ paramType f-          shape = arrayShape $ paramType f-          bound = MemArray bt shape NoUniqueness $ ArrayIn mem v_ixfun-      tell-        [ Let (Pat [PatElem v_copy bound]) (defAux ()) $-            BasicOp (Replicate mempty $ Var v)-        ]-      -- It is important that we treat this as a consumption, to-      -- avoid the Copy from being hoisted out of any enclosing-      -- loops.  Since we re-use (=overwrite) memory in the loop,-      -- the copy is critical for initialisation.  See issue #816.-      let uniqueMemInfo (MemArray pt pshape _ ret) =-            MemArray pt pshape Unique ret-          uniqueMemInfo info = info-      pure (uniqueMemInfo <$> f, Var v_copy)-    allocation (f, se) _ =-      pure (f, se)--doubleBufferResult ::-  (Constraints rep inner) =>-  [FParam rep] ->-  [DoubleBuffer] ->-  Body rep ->-  Body rep-doubleBufferResult valparams buffered (Body _ stms res) =-  let (ctx_res, val_res) = splitAt (length res - length valparams) res-      (copystms, val_res') =-        unzip $ zipWith3 buffer valparams buffered val_res-   in Body () (stms <> stmsFromList (catMaybes copystms)) $ ctx_res ++ val_res'+-- | The double buffering pass definition.+doubleBuffer :: (Mem rep inner) => String -> String -> OptimiseOp rep -> Pass rep rep+doubleBuffer name desc onOp =+  Pass+    { passName = name,+      passDescription = desc,+      passFunction = intraproceduralTransformation optimise+    }   where-    buffer _ (BufferAlloc bufname _ _ _) se =-      (Nothing, se {resSubExp = Var bufname})-    buffer fparam (BufferCopy bufname ixfun copyname _) (SubExpRes cs (Var v)) =-      -- To construct the copy we will need to figure out its type-      -- based on the type of the function parameter.-      let t = resultType $ paramType fparam-          summary = MemArray (elemType t) (arrayShape t) NoUniqueness $ ArrayIn bufname ixfun-          copystm =-            Let-              (Pat [PatElem copyname summary])-              (defAux ())-              (BasicOp $ Replicate mempty $ Var v)-       in (Just copystm, SubExpRes cs (Var copyname))-    buffer _ _ se =-      (Nothing, se)+    optimise scope stms = modifyNameSource $ \src ->+      let m =+            runDoubleBufferM $ localScope scope $ optimiseStms $ stmsToList stms+       in runState (runReaderT m env) src -    parammap = M.fromList $ zip (map paramName valparams) $ map resSubExp res+    env = Env mempty doNotTouchLoop onOp+    doNotTouchLoop pat merge body = pure (mempty, pat, merge, body) -    resultType t = t `setArrayDims` map substitute (arrayDims t)+-- | The pass for GPU kernels.+doubleBufferGPU :: Pass GPUMem GPUMem+doubleBufferGPU =+  doubleBuffer+    "Double buffer GPU"+    "Double buffer memory in sequential loops (GPU rep)."+    optimiseGPUOp -    substitute (Var v)-      | Just replacement <- M.lookup v parammap = replacement-    substitute se =-      se+-- | The pass for multicore+doubleBufferMC :: Pass MCMem MCMem+doubleBufferMC =+  doubleBuffer+    "Double buffer MC"+    "Double buffer memory in sequential loops (MC rep)."+    optimiseMCOp

@@ -42,11 +42,11 @@   where     table = consts_table <> mkTable (bodyStms (funDefBody fd))     mkSubst (Var v0)-      | Just (MemArray _ _ _ (ArrayIn mem0 ixfun0), BasicOp (Manifest _ v1)) <-+      | Just (MemArray _ _ _ (ArrayIn mem0 lmad0), BasicOp (Manifest _ v1)) <-           varInfo v0 table,-        Just (MemArray _ _ _ (ArrayIn mem1 ixfun1), _) <-+        Just (MemArray _ _ _ (ArrayIn mem1 lmad1), _) <-           varInfo v1 table,-        ixfun0 == ixfun1 =+        lmad0 == lmad1 =           M.fromList [(mem0, mem1), (v0, v1)]     mkSubst _ = mempty     onBody (Body dec stms res) =

@@ -213,10 +213,8 @@ -- precise. removeDeadBranchResult :: (BuilderOps rep) => BottomUpRuleMatch rep removeDeadBranchResult (_, used) pat _ (cond, cases, defbody, MatchDec rettype ifsort)-  | -- Only if there is no existential binding...-    all (`notNameIn` foldMap freeIn (patElems pat)) (patNames pat),-    -- Figure out which of the names in 'pat' are used...-    patused <- map (`UT.isUsedDirectly` used) $ patNames pat,+  | -- Figure out which of the names in 'pat' are used...+    patused <- map keep $ patNames pat,     -- If they are not all used, then this rule applies.     not (and patused) = do       -- Remove the parts of the branch-results that correspond to dead@@ -226,12 +224,23 @@           pick = map snd . filter fst . zip patused           pat' = pick $ patElems pat           rettype' = pick rettype+          -- We also need to adjust the existential references in the+          -- branch type.+          exts = scanl (+) 0 [if b then 1 else 0 | b <- patused]+          adjust = mapExt (exts !!)       Simplify $ do         cases' <- mapM (traverse $ onBody pick) cases         defbody' <- onBody pick defbody-        letBind (Pat pat') $ Match cond cases' defbody' $ MatchDec rettype' ifsort+        letBind (Pat pat') $ Match cond cases' defbody' $ MatchDec (map adjust rettype') ifsort   | otherwise = Skip   where+    usedDirectly v = v `UT.isUsedDirectly` used+    usedIndirectly v =+      any+        (\pe -> v `nameIn` freeIn pe && usedDirectly (patElemName pe))+        (patElems pat)+    keep v = usedDirectly v || usedIndirectly v+     onBody pick (Body _ stms res) = mkBodyM stms $ pick res  topDownRules :: (BuilderOps rep) => [TopDownRule rep]

@@ -23,7 +23,7 @@ import Data.List (foldl', zip4) import Data.Map qualified as M import Futhark.IR.GPU-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.SeqMem qualified as ExpMem import Futhark.MonadFreshNames import Futhark.Tools@@ -256,13 +256,13 @@ ---- Helpers for building the environment that binds array variable names to their index functions ---------------- -type IxFun = IxFun.IxFun (TPrimExp Int64 VName)+type LMAD = LMAD.LMAD (TPrimExp Int64 VName)  -- | Map from array variable names to their corresponding index functions. --   The info is not guaranteed to be exact, e.g., we assume ifs and loops --   return arrays layed out in normalized (row-major) form in memory. --   We only record aliasing statements, such as transposition, slice, etc.-type IxFnEnv = M.Map VName IxFun+type IxFnEnv = M.Map VName LMAD  type WithEnv = M.Map VName (Lambda GPU, [SubExp]) @@ -288,7 +288,7 @@        in (lam_op, ne) changeWithEnv with_env _ = pure with_env -composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> Maybe IxFun) -> TileM IxFnEnv+composeIxfuns :: IxFnEnv -> VName -> VName -> (LMAD -> Maybe LMAD) -> TileM IxFnEnv composeIxfuns env y x ixf_fun =   case ixf_fun =<< M.lookup x env of     Just ixf -> pure $ M.insert y ixf env@@ -296,27 +296,27 @@       tp <- lookupType x       pure $ case tp of         Array _ptp shp _u-          | Just ixf <- ixf_fun $ IxFun.iota $ map ExpMem.pe64 (shapeDims shp) ->+          | Just ixf <- ixf_fun $ LMAD.iota 0 $ map ExpMem.pe64 (shapeDims shp) ->               M.insert y ixf env         _ -> env  changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv changeIxFnEnv env y (BasicOp (Reshape ReshapeArbitrary shp_chg x)) =-  composeIxfuns env y x (`IxFun.reshape` fmap ExpMem.pe64 (shapeDims shp_chg))+  composeIxfuns env y x (`LMAD.reshape` fmap ExpMem.pe64 (shapeDims shp_chg)) changeIxFnEnv env y (BasicOp (Reshape ReshapeCoerce shp_chg x)) =-  composeIxfuns env y x (Just . (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg)))+  composeIxfuns env y x (Just . (`LMAD.coerce` fmap ExpMem.pe64 (shapeDims shp_chg))) changeIxFnEnv env y (BasicOp (Manifest perm x)) = do   tp <- lookupType x   case tp of     Array _ptp shp _u -> do       let shp' = map ExpMem.pe64 (shapeDims shp)-      let ixfn = IxFun.permute (IxFun.iota shp') perm+      let ixfn = LMAD.permute (LMAD.iota 0 shp') perm       pure $ M.insert y ixfn env     _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!" changeIxFnEnv env y (BasicOp (Rearrange perm x)) =-  composeIxfuns env y x (Just . (`IxFun.permute` perm))+  composeIxfuns env y x (Just . (`LMAD.permute` perm)) changeIxFnEnv env y (BasicOp (Index x slc)) =-  composeIxfuns env y x (Just . (`IxFun.slice` Slice (map (fmap ExpMem.pe64) $ unSlice slc)))+  composeIxfuns env y x (Just . (`LMAD.slice` Slice (map (fmap ExpMem.pe64) $ unSlice slc))) changeIxFnEnv env y (BasicOp (Opaque _ (Var x))) =   composeIxfuns env y x Just changeIxFnEnv env _ _ = pure env

@@ -8,6 +8,7 @@ import Control.Monad.Reader import Control.Monad.State import Control.Monad.Writer+import Data.Bifunctor import Data.Either (rights) import Data.List (find, foldl') import Data.Map.Strict qualified as M@@ -18,7 +19,7 @@ import Futhark.IR import Futhark.IR.GPU.Simplify qualified as GPU import Futhark.IR.GPUMem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.MonadFreshNames import Futhark.Optimise.Simplify.Rep (addScopeWisdom) import Futhark.Pass@@ -178,11 +179,10 @@         genericExpandedInvariantAllocations (const (shape, map le64 is)) invariant_allocs        scope <- askScope-      let scope' = scopeOf op_lam <> scope-      either throwError pure $-        runOffsetM scope' alloc_offsets $ do-          op_lam'' <- offsetMemoryInLambda op_lam'-          pure (alloc_stms, (shape, arrs, Just (op_lam'', nes)))+      let scope' = scopeOf op_lam <> scope <> scopeOf alloc_stms+      either throwError pure <=< runOffsetM scope' $ do+        op_lam'' <- offsetMemoryInLambda alloc_offsets op_lam'+        pure (alloc_stms, (shape, arrs, Just (op_lam'', nes))) transformExp e =   pure (mempty, e) @@ -248,9 +248,9 @@     else do       (lvl_stms, lvl', grid) <- ensureGridKnown lvl       allocsForBody variant_allocs invariant_allocs grid space kbody kbody' $-        \alloc_stms kbody'' -> do+        \offsets alloc_stms kbody'' -> do           ops'' <- forM ops' $ \op' ->-            localScope (scopeOf op') $ offsetMemoryInLambda op'+            localScope (scopeOf op') $ offsetMemoryInLambda offsets op'           pure (lvl_stms <> alloc_stms, (lvl', ops'', kbody''))   where     bound_in_kernel =@@ -267,7 +267,7 @@   SegSpace ->   KernelBody GPUMem ->   KernelBody GPUMem ->-  (Stms GPUMem -> KernelBody GPUMem -> OffsetM b) ->+  (RebaseMap -> Stms GPUMem -> KernelBody GPUMem -> OffsetM b) ->   ExpandM b allocsForBody variant_allocs invariant_allocs grid space kbody kbody' m = do   (alloc_offsets, alloc_stms) <-@@ -279,11 +279,10 @@       invariant_allocs    scope <- askScope-  let scope' = scopeOfSegSpace space <> scope-  either throwError pure $-    runOffsetM scope' alloc_offsets $ do-      kbody'' <- offsetMemoryInKernelBody kbody'-      m alloc_stms kbody''+  let scope' = scopeOfSegSpace space <> scope <> scopeOf alloc_stms+  either throwError pure <=< runOffsetM scope' $ do+    kbody'' <- offsetMemoryInKernelBody alloc_offsets kbody'+    m alloc_offsets alloc_stms kbody''  memoryRequirements ::   KernelGrid ->@@ -551,160 +550,202 @@ -- | A map from memory block names to index function embeddings.. type RebaseMap = M.Map VName ([Exp64] -> Expansion) +--- Modifying the index functions of code.+ newtype OffsetM a-  = OffsetM-      ( ReaderT-          (Scope GPUMem)-          (ReaderT RebaseMap (Either String))-          a-      )+  = OffsetM (BuilderT GPUMem (StateT VNameSource (Either String)) a)   deriving     ( Applicative,       Functor,       Monad,       HasScope GPUMem,       LocalScope GPUMem,-      MonadError String+      MonadError String,+      MonadFreshNames     ) -runOffsetM :: Scope GPUMem -> RebaseMap -> OffsetM a -> Either String a-runOffsetM scope offsets (OffsetM m) =-  runReaderT (runReaderT m scope) offsets+instance MonadBuilder OffsetM where+  type Rep OffsetM = GPUMem+  mkExpDecM pat e = OffsetM $ mkExpDecM pat e+  mkBodyM stms res = OffsetM $ mkBodyM stms res+  mkLetNamesM pat e = OffsetM $ mkLetNamesM pat e -askRebaseMap :: OffsetM RebaseMap-askRebaseMap = OffsetM $ lift ask+  addStms = OffsetM . addStms+  collectStms (OffsetM m) = OffsetM $ collectStms m -localRebaseMap :: (RebaseMap -> RebaseMap) -> OffsetM a -> OffsetM a-localRebaseMap f (OffsetM m) = OffsetM $ do-  scope <- ask-  lift $ local f $ runReaderT m scope+runOffsetM ::+  (MonadFreshNames m) =>+  Scope GPUMem ->+  OffsetM a ->+  m (Either String a)+runOffsetM scope (OffsetM m) = modifyNameSource $ \src ->+  case runStateT (runBuilderT m scope) src of+    Left e -> (Left e, src)+    Right (x, src') -> (Right (fst x), src') -lookupNewBase :: VName -> [Exp64] -> OffsetM (Maybe Expansion)-lookupNewBase name x = do-  offsets <- askRebaseMap-  pure $ ($ x) <$> M.lookup name offsets+lookupNewBase :: VName -> [Exp64] -> RebaseMap -> Maybe Expansion+lookupNewBase name x offsets =+  ($ x) <$> M.lookup name offsets -offsetMemoryInKernelBody :: KernelBody GPUMem -> OffsetM (KernelBody GPUMem)-offsetMemoryInKernelBody kbody = do-  scope <- askScope+offsetMemoryInKernelBody :: RebaseMap -> KernelBody GPUMem -> OffsetM (KernelBody GPUMem)+offsetMemoryInKernelBody offsets kbody = do   stms' <--    stmsFromList . snd-      <$> mapAccumLM-        (\scope' -> localScope scope' . offsetMemoryInStm)-        scope-        (stmsToList $ kernelBodyStms kbody)+    collectStms_ $+      mapM_ (addStm <=< offsetMemoryInStm offsets) (kernelBodyStms kbody)   pure kbody {kernelBodyStms = stms'} -offsetMemoryInBody :: Body GPUMem -> OffsetM (Body GPUMem)-offsetMemoryInBody (Body dec stms res) = do-  scope <- askScope-  stms' <--    stmsFromList . snd-      <$> mapAccumLM-        (\scope' -> localScope scope' . offsetMemoryInStm)-        scope-        (stmsToList stms)-  pure $ Body dec stms' res+offsetMemoryInBody :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)+offsetMemoryInBody offsets (Body _ stms res) = do+  buildBody_ $ do+    mapM_ (addStm <=< offsetMemoryInStm offsets) stms+    pure res -offsetMemoryInStm :: Stm GPUMem -> OffsetM (Scope GPUMem, Stm GPUMem)-offsetMemoryInStm (Let pat dec e) = do-  e' <- offsetMemoryInExp e-  pat' <--    offsetMemoryInPat pat-      =<< maybe (throwError "offsetMemoryInStm: ill-typed") pure-      =<< expReturns e'-  scope <- askScope-  -- Try to recompute the index function.  Fall back to creating rebase-  -- operations with the RebaseMap.-  rts <--    maybe (throwError "offsetMemoryInStm: ill-typed") pure $-      runReader (expReturns e') scope-  let pat'' = Pat $ zipWith pick (patElems pat') rts-      stm = Let pat'' dec e'-  let scope' = scopeOf stm <> scope-  pure (scope', stm)+argsContext :: [SubExp] -> OffsetM [SubExp]+argsContext = fmap concat . mapM resCtx   where-    pick ::-      PatElem (MemInfo SubExp NoUniqueness MemBind) ->-      ExpReturns ->-      PatElem (MemInfo SubExp NoUniqueness MemBind)-    pick-      (PatElem name (MemArray pt s u _ret))-      (MemArray _ _ _ (Just (ReturnsInBlock m extixfun)))-        | Just ixfun <- instantiateIxFun extixfun =-            PatElem name (MemArray pt s u (ArrayIn m ixfun))-    pick p _ = p+    resCtx se = do+      v_t <- subExpMemInfo se+      case v_t of+        MemArray _ _ _ (ArrayIn mem lmad) -> do+          ctxs <- mapM (letSubExp "ctx" <=< toExp) (LMAD.existentialized lmad)+          pure $ Var mem : ctxs+        _ -> pure [] -    instantiateIxFun :: ExtIxFun -> Maybe IxFun-    instantiateIxFun = traverse (traverse inst)-      where-        inst Ext {} = Nothing-        inst (Free x) = pure x+offsetMemoryInBodyReturnCtx :: RebaseMap -> Body GPUMem -> OffsetM (Body GPUMem)+offsetMemoryInBodyReturnCtx offsets (Body _ stms res) = do+  buildBody_ $ do+    mapM_ (addStm <=< offsetMemoryInStm offsets) stms+    ctx <- argsContext $ map resSubExp res+    pure $ res <> subExpsRes ctx -offsetMemoryInPat :: Pat LetDecMem -> [ExpReturns] -> OffsetM (Pat LetDecMem)-offsetMemoryInPat (Pat pes) rets = do-  Pat <$> zipWithM onPE pes rets+lmadFrom :: LMAD.Shape num -> [num] -> LMAD.LMAD num+lmadFrom shape xs =+  LMAD.LMAD (head xs) $ zipWith LMAD.LMADDim (drop 1 xs) shape++-- | Append pattern elements corresponding to memory and index+-- function components for every array bound in the pattern.+addPatternContext :: Pat LetDecMem -> OffsetM (Pat LetDecMem)+addPatternContext (Pat pes) = localScope (scopeOfPat (Pat pes)) $ do+  (pes_ctx, pes') <- mapAccumLM onType [] pes+  pure $ Pat $ pes' <> pes_ctx   where+    onType+      acc+      (PatElem pe_v (MemArray pt pe_shape pe_u (ArrayIn pe_mem lmad))) = do+        space <- lookupMemSpace pe_mem+        pe_mem' <- newVName $ baseString pe_mem <> "_ext"+        let num_exts = length (LMAD.existentialized lmad)+        lmad_exts <-+          replicateM num_exts $+            PatElem <$> newVName "ext" <*> pure (MemPrim int64)+        let pe_lmad' = lmadFrom (LMAD.shape lmad) $ map (le64 . patElemName) lmad_exts+        pure+          ( acc ++ PatElem pe_mem' (MemMem space) : lmad_exts,+            PatElem pe_v $ MemArray pt pe_shape pe_u $ ArrayIn pe_mem' pe_lmad'+          )+    onType acc t = pure (acc, t)++-- | Append pattern elements corresponding to memory and index+-- function components for every array bound in the parameters.+addParamsContext :: [Param FParamMem] -> OffsetM [Param FParamMem]+addParamsContext ps = localScope (scopeOfFParams ps) $ do+  (ps_ctx, ps') <- mapAccumLM onType [] ps+  pure $ ps' <> ps_ctx+  where+    onType acc (Param attr v (MemArray pt shape u (ArrayIn mem lmad))) = do+      space <- lookupMemSpace mem+      mem' <- newVName $ baseString mem <> "_ext"+      let num_exts = length (LMAD.existentialized lmad)+      lmad_exts <-+        replicateM num_exts $+          Param mempty <$> newVName "ext" <*> pure (MemPrim int64)+      let lmad' = lmadFrom (LMAD.shape lmad) $ map (le64 . paramName) lmad_exts+      pure+        ( acc ++ Param mempty mem' (MemMem space) : lmad_exts,+          Param attr v $ MemArray pt shape u $ ArrayIn mem' lmad'+        )+    onType acc t = pure (acc, t)++offsetBranch ::+  Pat LetDecMem ->+  [BranchTypeMem] ->+  OffsetM (Pat LetDecMem, [BranchTypeMem])+offsetBranch (Pat pes) ts = do+  ((pes_ctx, ts_ctx), (pes', ts')) <-+    bimap unzip unzip <$> mapAccumLM onType [] (zip pes ts)+  pure (Pat $ pes' <> pes_ctx, ts' <> ts_ctx)+  where+    onType+      acc+      ( PatElem pe_v (MemArray _ pe_shape pe_u (ArrayIn pe_mem pe_lmad)),+        MemArray pt shape u meminfo+        ) = do+        (space, lmad) <- case meminfo of+          ReturnsInBlock mem lmad -> do+            space <- lookupMemSpace mem+            pure (space, lmad)+          ReturnsNewBlock space _ lmad ->+            pure (space, lmad)+        pe_mem' <- newVName $ baseString pe_mem <> "_ext"+        let start = length ts + length acc+            num_exts = length (LMAD.existentialized lmad)+            ext (Free se) = Free <$> pe64 se+            ext (Ext i) = le64 (Ext i)+        lmad_exts <-+          replicateM num_exts $+            PatElem <$> newVName "ext" <*> pure (MemPrim int64)+        let pe_lmad' = lmadFrom (LMAD.shape pe_lmad) $ map (le64 . patElemName) lmad_exts+        pure+          ( acc+              ++ (PatElem pe_mem' $ MemMem space, MemMem space)+              : map (,MemPrim int64) lmad_exts,+            ( PatElem pe_v $ MemArray pt pe_shape pe_u $ ArrayIn pe_mem' pe_lmad',+              MemArray pt shape u . ReturnsNewBlock space start . fmap ext $+                LMAD.mkExistential (shapeDims shape) (1 + start)+            )+          )+    onType acc t = pure (acc, t)++offsetMemoryInPat :: RebaseMap -> Pat LetDecMem -> [ExpReturns] -> Pat LetDecMem+offsetMemoryInPat offsets (Pat pes) rets = do+  Pat $ zipWith onPE pes rets+  where     onPE       (PatElem name (MemArray pt shape u (ArrayIn mem _)))       (MemArray _ _ _ info)-        | Just ixfun <- getIxFun info =-            pure . PatElem name . MemArray pt shape u . ArrayIn mem $-              fmap (fmap unExt) ixfun-    onPE pe _ = do-      new_dec <- offsetMemoryInMemBound (patElemName pe) $ patElemDec pe-      pure pe {patElemDec = new_dec}+        | Just lmad <- getLMAD info =+            PatElem name . MemArray pt shape u . ArrayIn mem $+              fmap (fmap unExt) lmad+    onPE pe _ =+      offsetMemoryInMemBound offsets <$> pe     unExt (Ext i) = patElemName (pes !! i)     unExt (Free v) = v-    getIxFun (Just (ReturnsNewBlock _ _ ixfun)) = Just ixfun-    getIxFun (Just (ReturnsInBlock _ ixfun)) = Just ixfun-    getIxFun _ = Nothing+    getLMAD (Just (ReturnsNewBlock _ _ lmad)) = Just lmad+    getLMAD (Just (ReturnsInBlock _ lmad)) = Just lmad+    getLMAD _ = Nothing -offsetMemoryInParam :: Param (MemBound u) -> OffsetM (Param (MemBound u))-offsetMemoryInParam fparam = do-  fparam' <- offsetMemoryInMemBound (paramName fparam) $ paramDec fparam-  pure fparam {paramDec = fparam'}+offsetMemoryInParam :: RebaseMap -> Param (MemBound u) -> Param (MemBound u)+offsetMemoryInParam offsets = fmap $ offsetMemoryInMemBound offsets -offsetMemoryInMemBound :: VName -> MemBound u -> OffsetM (MemBound u)-offsetMemoryInMemBound v summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do-  embedding <- lookupNewBase mem $ IxFun.shape ixfun-  case embedding of-    Nothing -> pure summary-    Just (o, p) -> do-      let problem =-            throwError . unlines $-              [ "offsetMemoryInMemBound",-                prettyString v,-                prettyString (o, p),-                prettyString ixfun-              ]-      ixfun' <- maybe problem pure $ IxFun.expand o p ixfun-      pure $ MemArray pt shape u $ ArrayIn mem ixfun'-offsetMemoryInMemBound _ summary = pure summary+offsetMemoryInMemBound :: RebaseMap -> MemBound u -> MemBound u+offsetMemoryInMemBound offsets (MemArray pt shape u (ArrayIn mem lmad))+  | Just (o, p) <- lookupNewBase mem (LMAD.shape lmad) offsets =+      MemArray pt shape u $ ArrayIn mem $ LMAD.expand o p lmad+offsetMemoryInMemBound _ info = info -offsetMemoryInBodyReturns :: BodyReturns -> OffsetM BodyReturns-offsetMemoryInBodyReturns br@(MemArray pt shape u (ReturnsInBlock mem ixfun))-  | Just ixfun' <- isStaticIxFun ixfun = do-      embedding <- lookupNewBase mem $ IxFun.shape ixfun'-      case embedding of-        Nothing -> pure br-        Just (o, p) -> do-          let problem =-                throwError . unlines $-                  [ "offsetMemoryInBodyReturns",-                    prettyString (o, p),-                    prettyString ixfun-                  ]-          ixfun'' <--            maybe problem pure $-              IxFun.expand (Free <$> o) (fmap Free p) ixfun-          pure $ MemArray pt shape u $ ReturnsInBlock mem ixfun''-offsetMemoryInBodyReturns br = pure br+offsetMemoryInBodyReturns :: RebaseMap -> BodyReturns -> BodyReturns+offsetMemoryInBodyReturns offsets (MemArray pt shape u (ReturnsInBlock mem lmad))+  | Just lmad' <- isStaticLMAD lmad,+    Just (o, p) <- lookupNewBase mem (LMAD.shape lmad') offsets =+      MemArray pt shape u $+        ReturnsInBlock mem $+          LMAD.expand (Free <$> o) (fmap Free p) lmad+offsetMemoryInBodyReturns _ br = br -offsetMemoryInLambda :: Lambda GPUMem -> OffsetM (Lambda GPUMem)-offsetMemoryInLambda lam = do-  body <- inScopeOf lam $ offsetMemoryInBody $ lambdaBody lam-  params <- mapM offsetMemoryInParam $ lambdaParams lam+offsetMemoryInLambda :: RebaseMap -> Lambda GPUMem -> OffsetM (Lambda GPUMem)+offsetMemoryInLambda offsets lam = do+  body <- inScopeOf lam $ offsetMemoryInBody offsets $ lambdaBody lam+  let params = map (offsetMemoryInParam offsets) $ lambdaParams lam   pure $ lam {lambdaBody = body, lambdaParams = params}  -- A loop may have memory parameters, and those memory blocks may@@ -712,35 +753,32 @@ -- initial value of a loop parameter is an expanded memory block, -- then so will the result be. offsetMemoryInLoopParams ::+  RebaseMap ->   [(FParam GPUMem, SubExp)] ->-  ([(FParam GPUMem, SubExp)] -> OffsetM a) ->+  (RebaseMap -> [(FParam GPUMem, SubExp)] -> OffsetM a) ->   OffsetM a-offsetMemoryInLoopParams merge f = do+offsetMemoryInLoopParams offsets merge f = do   let (params, args) = unzip merge-  localRebaseMap extend $ do-    params' <- mapM offsetMemoryInParam params-    f $ zip params' args+  params' <- addParamsContext params+  args' <- (args <>) <$> argsContext args+  f offsets' $ zip params' args'   where+    offsets' = extend offsets     extend rm = foldl' onParamArg rm merge     onParamArg rm (param, Var arg)       | Just x <- M.lookup arg rm =           M.insert (paramName param) x rm     onParamArg rm _ = rm -offsetMemoryInExp :: Exp GPUMem -> OffsetM (Exp GPUMem)-offsetMemoryInExp (Loop merge form body) = do-  offsetMemoryInLoopParams merge $ \merge' -> do-    body' <--      localScope-        (scopeOfFParams (map fst merge') <> scopeOfLoopForm form)-        (offsetMemoryInBody body)-    pure $ Loop merge' form body'-offsetMemoryInExp e = mapExpM recurse e+-- | Handles only the expressions where we do not change the number of+-- results; meaning anything except Loop, Match, and nonscalar Apply.+offsetMemoryInExp :: RebaseMap -> Exp GPUMem -> OffsetM (Exp GPUMem)+offsetMemoryInExp offsets = mapExpM recurse   where     recurse =       (identityMapper @GPUMem)-        { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody,-          mapOnBranchType = offsetMemoryInBodyReturns,+        { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody offsets,+          mapOnBranchType = pure . offsetMemoryInBodyReturns offsets,           mapOnOp = onOp         }     onOp (Inner (SegOp op)) =@@ -749,10 +787,56 @@       where         segOpMapper =           identitySegOpMapper-            { mapOnSegOpBody = offsetMemoryInKernelBody,-              mapOnSegOpLambda = offsetMemoryInLambda+            { mapOnSegOpBody = offsetMemoryInKernelBody offsets,+              mapOnSegOpLambda = offsetMemoryInLambda offsets             }     onOp op = pure op++offsetMemoryInStm :: RebaseMap -> Stm GPUMem -> OffsetM (Stm GPUMem)+offsetMemoryInStm offsets (Let pat dec (Match cond cases defbody (MatchDec ts kind))) = do+  cases' <- forM cases $ \(Case vs body) ->+    Case vs <$> offsetMemoryInBodyReturnCtx offsets body+  defbody' <- offsetMemoryInBodyReturnCtx offsets defbody+  (pat', ts') <- offsetBranch pat ts+  pure $ Let pat' dec $ Match cond cases' defbody' $ MatchDec ts' kind+offsetMemoryInStm offsets (Let pat dec (Loop merge form body)) = do+  loop' <-+    offsetMemoryInLoopParams offsets merge $ \offsets' merge' -> do+      body' <-+        localScope+          (scopeOfFParams (map fst merge') <> scopeOfLoopForm form)+          (offsetMemoryInBodyReturnCtx offsets' body)+      pure $ Loop merge' form body'+  pat' <- addPatternContext pat+  pure $ Let pat' dec loop'+offsetMemoryInStm offsets (Let pat dec e) = do+  e' <- offsetMemoryInExp offsets e+  pat' <-+    offsetMemoryInPat offsets pat+      <$> ( maybe (throwError "offsetMemoryInStm: ill-typed") pure+              =<< expReturns e'+          )+  scope <- askScope+  -- Try to recompute the index function.  Fall back to creating rebase+  -- operations with the RebaseMap.+  rts <-+    maybe (throwError "offsetMemoryInStm: ill-typed") pure $+      runReader (expReturns e') scope+  let pat'' = Pat $ zipWith pick (patElems pat') rts+  pure $ Let pat'' dec e'+  where+    pick+      (PatElem name (MemArray pt s u _ret))+      (MemArray _ _ _ (Just (ReturnsInBlock m extlmad)))+        | Just lmad <- instantiateLMAD extlmad =+            PatElem name (MemArray pt s u (ArrayIn m lmad))+    pick p _ = p++    instantiateLMAD :: ExtLMAD -> Maybe LMAD+    instantiateLMAD = traverse (traverse inst)+      where+        inst Ext {} = Nothing+        inst (Free x) = pure x  ---- Slicing allocation sizes out of a kernel. 

@@ -18,7 +18,6 @@     allocInStms,     allocForArray,     simplifiable,-    arraySizeInBytesExp,     mkLetNamesB',     mkLetNamesB'', @@ -48,7 +47,7 @@ import Futhark.Analysis.SymbolTable (IndexOp) import Futhark.Analysis.UsageTable qualified as UT import Futhark.IR.Mem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.IR.Prop.Aliases (AliasedOp) import Futhark.MonadFreshNames import Futhark.Optimise.Simplify.Engine (SimpleOps (..))@@ -75,10 +74,7 @@   )  data AllocEnv fromrep torep = AllocEnv-  { -- | Aggressively try to reuse memory in do-loops --    -- should be True inside kernels, False outside.-    aggressiveReuse :: Bool,-    -- | When allocating memory, put it in this memory space.+  { -- | When allocating memory, put it in this memory space.     -- This is primarily used to ensure that group-wide     -- statements store their results in local memory.     allocSpace :: Space,@@ -140,8 +136,7 @@   where     env =       AllocEnv-        { aggressiveReuse = False,-          allocSpace = space,+        { allocSpace = space,           envConsts = mempty,           allocInOp = handleOp,           envExpHints = hints@@ -154,17 +149,8 @@ arraySizeInBytesExp t =   untyped $ foldl' (*) (elemSize t) $ map pe64 (arrayDims t) -arraySizeInBytesExpM :: (MonadBuilder m) => Type -> m (PrimExp VName)-arraySizeInBytesExpM t = do-  let dim_prod_i64 = product $ map pe64 (arrayDims t)-      elm_size_i64 = elemSize t-  pure $-    BinOpExp (SMax Int64) (ValueExp $ IntValue $ Int64Value 0) $-      untyped $-        dim_prod_i64 * elm_size_i64- arraySizeInBytes :: (MonadBuilder m) => Type -> m SubExp-arraySizeInBytes = letSubExp "bytes" <=< toExp <=< arraySizeInBytesExpM+arraySizeInBytes = letSubExp "bytes" <=< toExp . arraySizeInBytesExp  allocForArray' ::   (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m)) =>@@ -270,15 +256,15 @@         pure $ PatElem (identName ident) summary       MemMem space ->         pure $ PatElem (identName ident) $ MemMem space-      MemArray bt _ u (Just (ReturnsInBlock mem extixfun)) -> do-        let ixfn = instantiateExtIxFun idents extixfun+      MemArray bt _ u (Just (ReturnsInBlock mem extlmad)) -> do+        let ixfn = instantiateExtLMAD idents extlmad         pure . PatElem (identName ident) . MemArray bt ident_shape u $ ArrayIn mem ixfn       MemArray _ extshape _ Nothing         | Just _ <- knownShape extshape -> do             summary <- summaryForBindage def_space (identType ident) hint             pure $ PatElem (identName ident) summary       MemArray bt _ u (Just (ReturnsNewBlock _ i extixfn)) -> do-        let ixfn = instantiateExtIxFun idents extixfn+        let ixfn = instantiateExtLMAD idents extixfn         pure . PatElem (identName ident) . MemArray bt ident_shape u $           ArrayIn (getIdent idents i) ixfn       MemAcc acc ispace ts u ->@@ -295,15 +281,15 @@         Nothing ->           error $ "getIdent: Ext " <> show i <> " but pattern has " <> show (length idents) <> " elements: " <> prettyString idents -    instantiateExtIxFun idents = fmap $ fmap inst+    instantiateExtLMAD idents = fmap $ fmap inst       where         inst (Free v) = v         inst (Ext i) = getIdent idents i -instantiateIxFun :: (Monad m) => ExtIxFun -> m IxFun-instantiateIxFun = traverse $ traverse inst+instantiateLMAD :: (Monad m) => ExtLMAD -> m LMAD+instantiateLMAD = traverse $ traverse inst   where-    inst Ext {} = error "instantiateIxFun: not yet"+    inst Ext {} = error "instantiateLMAD: not yet"     inst (Free x) = pure x  summaryForBindage ::@@ -320,16 +306,13 @@   pure $ MemAcc acc ispace ts u summaryForBindage def_space t@(Array pt shape u) NoHint = do   m <- allocForArray' t def_space-  pure $ MemArray pt shape u $ ArrayIn m $ IxFun.iota $ map pe64 $ arrayDims t-summaryForBindage _ t@(Array pt _ _) (Hint ixfun space) = do+  pure $ MemArray pt shape u $ ArrayIn m $ LMAD.iota 0 $ map pe64 $ arrayDims t+summaryForBindage _ t@(Array pt _ _) (Hint lmad space) = do   bytes <-     letSubExp "bytes" <=< toExp . untyped $-      product-        [ product $ IxFun.base ixfun,-          fromIntegral (primByteSize pt :: Int64)-        ]+      primByteSize pt * (1 + LMAD.range lmad)   m <- letExp "mem" $ Op $ Alloc bytes space-  pure $ MemArray pt (arrayShape t) NoUniqueness $ ArrayIn m ixfun+  pure $ MemArray pt (arrayShape t) NoUniqueness $ ArrayIn m lmad  allocInFParams ::   (Allocable fromrep torep inner) =>@@ -355,10 +338,10 @@   case paramDeclType param of     Array pt shape u -> do       let memname = baseString (paramName param) <> "_mem"-          ixfun = IxFun.iota $ map pe64 $ shapeDims shape+          lmad = LMAD.iota 0 $ map pe64 $ shapeDims shape       mem <- lift $ newVName memname       tell ([Param (paramAttrs param) mem $ MemMem pspace], [])-      pure param {paramDec = MemArray pt shape u $ ArrayIn mem ixfun}+      pure param {paramDec = MemArray pt shape u $ ArrayIn mem lmad}     Prim pt ->       pure param {paramDec = MemPrim pt}     Mem space ->@@ -372,12 +355,11 @@   VName ->   AllocM fromrep torep (VName, VName) ensureRowMajorArray space_ok v = do-  (mem, ixfun) <- lookupArraySummary v+  (mem, _) <- lookupArraySummary v   mem_space <- lookupMemSpace mem   default_space <- askDefaultSpace   let space = fromMaybe default_space space_ok-  if length (IxFun.base ixfun) == IxFun.rank ixfun-    && maybe True (== mem_space) space_ok+  if maybe True (== mem_space) space_ok     then pure (mem, v)     else allocLinearArray space (baseString v) v @@ -390,8 +372,8 @@   error $ "ensureArrayIn: " ++ prettyString v ++ " cannot be an array." ensureArrayIn space (Var v) = do   (mem', v') <- lift $ ensureRowMajorArray (Just space) v-  (_, ixfun) <- lift $ lookupArraySummary v'-  ctx <- lift $ mapM (letSubExp "ixfun_arg" <=< toExp) (IxFun.existentialized ixfun)+  (_, lmad) <- lift $ lookupArraySummary v'+  ctx <- lift $ mapM (letSubExp "lmad_arg" <=< toExp) (LMAD.existentialized lmad)   tell ([Var mem'], ctx)   pure $ Var v' @@ -422,15 +404,15 @@     param_names = namesFromList $ map (paramName . fst) merge     anyIsLoopParam names = names `namesIntersect` param_names -    scalarRes param_t v_mem_space v_ixfun (Var res) = do+    scalarRes param_t v_mem_space v_lmad (Var res) = do       -- Try really hard to avoid copying needlessly, but the result       -- _must_ be in ScalarSpace and have the right index function.-      (res_mem, res_ixfun) <- lift $ lookupArraySummary res+      (res_mem, res_lmad) <- lift $ lookupArraySummary res       res_mem_space <- lift $ lookupMemSpace res_mem       (res_mem', res') <--        if (res_mem_space, res_ixfun) == (v_mem_space, v_ixfun)+        if (res_mem_space, res_lmad) == (v_mem_space, v_lmad)           then pure (res_mem, res)-          else lift $ arrayWithIxFun v_mem_space v_ixfun (fromDecl param_t) res+          else lift $ arrayWithLMAD v_mem_space v_lmad (fromDecl param_t) res       tell ([Var res_mem'], [])       pure $ Var res'     scalarRes _ _ _ se = pure se@@ -447,7 +429,7 @@         )     allocInMergeParam (mergeparam, Var v)       | param_t@(Array pt shape u) <- paramDeclType mergeparam = do-          (v_mem, v_ixfun) <- lift $ lookupArraySummary v+          (v_mem, v_lmad) <- lift $ lookupArraySummary v           v_mem_space <- lift $ lookupMemSpace v_mem            -- Loop-invariant array parameters that are in scalar space@@ -467,33 +449,33 @@                   tell ([p], [])                    pure-                    ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName p) v_ixfun},+                    ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName p) v_lmad},                       Var v,-                      scalarRes param_t v_mem_space v_ixfun+                      scalarRes param_t v_mem_space v_lmad                     )             _ -> do               (v_mem', v') <- lift $ ensureRowMajorArray Nothing v-              let ixfun_ext =-                    IxFun.existentialize 0 $ IxFun.iota $ map pe64 $ shapeDims shape+              let lmad_ext =+                    LMAD.existentialize 0 $ LMAD.iota 0 $ map pe64 $ shapeDims shape                v_mem_space' <- lift $ lookupMemSpace v_mem'                ctx_params <--                replicateM (length (IxFun.existentialized ixfun_ext)) $+                replicateM (length (LMAD.existentialized lmad_ext)) $                   newParam "ctx_param_ext" (MemPrim int64) -              param_ixfun <--                instantiateIxFun $-                  IxFun.substituteInIxFun+              param_lmad <-+                instantiateLMAD $+                  LMAD.substitute                     ( M.fromList . zip (fmap Ext [0 ..]) $                         map (le64 . Free . paramName) ctx_params                     )-                    ixfun_ext+                    lmad_ext                mem_param <- newParam "mem_param" $ MemMem v_mem_space'               tell ([mem_param], ctx_params)               pure-                ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName mem_param) param_ixfun},+                ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName mem_param) param_lmad},                   Var v',                   ensureArrayIn v_mem_space'                 )@@ -503,18 +485,18 @@       mergeparam' <- allocInFParam mergeparam space       pure (mergeparam', se, linearFuncallArg (paramType mergeparam) space) -arrayWithIxFun ::+arrayWithLMAD ::   (MonadBuilder m, Op (Rep m) ~ MemOp inner (Rep m), LetDec (Rep m) ~ LetDecMem) =>   Space ->-  IxFun ->+  LMAD ->   Type ->   VName ->   m (VName, VName)-arrayWithIxFun space ixfun v_t v = do+arrayWithLMAD space lmad v_t v = do   let Array pt shape u = v_t   mem <- allocForArray' v_t space   v_copy <- newVName $ baseString v <> "_scalcopy"-  let pe = PatElem v_copy $ MemArray pt shape u $ ArrayIn mem ixfun+  let pe = PatElem v_copy $ MemArray pt shape u $ ArrayIn mem lmad   letBind (Pat [pe]) $ BasicOp $ Replicate mempty $ Var v   pure (mem, v_copy) @@ -524,10 +506,10 @@   VName ->   AllocM fromrep torep (VName, VName) ensureDirectArray space_ok v = do-  (mem, ixfun) <- lookupArraySummary v+  (mem, lmad) <- lookupArraySummary v   mem_space <- lookupMemSpace mem   default_space <- askDefaultSpace-  if IxFun.isDirect ixfun && maybe True (== mem_space) space_ok+  if LMAD.isDirect lmad && maybe True (== mem_space) space_ok     then pure (mem, v)     else needCopy (fromMaybe default_space space_ok)   where@@ -551,7 +533,7 @@       v' <- newVName $ s <> "_desired_form"       let info =             MemArray pt shape u . ArrayIn mem $-              IxFun.permute (IxFun.iota $ map pe64 $ arrayDims t) perm+              LMAD.permute (LMAD.iota 0 $ map pe64 $ arrayDims t) perm           pat = Pat [PatElem v' info]       addStm $ Let pat (defAux ()) $ BasicOp $ Manifest perm v       pure (mem, v')@@ -565,11 +547,10 @@   VName ->   AllocM fromrep torep (VName, VName) ensurePermArray space_ok perm v = do-  (mem, ixfun) <- lookupArraySummary v+  (mem, _) <- lookupArraySummary v   mem_space <- lookupMemSpace mem   default_space <- askDefaultSpace-  if length (IxFun.base ixfun) == length (IxFun.shape ixfun)-    && maybe True (== mem_space) space_ok+  if maybe True (== mem_space) space_ok     then pure (mem, v)     else allocPermArray (fromMaybe default_space space_ok) perm (baseString v) v @@ -668,9 +649,7 @@       i <- get <* modify (+ 1)       let shape' = fmap shift shape       pure . MemArray pt shape' u . ReturnsNewBlock space i $-        IxFun.iota $-          map convert $-            shapeDims shape'+        LMAD.iota 0 (map convert $ shapeDims shape')     addMem (Acc acc ispace ts u) = pure $ MemAcc acc ispace ts u      convert (Ext i) = le64 $ Ext i@@ -757,14 +736,14 @@   mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body  data MemReq-  = MemReq Space Rank+  = MemReq Space   | NeedsNormalisation Space   deriving (Eq, Show)  combMemReqs :: MemReq -> MemReq -> MemReq combMemReqs x@NeedsNormalisation {} _ = x combMemReqs _ y@NeedsNormalisation {} = y-combMemReqs x@(MemReq x_space _) y@MemReq {} =+combMemReqs x@(MemReq x_space) y@MemReq {} =   if x == y then x else NeedsNormalisation x_space  type MemReqType = MemInfo (Ext SubExp) NoUniqueness MemReq@@ -775,20 +754,17 @@ combMemReqTypes x _ = x  contextRets :: MemReqType -> [MemInfo d u r]-contextRets (MemArray _ shape _ (MemReq space (Rank base_rank))) =-  -- Memory + offset + base_rank + stride*rank.-  MemMem space-    : MemPrim int64-    : replicate base_rank (MemPrim int64)+contextRets (MemArray _ shape _ (MemReq space)) =+  -- Memory + offset + stride*rank.+  [MemMem space, MemPrim int64]     ++ replicate (shapeRank shape) (MemPrim int64) contextRets (MemArray _ shape _ (NeedsNormalisation space)) =-  -- Memory + offset + (base,stride)*rank.-  MemMem space-    : MemPrim int64-    : replicate (2 * shapeRank shape) (MemPrim int64)+  -- Memory + offset + stride*rank.+  [MemMem space, MemPrim int64]+    ++ replicate (shapeRank shape) (MemPrim int64) contextRets _ = [] --- Add memory information to the body, but do not return memory/ixfun+-- Add memory information to the body, but do not return memory/lmad -- information.  Instead, return restrictions on what the index -- function should look like.  We will then (crudely) unify these -- restrictions across all bodies.@@ -805,10 +781,9 @@     restriction t se = do       v_info <- subExpMemInfo se       case (t, v_info) of-        (Array pt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) -> do+        (Array pt shape u, MemArray _ _ _ (ArrayIn mem _)) -> do           space <- lookupMemSpace mem-          pure . MemArray pt shape u $-            MemReq space (Rank $ length $ IxFun.base ixfun)+          pure $ MemArray pt shape u $ MemReq space         (_, MemMem space) -> pure $ MemMem space         (_, MemPrim pt) -> pure $ MemPrim pt         (_, MemAcc acc ispace ts u) -> pure $ MemAcc acc ispace ts u@@ -829,17 +804,17 @@         res_rets_acc ++ [inspect ctx_offset req]       ) -    arrayInfo rank (NeedsNormalisation space) =-      (space, rank)-    arrayInfo _ (MemReq space (Rank base_rank)) =-      (space, base_rank)+    arrayInfo (NeedsNormalisation space) =+      space+    arrayInfo (MemReq space) =+      space      inspect ctx_offset (MemArray pt shape u req) =       let shape' = fmap (adjustExt num_new_ctx) shape-          (space, base_rank) = arrayInfo (shapeRank shape) req+          space = arrayInfo req        in MemArray pt shape' u . ReturnsNewBlock space ctx_offset $             convert-              <$> IxFun.mkExistential base_rank (shapeDims shape') (ctx_offset + 1)+              <$> LMAD.mkExistential (shapeDims shape') (ctx_offset + 1)     inspect _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u     inspect _ (MemPrim pt) = MemPrim pt     inspect _ (MemMem space) = MemMem space@@ -875,9 +850,9 @@         MemPrim {} -> pure []         MemAcc {} -> pure []         MemMem {} -> pure [] -- should not happen-        MemArray _ _ _ (ArrayIn mem ixfun) -> do-          ixfun_exts <- mapM (letSubExp "ixfun_ext" <=< toExp) $ IxFun.existentialized ixfun-          pure $ subExpRes (Var mem) : subExpsRes ixfun_exts+        MemArray _ _ _ (ArrayIn mem lmad) -> do+          lmad_exts <- mapM (letSubExp "lmad_ext" <=< toExp) $ LMAD.existentialized lmad+          pure $ subExpRes (Var mem) : subExpsRes lmad_exts  -- Do a a simple form of invariance analysis to simplify a Match.  It -- is unfortunate that we have to do it here, but functions such as@@ -991,8 +966,8 @@           (lambdaBody lam)       pure (lam', nes) -    mkP attrs p pt shape u mem ixfun is =-      Param attrs p . MemArray pt shape u . ArrayIn mem . IxFun.slice ixfun $+    mkP attrs p pt shape u mem lmad is =+      Param attrs p . MemArray pt shape u . ArrayIn mem . LMAD.slice lmad $         fmap pe64 $           Slice $             is ++ map sliceDim (shapeDims shape)@@ -1000,8 +975,8 @@     onXParam _ (Param attrs p (Prim t)) _ =       pure $ Param attrs p (MemPrim t)     onXParam is (Param attrs p (Array pt shape u)) arr = do-      (mem, ixfun) <- lookupArraySummary arr-      pure $ mkP attrs p pt shape u mem ixfun is+      (mem, lmad) <- lookupArraySummary arr+      pure $ mkP attrs p pt shape u mem lmad is     onXParam _ p _ =       error $ "Cannot handle MkAcc param: " ++ prettyString p @@ -1012,8 +987,8 @@       space <- askDefaultSpace       mem <- allocForArray arr_t space       let base_dims = map pe64 $ arrayDims arr_t-          ixfun = IxFun.iota base_dims-      pure $ mkP attrs p pt shape u mem ixfun is+          lmad = LMAD.iota 0 base_dims+      pure $ mkP attrs p pt shape u mem lmad is     onYParam _ p _ =       error $ "Cannot handle MkAcc param: " ++ prettyString p allocInExp e = mapExpM alloc e@@ -1144,7 +1119,7 @@  data ExpHint   = NoHint-  | Hint IxFun Space+  | Hint LMAD Space  defaultExpHints :: (ASTRep rep, HasScope rep m) => Exp rep -> m [ExpHint] defaultExpHints e = map (const NoHint) <$> expExtType e

@@ -12,7 +12,7 @@ import Data.Set qualified as S import Futhark.IR.GPU import Futhark.IR.GPUMem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Pass.ExplicitAllocations import Futhark.Pass.ExplicitAllocations.SegOp @@ -25,7 +25,6 @@ allocAtLevel lvl = local $ \env ->   env     { allocSpace = space,-      aggressiveReuse = True,       allocInOp = handleHostOp (Just lvl)     }   where@@ -102,8 +101,8 @@   dims <- arrayDims <$> lookupType v   let perm_inv = rearrangeInverse perm       dims' = rearrangeShape perm dims-      ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv-  pure [Hint ixfun $ Space "device"]+      lmad = LMAD.permute (LMAD.iota 0 $ map pe64 dims') perm_inv+  pure [Hint lmad $ Space "device"] kernelExpHints (Op (Inner (SegOp (SegMap lvl@(SegThread _ _) space ts body)))) =   zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body kernelExpHints (Op (Inner (SegOp (SegRed lvl@(SegThread _ _) space reds ts body)))) =@@ -133,7 +132,7 @@           pure $ Hint (innermost space_dims (arrayDims t)) $ Space "device"     hint _ _ = pure NoHint -innermost :: [SubExp] -> [SubExp] -> IxFun+innermost :: [SubExp] -> [SubExp] -> LMAD innermost space_dims t_dims =   let r = length t_dims       dims = space_dims ++ t_dims@@ -142,9 +141,9 @@           ++ [0 .. length space_dims - 1]       perm_inv = rearrangeInverse perm       dims_perm = rearrangeShape perm dims-      ixfun_base = IxFun.iota $ map pe64 dims_perm-      ixfun_rearranged = IxFun.permute ixfun_base perm_inv-   in ixfun_rearranged+      lmad_base = LMAD.iota 0 $ map pe64 dims_perm+      lmad_rearranged = LMAD.permute lmad_base perm_inv+   in lmad_rearranged  semiStatic :: S.Set VName -> SubExp -> Bool semiStatic _ Constant {} = True@@ -163,9 +162,8 @@                   dims = seg_dims ++ map pe64 (arrayDims t)                   nilSlice d = DimSlice 0 d 0                in Hint-                    ( IxFun.slice (IxFun.iota dims) $-                        fullSliceNum dims $-                          map nilSlice seg_dims+                    ( LMAD.slice (LMAD.iota 0 dims) $+                        fullSliceNum dims (map nilSlice seg_dims)                     )                     $ ScalarSpace (arrayDims t)                     $ elemType t@@ -183,8 +181,8 @@     maybePrivate consts t       | Just (Array pt shape _) <- hasStaticShape t,         all (semiStatic consts) $ shapeDims shape = do-          let ixfun = IxFun.iota $ map pe64 $ shapeDims shape-          pure $ Hint ixfun $ ScalarSpace (shapeDims shape) pt+          let lmad = LMAD.iota 0 $ map pe64 $ shapeDims shape+          pure $ Hint lmad $ ScalarSpace (shapeDims shape) pt       | otherwise =           pure NoHint 

@@ -9,7 +9,7 @@  import Control.Monad import Futhark.IR.GPUMem-import Futhark.IR.Mem.IxFun qualified as IxFun+import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Pass.ExplicitAllocations  instance SizeSubst (SegOp lvl rep)@@ -51,19 +51,19 @@               BinOp (Mul Int64 OverflowUndef) num_threads (intConst Int64 2)           let t = paramType x `arrayOfRow` twice_num_threads           mem <- allocForArray t =<< askDefaultSpace-          -- XXX: this iota ixfun is a bit inefficient; leading to+          -- XXX: this iota lmad is a bit inefficient; leading to           -- uncoalesced access.           let base_dims = map pe64 $ arrayDims t-              ixfun_base = IxFun.iota base_dims-              ixfun_x =-                IxFun.slice ixfun_base $+              lmad_base = LMAD.iota 0 base_dims+              lmad_x =+                LMAD.slice lmad_base $                   fullSliceNum base_dims [DimFix my_id]-              ixfun_y =-                IxFun.slice ixfun_base $+              lmad_y =+                LMAD.slice lmad_base $                   fullSliceNum base_dims [DimFix other_id]           pure-            ( x {paramDec = MemArray pt shape u $ ArrayIn mem ixfun_x},-              y {paramDec = MemArray pt shape u $ ArrayIn mem ixfun_y}+            ( x {paramDec = MemArray pt shape u $ ArrayIn mem lmad_x},+              y {paramDec = MemArray pt shape u $ ArrayIn mem lmad_y}             )         Prim bt ->           pure

@@ -43,6 +43,9 @@       <*> o JSON..: "duration"       <*> o JSON..: "description" +-- | A profiling report contains all profiling information for a+-- single benchmark (meaning a single invocation on an entry point on+-- a specific dataset). data ProfilingReport = ProfilingReport   { profilingEvents :: [ProfilingEvent],     -- | Mapping memory spaces to bytes.@@ -63,8 +66,10 @@       <$> o JSON..: "events"       <*> (JSON.toMapText <$> o JSON..: "memory") +-- | Read a profiling report from a bytestring containing JSON. decodeProfilingReport :: LBS.ByteString -> Maybe ProfilingReport decodeProfilingReport = JSON.decode +-- | Read a profiling report from a text containing JSON. profilingReportFromText :: T.Text -> Maybe ProfilingReport profilingReportFromText = JSON.decode . toLazyByteString . encodeUtf8Builder

@@ -21,6 +21,7 @@     partitionMaybe,     maybeNth,     maybeHead,+    lookupWithIndex,     splitFromEnd,     splitAt3,     focusNth,@@ -160,7 +161,7 @@ mapEither :: (a -> Either b c) -> [a] -> ([b], [c]) mapEither f l = partitionEithers $ map f l --- | A combination of 'partition' and 'mapMaybe'+-- | A combination of 'Data.List.partition' and 'mapMaybe' partitionMaybe :: (a -> Maybe b) -> [a] -> ([b], [a]) partitionMaybe f = helper ([], [])   where@@ -181,6 +182,11 @@ maybeHead [] = Nothing maybeHead (x : _) = Just x +-- | Lookup a value, returning also the index at which it appears.+lookupWithIndex :: (Eq a) => a -> [(a, b)] -> Maybe (Int, b)+lookupWithIndex needle haystack =+  lookup needle $ zip (map fst haystack) (zip [0 ..] (map snd haystack))+ -- | Like 'splitAt', but from the end. splitFromEnd :: Int -> [a] -> ([a], [a]) splitFromEnd i l = splitAt (length l - i) l@@ -469,13 +475,15 @@   let x' = f x    in if x' == x then x else fixPoint f x' +-- | Like 'concatMap', but monoidal and monadic. concatMapM :: (Monad m, Monoid b) => (a -> m b) -> [a] -> m b concatMapM f xs = mconcat <$> mapM f xs --- | Topological sorting of an array with an adjancency function,--- if there is a cycle, it cause an error--- @a `dep` b@ means 'a -> b', and the returned array guarantee that for i < j,--- @not ( (ret !! j) `dep` (ret !! i) )@.+-- | Topological sorting of an array with an adjancency function, if+-- there is a cycle, it causes an error. @dep a b@ means @a -> b@,+-- and the returned array guarantee that for i < j:+--+-- @not ( dep (ret !! j) (ret !! i) )@. topologicalSort :: (a -> a -> Bool) -> [a] -> [a] topologicalSort dep nodes =   fst $ execState (mapM_ (sorting . snd) nodes_idx) (mempty, mempty)

@@ -24,6 +24,8 @@     hypotf,     ldexp,     ldexpf,+    copysign,+    copysignf,   ) where @@ -162,3 +164,15 @@ -- | The system-level @ldexpf@ function. ldexpf :: Float -> CInt -> Float ldexpf = c_ldexpf++foreign import ccall "copysign" c_copysign :: Double -> Double -> Double++foreign import ccall "copysignf" c_copysignf :: Float -> Float -> Float++-- | The system-level @copysign@ function.+copysign :: Double -> Double -> Double+copysign = c_copysign++-- | The system-level @copysignf@ function.+copysignf :: Float -> Float -> Float+copysignf = c_copysignf

@@ -495,8 +495,7 @@              : '(' TypeExp ')'                { TEParens $2 (srcspan $1 $>) }              | '(' ')'                        { TETuple [] (srcspan $1 $>) }              | '(' TypeExp ',' TupleTypes ')' { TETuple ($2:$4) (srcspan $1 $>) }-             | '{' '}'                        { TERecord [] (srcspan $1 $>) }-             | '{' FieldTypes1 '}'            { TERecord $2 (srcspan $1 $>) }+             | '{' FieldTypes '}'             { TERecord $2 (srcspan $1 $>) }              | SizeExp TypeExpTerm            { TEArray $1 $2 (srcspan $1 $>) }              | QualName                       { TEVar (fst $1) (srclocOf (snd $1)) } @@ -512,12 +511,14 @@ FieldType :: { (Name, UncheckedTypeExp) } FieldType : FieldId ':' TypeExp { (fst $1, $3) } -FieldTypes1 :: { [(Name, UncheckedTypeExp)] }-FieldTypes1 : FieldType                 { [$1] }-            | FieldType ',' FieldTypes1 { $1 : $3 }+FieldTypes :: { [(Name, UncheckedTypeExp)] }+FieldTypes :                          { [] }+           | FieldType                { [$1] }+           | FieldType ',' FieldTypes { $1 : $3 }  TupleTypes :: { [UncheckedTypeExp] }             : TypeExp                { [$1] }+            | TypeExp ','            { [$1] }             | TypeExp ',' TupleTypes { $1 : $3 }  @@ -658,6 +659,7 @@  Exps1_ :: { [UncheckedExp] }         : Exps1_ ',' Exp { $3 : $1 }+        | Exps1_ ','     { $1 }         | Exp            { [$1] }  FieldAccesses :: { [(Name, Loc)] }@@ -672,12 +674,9 @@        | id              { let L loc (ID s) = $1 in RecordFieldImplicit s NoInfo (srclocOf loc) }  Fields :: { [FieldBase NoInfo Name] }-        : Fields1 { $1 }-        |         { [] }--Fields1 :: { [FieldBase NoInfo Name] }-        : Field ',' Fields1 { $1 : $3 }-        | Field             { [$1] }+       : Field ',' Fields { $1 : $3 }+       | Field            { [$1] }+       |                  { [] }  LetExp :: { UncheckedExp }      : let SizeBinders1 Pat '=' Exp LetBody@@ -807,7 +806,8 @@                                         in PatConstr n NoInfo [] (srclocOf loc) }  Pats1 :: { [PatBase NoInfo Name StructType] }-           : Pat               { [$1] }+           : Pat            { [$1] }+           | Pat ','       { [$1] }            | Pat ',' Pats1 { $1 : $3 }  InnerPat :: { PatBase NoInfo Name StructType }@@ -840,7 +840,8 @@  CFieldPats1 :: { [(Name, PatBase NoInfo Name StructType)] }                  : CFieldPat ',' CFieldPats1 { $1 : $3 }-                 | CFieldPat                    { [$1] }+                 | CFieldPat ','             { [$1] }+                 | CFieldPat                 { [$1] }  PatLiteralNoNeg :: { (PatLit, Loc) }              : charlit  { let L loc (CHARLIT x) = $1@@ -877,12 +878,9 @@          |      ':'      ':' Exp2 { DimSlice Nothing Nothing (Just $3) }  DimIndices :: { [UncheckedDimIndex] }-            :             { [] }-            | DimIndices1 { fst $1 : snd $1 }--DimIndices1 :: { (UncheckedDimIndex, [UncheckedDimIndex]) }-             : DimIndex                 { ($1, []) }-             | DimIndex ',' DimIndices1 { ($1, fst $3 : snd $3) }+             :                         { [] }+             | DimIndex                { [$1] }+             | DimIndex ',' DimIndices { $1 : $3 }  VarId :: { IdentBase NoInfo Name StructType } VarId : id { let L loc (ID name) = $1 in Ident name NoInfo (srclocOf loc) }@@ -901,9 +899,9 @@  AttrInfo :: { AttrInfo Name }          : AttrAtom         { let (x,y) = $1 in AttrAtom x (srclocOf y) }-         | id '('       ')' { let L _ (ID s) = $1 in AttrComp s [] (srcspan $1 $>) }          | id '(' Attrs ')' { let L _ (ID s) = $1 in AttrComp s $3 (srcspan $1 $>) }  Attrs :: { [AttrInfo Name] }-       : AttrInfo           { [$1] }+       :                    { [] }+       | AttrInfo           { [$1] }        | AttrInfo ',' Attrs { $1 : $3 }

@@ -7,6 +7,7 @@     prettyTuple,     leadingOperator,     IsName (..),+    prettyNameString,     Annot (..),   ) where@@ -53,6 +54,10 @@ instance IsName Name where   prettyName = pretty   toName = id++-- | Prettyprint name as string.  Only use this for debugging.+prettyNameString :: (IsName v) => v -> String+prettyNameString = T.unpack . docText . prettyName  -- | Class for type constructors that represent annotations.  Used in -- the prettyprinter to either print the original AST, or the computed

@@ -1316,6 +1316,10 @@       f32 "erfc32" erfcf,       f64 "erfc64" erfc,       --+      f16_2 "copysign16" $ \x y -> convFloat (copysign (convFloat x) (convFloat y)),+      f32_2 "copysign32" copysignf,+      f64_2 "copysign64" copysign,+      --       i8 "clz8" $ IntValue . Int32Value . fromIntegral . countLeadingZeros,       i16 "clz16" $ IntValue . Int32Value . fromIntegral . countLeadingZeros,       i32 "clz32" $ IntValue . Int32Value . fromIntegral . countLeadingZeros,

@@ -749,6 +749,7 @@ niceTypeExp (TEApply te TypeArgExpSize {} _) = niceTypeExp te niceTypeExp (TEArray _ te _) = niceTypeExp te niceTypeExp (TEUnique te _) = niceTypeExp te+niceTypeExp (TEDim _ te _) = niceTypeExp te niceTypeExp _ = False  checkOneDec :: DecBase NoInfo Name -> TypeM (TySet, Env, DecBase Info VName)

@@ -557,8 +557,9 @@ -- functions. type Loop = (Pat ParamType, Exp, LoopFormBase Info VName, Exp) --- | Mark bindings of consumed names as Consume.-updateParamDiet :: Names -> Pat ParamType -> Pat ParamType+-- | Mark bindings of consumed names as Consume, except those under a+-- 'PatAscription', which are left unchanged.+updateParamDiet :: (VName -> Bool) -> Pat ParamType -> Pat ParamType updateParamDiet cons = recurse   where     recurse (Wildcard (Info t) wloc) =@@ -568,7 +569,7 @@     recurse (PatAttr attr p ploc) =       PatAttr attr (recurse p) ploc     recurse (Id name (Info t) iloc)-      | name `S.member` cons =+      | cons name =           let t' = t `setUniqueness` Consume            in Id name (Info t') iloc       | otherwise =@@ -587,7 +588,7 @@ convergeLoopParam :: Loc -> Pat ParamType -> Names -> TypeAliases -> CheckM (Pat ParamType) convergeLoopParam loop_loc param body_cons body_als = do   let -- Make the pattern Consume where needed.-      param' = updateParamDiet (S.filter (`elem` patNames param) body_cons) param+      param' = updateParamDiet (`S.member` S.filter (`elem` patNames param) body_cons) param    -- Check that the new values of consumed merge parameters do not   -- alias something bound outside the loop, AND that anything@@ -653,7 +654,7 @@   -- use to infer the proper diet of the parameter.   ((body', body_cons), body_als) <-     noConsumable-      . bindingParam (fmap (second (const Consume)) param)+      . bindingParam (updateParamDiet (const True) param)       . bindingLoopForm form'       $ do         ((body', body_als), body_cons) <- contain $ checkExp body

@@ -614,9 +614,9 @@   -- Not technically an ascription, but we want the pattern to have   -- exactly the type of 'e'.   t <- expType e'-  incLevel . bindingSizes sizes $ \sizes' ->-    bindingPat sizes' pat t $ \pat' -> do-      body' <- checkExp body+  bindingSizes sizes $ \sizes' ->+    incLevel . bindingPat sizes' pat t $ \pat' -> do+      body' <- incLevel $ checkExp body       body_t <- expTypeFully body'        -- If the bound expression is of type i64, then we replace the

@@ -346,4 +346,4 @@             binding (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps         descend ps' [] = m tps' $ reverse ps' -    descend [] orig_ps+    incLevel $ descend [] orig_ps

@@ -817,31 +817,33 @@    modifyConstraints $ M.insert vn (lvl, Size (Just e) usage)   where+    checkVar _ dim'+      | vn == dim' = do+          notes <- dimNotes usage e+          unifyError usage notes bcs $+            "Occurs check: cannot instantiate"+              <+> dquotes (prettyName vn)+              <+> "with"+              <+> dquotes (pretty e)+              <+> "."     checkVar constraints dim'       | Just (dim_lvl, c) <- dim' `M.lookup` constraints,-        dim_lvl > lvl =+        dim_lvl >= lvl =           case c of             ParamSize {} -> do               notes <- dimNotes usage e               unifyError usage notes bcs $-                "Cannot unify size variable"-                  <+> dquotes (pretty e)-                  <+> "with"+                "Cannot link size"                   <+> dquotes (prettyName vn)+                  <+> "to"+                  <+> dquotes (pretty e)                   <+> "(scope violation)."                   </> "This is because"                   <+> dquotes (pretty $ qualName dim')-                  <+> "is rigidly bound in a deeper scope."+                  <+> "is not in scope when"+                  <+> dquotes (prettyName vn)+                  <+> "is introduced."             _ -> modifyConstraints $ M.insert dim' (lvl, c)-    checkVar _ dim'-      | vn == dim' = do-          notes <- dimNotes usage e-          unifyError usage notes bcs $-            "Occurs check: cannot instantiate"-              <+> dquotes (prettyName vn)-              <+> "with"-              <+> dquotes (pretty e)-              <+> "."     checkVar _ _ = pure ()  -- | Assert that this type must be one of the given primitive types.@@ -934,8 +936,6 @@           pure () -- All primtypes support equality.         Just (_, Equality {}) ->           pure ()-        Just (_, HasConstrs _ cs _) ->-          mapM_ (equalityType usage) $ concat $ M.elems cs         _ ->           unifyError usage mempty noBreadCrumbs $             "Type" <+> prettyName vn <+> "does not support equality."@@ -1221,7 +1221,7 @@           M.insert dim (0, Size Nothing usage)     pure dim -  curLevel = pure 0+  curLevel = pure 1    unifyError loc notes bcs doc =     throwError $ TypeError (locOf loc) notes $ doc <> pretty bcs@@ -1243,8 +1243,8 @@     constraints =       M.fromList $         map nonrigid nonrigid_tparams <> map rigid rigid_tparams-    nonrigid (TypeParamDim p loc) = (p, (0, Size Nothing $ Usage Nothing loc))-    nonrigid (TypeParamType l p loc) = (p, (0, NoConstraint l $ Usage Nothing loc))+    nonrigid (TypeParamDim p loc) = (p, (1, Size Nothing $ Usage Nothing loc))+    nonrigid (TypeParamType l p loc) = (p, (1, NoConstraint l $ Usage Nothing loc))     rigid (TypeParamDim p loc) = (p, (0, ParamSize loc))     rigid (TypeParamType l p loc) = (p, (0, ParamType l loc)) 

@@ -28,6 +28,7 @@ instance Monoid Warnings where   mempty = Warnings mempty +-- | Prettyprint warnings, making use of colours and such. prettyWarnings :: Warnings -> Doc AnsiStyle prettyWarnings (Warnings []) = mempty prettyWarnings (Warnings ws) =

@@ -11,10 +11,10 @@ import Data.Map qualified as M import Data.Text qualified as T import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.Mem.IxFun qualified as IxFunLMAD import Futhark.IR.Mem.IxFun.Alg qualified as IxFunAlg import Futhark.IR.Mem.IxFunWrapper import Futhark.IR.Mem.IxFunWrapper qualified as IxFunWrap+import Futhark.IR.Mem.LMAD qualified as IxFunLMAD import Futhark.IR.Prop import Futhark.IR.Syntax import Futhark.IR.Syntax.Core ()@@ -49,7 +49,7 @@           ([], x)           strides -compareIxFuns :: Maybe (IxFunLMAD.IxFun Int) -> IxFunAlg.IxFun Int -> Assertion+compareIxFuns :: Maybe (IxFunLMAD.LMAD Int) -> IxFunAlg.IxFun Int -> Assertion compareIxFuns (Just ixfunLMAD) ixfunAlg =   let lmadShape = IxFunLMAD.shape ixfunLMAD       algShape = IxFunAlg.shape ixfunAlg

@@ -13,8 +13,8 @@ where  import Control.Monad (join)-import Futhark.IR.Mem.IxFun qualified as I import Futhark.IR.Mem.IxFun.Alg qualified as IA+import Futhark.IR.Mem.LMAD qualified as I import Futhark.IR.Syntax (FlatSlice, Slice) import Futhark.Util.IntegralExp @@ -22,16 +22,15 @@  type Permutation = [Int] -type IxFun num = (Maybe (I.IxFun num), IA.IxFun num)+type IxFun num = (Maybe (I.LMAD num), IA.IxFun num)  iota ::   (IntegralExp num) =>   Shape num ->   IxFun num-iota x = (Just $ I.iota x, IA.iota x)+iota x = (Just $ I.iota 0 x, IA.iota x)  permute ::-  (IntegralExp num) =>   IxFun num ->   Permutation ->   IxFun num@@ -45,7 +44,6 @@ reshape (l, a) x = (join (I.reshape <$> l <*> pure x), IA.reshape a x)  coerce ::-  (Eq num, IntegralExp num) =>   IxFun num ->   Shape num ->   IxFun num@@ -59,16 +57,16 @@ slice (l, a) x = (I.slice <$> l <*> pure x, IA.slice a x)  flatSlice ::-  (Eq num, IntegralExp num) =>+  (IntegralExp num) =>   IxFun num ->   FlatSlice num ->   IxFun num flatSlice (l, a) x = (I.flatSlice <$> l <*> pure x, IA.flatSlice a x)  expand ::-  (Eq num, IntegralExp num) =>+  (IntegralExp num) =>   num ->   num ->   IxFun num ->   IxFun num-expand o p (lf, af) = (I.expand o p =<< lf, IA.expand o p af)+expand o p (lf, af) = (Just . I.expand o p =<< lf, IA.expand o p af)

@@ -48,10 +48,9 @@               ]           )           @?= ( [Pure "i64"],-                M.fromList-                  [ ("foo", (1, [0])),-                    ("bar", (0, [0]))-                  ]+                [ ("bar", [0]),+                  ("foo", [0])+                ]               ),       testCase "Dedup of array" $         internaliseConstructors@@ -61,10 +60,9 @@               ]           )           @?= ( [Pure "[?0]i64"],-                M.fromList-                  [ ("foo", (1, [0])),-                    ("bar", (0, [0]))-                  ]+                [ ("bar", [0]),+                  ("foo", [0])+                ]               ),       testCase         "Dedup of array of tuple"@@ -75,10 +73,9 @@               ]           )           @?= ( [Pure "[?0]i64", Free [Pure "[?0]i64", Pure "[?0]i64"]],-                M.fromList-                  [ ("foo", (1, [1, 2])),-                    ("bar", (0, [0]))-                  ]+                [ ("bar", [0]),+                  ("foo", [1, 2])+                ]               )     ]