futhark 0.23.1 → 0.24.1

raw patch · 94 files changed

+1593/−1330 lines, 94 filesdep −processdep ~futhark-manifestdep ~futhark-serverPVP ok

version bump matches the API change (PVP)

Dependencies removed: process

Dependency ranges changed: futhark-manifest, futhark-server

API changes (from Hackage documentation)

Files

• docs/c-api.rst +11/−12
• docs/error-index.rst +35/−0
• docs/installation.rst +1/−1
• docs/language-reference.rst +20/−8
• docs/server-protocol.rst +16/−2
• futhark.cabal +3/−4
• prelude/ad.fut +2/−2
• prelude/array.fut +3/−3
• prelude/soacs.fut +15/−15
• prelude/zip.fut +2/−2
• rts/c/backends/cuda.h +0/−2
• rts/c/backends/opencl.h +0/−2
• rts/c/context.h +1/−0
• rts/c/server.h +41/−0
• src/Futhark/AD/Rev.hs +5/−1
• src/Futhark/AD/Rev/Map.hs +20/−6
• src/Futhark/AD/Rev/Monad.hs +3/−0
• src/Futhark/AD/Rev/SOAC.hs +18/−2
• src/Futhark/AD/Rev/Scan.hs +1/−1
• src/Futhark/Analysis/PrimExp.hs +22/−15
• src/Futhark/CLI/Autotune.hs +30/−26
• src/Futhark/CLI/Dataset.hs +2/−2
• src/Futhark/CLI/Dev.hs +1/−1
• src/Futhark/CLI/Doc.hs +2/−1
• src/Futhark/CLI/Literate.hs +1/−1
• src/Futhark/CLI/Misc.hs +2/−1
• src/Futhark/CLI/REPL.hs +7/−5
• src/Futhark/CLI/Test.hs +1/−1
• src/Futhark/CodeGen/Backends/CCUDA.hs +3/−3
• src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +1/−15
• src/Futhark/CodeGen/Backends/COpenCL.hs +2/−2
• src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +2/−34
• src/Futhark/CodeGen/Backends/GenericC.hs +48/−7
• src/Futhark/CodeGen/Backends/GenericC/CLI.hs +1/−1
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +1/−1
• src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs +15/−4
• src/Futhark/CodeGen/Backends/GenericC/Fun.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericC/Server.hs +9/−1
• src/Futhark/CodeGen/Backends/GenericPython.hs +1/−1
• src/Futhark/CodeGen/Backends/MulticoreC.hs +1/−2
• src/Futhark/CodeGen/Backends/MulticoreC/Boilerplate.hs +0/−5
• src/Futhark/CodeGen/Backends/MulticoreISPC.hs +8/−8
• src/Futhark/CodeGen/Backends/MulticoreWASM.hs +1/−0
• src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +4/−4
• src/Futhark/CodeGen/Backends/SequentialC.hs +10/−1
• src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs +0/−5
• src/Futhark/CodeGen/Backends/SequentialWASM.hs +1/−0
• src/Futhark/CodeGen/ImpCode.hs +6/−1
• src/Futhark/CodeGen/ImpCode/OpenCL.hs +1/−1
• src/Futhark/CodeGen/ImpGen.hs +1/−1
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +19/−11
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +33/−5
• src/Futhark/Compiler/Program.hs +19/−19
• src/Futhark/Doc/Generator.hs +48/−37
• src/Futhark/IR/Mem/IxFun.hs +14/−10
• src/Futhark/IR/Mem/Simplify.hs +1/−30
• src/Futhark/IR/Pretty.hs +5/−5
• src/Futhark/IR/SOACS/Simplify.hs +3/−0
• src/Futhark/IR/SegOp.hs +2/−2
• src/Futhark/Internalise/Defunctionalise.hs +176/−194
• src/Futhark/Internalise/Defunctorise.hs +6/−6
• src/Futhark/Internalise/Entry.hs +7/−7
• src/Futhark/Internalise/Exps.hs +50/−53
• src/Futhark/Internalise/LiftLambdas.hs +6/−4
• src/Futhark/Internalise/Monomorphise.hs +36/−48
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +1/−1
• src/Futhark/Optimise/EntryPointMem.hs +1/−1
• src/Futhark/Optimise/MergeGPUBodies.hs +1/−1
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +1/−3
• src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +3/−7
• src/Futhark/Pass/ExpandAllocations.hs +2/−2
• src/Futhark/Pass/ExplicitAllocations.hs +1/−1
• src/Futhark/Pass/ExtractMulticore.hs +3/−3
• src/Futhark/Tools.hs +23/−0
• src/Futhark/Util.hs +15/−9
• src/Language/Futhark/FreeVars.hs +2/−2
• src/Language/Futhark/Interpreter.hs +102/−95
• src/Language/Futhark/Parser/Monad.hs +1/−2
• src/Language/Futhark/Parser/Parser.y +12/−33
• src/Language/Futhark/Pretty.hs +25/−15
• src/Language/Futhark/Prop.hs +222/−193
• src/Language/Futhark/Query.hs +7/−9
• src/Language/Futhark/Semantic.hs +7/−17
• src/Language/Futhark/Syntax.hs +120/−78
• src/Language/Futhark/Traversals.hs +39/−24
• src/Language/Futhark/TypeChecker.hs +13/−13
• src/Language/Futhark/TypeChecker/Modules.hs +2/−2
• src/Language/Futhark/TypeChecker/Monad.hs +9/−9
• src/Language/Futhark/TypeChecker/Terms.hs +81/−90
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +12/−17
• src/Language/Futhark/TypeChecker/Types.hs +62/−46
• src/Language/Futhark/TypeChecker/Unify.hs +6/−6
• unittests/Language/Futhark/SyntaxTests.hs +11/−4
• unittests/Language/Futhark/TypeChecker/TypesTests.hs +1/−1

@@ -133,7 +133,7 @@    or stale, the program performs initialisation from scratch.  There    is no machine-readable way to get information about whether the    cache was hit succesfully, but you can enable logging to see what-   hapens.+   happens.     The lifespan of ``fname`` must exceed the lifespan of the    configuration object.  Pass ``NULL`` to disable caching (this is@@ -239,7 +239,7 @@ that are created with a ``new`` function, or returned from an entry point, *must* at some point be freed manually.  Values are internally reference counted, so even for entry points that return their input-unchanged, you should still free both the input and the output - this+unchanged, you must still free both the input and the output - this will not result in a double free.  .. c:struct:: futhark_i32_1d@@ -281,23 +281,22 @@     Return a pointer to the shape of the array, with one element per    dimension.  The lifetime of the shape is the same as ``arr``, and-   should *not* be manually freed.  Assuming ``arr`` is a valid+   must *not* be manually freed.  Assuming ``arr`` is a valid    object, this function cannot fail.  .. _opaques: -Opaque values+Opaque Values ~~~~~~~~~~~~~  Each instance of a complex type in an entry point (records, nested tuples, etc) is represented by an opaque C struct named ``futhark_opaque_foo``.  In the general case, ``foo`` will be a hash-of the internal representation.  However, if you insert explicit type-annotations in the entry point (and the type name contains only-characters valid for C identifiers), the indicated name will be used.-Note that arrays contain brackets, which are usually not valid in-identifiers.  Defining a simple type abbreviation is the best way-around this.+of the internal representation.  However, if you insert an explicit+type annotations in the entry point (and the type name contains only+characters valid in C identifiers), that name will be used.  Note that+arrays contain brackets, which are not valid in identifiers.  Defining+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@@ -354,7 +353,7 @@ ~~~~~~~  A record is an opaque type (see above) that supports additional-functions to *project* individual fields (read their values) to+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. @@ -501,7 +500,7 @@    Use exactly this command queue for the context.  If this is set,    all other device/platform configuration options are ignored.  Once    the context is active, the command queue belongs to Futhark and-   should not be used by anything else.  This is useful for+   must not be used by anything else.  This is useful for    implementing custom device selection logic in application code.  .. c:function:: cl_command_queue futhark_context_get_command_queue(struct futhark_context *ctx)

@@ -149,6 +149,41 @@    def main (xs: *[]i32) : (*[]i32, *[]i32) = (xs, copy xs) +.. _self-aliasing-arg:++"Argument passed for consuming parameter is self-aliased."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++Caused by programs like the following:++.. code-block:: futhark++  def g (t: *([]i64, []i64)) = 0++  def f n =+    let x = iota n+    in g (x,x)++The function ``g`` expects to consume two separate ``[]i64`` arrays,+but ``f`` passes it a tuple containing two references to the same+physical array.  This is not allowed, as ``g`` must be allowed to+assume that components of consuming record- or tuple parameters have+no internal aliases.  We can fix this by inserting copies to break the+aliasing:++.. code-block:: futhark++  def f n =+    let x = iota n+    in g (copy (x,x))++Alternative, we could duplicate the expression producing the array:++.. code-block:: futhark++  def f n =+    g (iota n, iota n))+ .. _consuming-parameter:  "Consuming parameter passed non-unique argument"

@@ -97,7 +97,7 @@   `futhark-nightly-linux-x86_64.tar.xz <https://futhark-lang.org/releases/futhark-nightly-linux-x86_64.tar.xz>`_  **macOS (x86_64)**-  `futhark-nightly-macos-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-x86_64.zip>`_+  `futhark-nightly-macos-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-x86_64.tar.xz>`_  **Windows (x86_64)**   `futhark-nightly-windows-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-windows-x86_64.zip>`_

@@ -1434,15 +1434,27 @@   def modify (a: *[]i32) (i: i32) (x: i32): *[]i32 =     a with [i] = a[i] + x -For bulk in-place updates with multiple values, use the ``scatter``-function in the basis library.  In the parameter declaration ``a:-*[i32]``, the asterisk means that the function ``modify`` has been-given "ownership" of the array ``a``, meaning that any caller of-``modify`` will never reference array ``a`` after the call again.-This allows the ``with`` expression to perform an in-place update.+A parameter that is not consuming is called *observing*.  In the+parameter declaration ``a: *[i32]``, the asterisk means that the+function ``modify`` has been given "ownership" of the array ``a``,+meaning that any caller of ``modify`` will never reference array ``a``+after the call again.  This allows the ``with`` expression to perform+an in-place update.  After a call ``modify a i x``, neither ``a`` or+any variable that *aliases* ``a`` may be used on any following+execution path. -After a call ``modify a i x``, neither ``a`` or any variable that-*aliases* ``a`` may be used on any following execution path.+If an asterisk is present at *any point* inside a tuple parameter+type, the parameter as a whole is considered consuming.  For example::++  def consumes_both ((a,b): (*[]i32,[]i32)) = ...++This is usually not desirable behaviour.  Use multiple parameters+instead::++  def consumes_first_arg (a: *[]i32) (b: []i32) = ...++For bulk in-place updates with multiple values, use the ``scatter``+function in the basis library.  Alias Analysis ~~~~~~~~~~~~~~

@@ -146,10 +146,24 @@  Corresponds to :c:func:`futhark_context_report`. -``set_tuning_param``-....................+``set_tuning_param`` *param* *value*+....................................  Corresponds to :c:func:`futhark_context_config_set_tuning_param`.++``tuning_params`` *entry*+.........................++For each tuning parameters relevant to the given entry point, print+its name, then a space, then its class.++This is similar to on :c:func:`futhark_tuning_params_for_sum`, but+note that this command prints *names* and not *integers*.++``tuning_param_class`` *param*+..............................++Corresponds to :c:func:`futhark_get_tuning_param_class`.  Record Commands ~~~~~~~~~~~~~~~

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.23.1+version:        0.24.1 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -483,8 +483,8 @@     , filepath >=1.4.1.1     , free >=4.12.4     , futhark-data >= 1.1.0.0-    , futhark-server >= 1.2.1.0-    , futhark-manifest >= 1.1.0.0+    , futhark-server >= 1.2.2.0+    , futhark-manifest >= 1.2.0.1     , githash >=0.1.6.1     , half >= 0.3     , haskeline@@ -498,7 +498,6 @@     , neat-interpolation >=0.3     , parallel >=3.2.1.0     , random >= 1.2.0-    , process >=1.4.3.0     , process-extras >=0.7.2     , regex-tdfa >=1.2     , srcloc >=0.4

@@ -7,12 +7,12 @@ -- | Jacobian-Vector Product ("forward mode"), producing also the -- primal result as the first element of the result tuple. let jvp2 'a 'b (f: a -> b) (x: a) (x': a) : (b, b) =-  intrinsics.jvp2 (f, x, x')+  intrinsics.jvp2 f x x'  -- | Vector-Jacobian Product ("reverse mode"), producing also the -- primal result as the first element of the result tuple. let vjp2 'a 'b (f: a -> b) (x: a) (y': b) : (b, a) =-  intrinsics.vjp2 (f, x, y')+  intrinsics.vjp2 f x y'  -- | Jacobian-Vector Product ("forward mode"). let jvp 'a 'b (f: a -> b) (x: a) (x': a) : b =

@@ -62,7 +62,7 @@ -- **Work:** O(n). -- -- **Span:** O(1).-def (++) [n] [m] 't (xs: [n]t) (ys: [m]t): *[]t = intrinsics.concat (xs, ys)+def (++) [n] [m] 't (xs: [n]t) (ys: [m]t): *[]t = intrinsics.concat xs ys  -- | An old-fashioned way of saying `++`. def concat [n] [m] 't (xs: [n]t) (ys: [m]t): *[]t = xs ++ ys@@ -83,7 +83,7 @@ -- -- Note: In most cases, `rotate` will be fused with subsequent -- operations such as `map`, in which case it is free.-def rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs)+def rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate r xs  -- | Construct an array of consecutive integers of the given length, -- starting at 0.@@ -143,7 +143,7 @@ -- -- **Complexity:** O(1). def unflatten [p] 't (n: i64) (m: i64) (xs: [p]t): [n][m]t =-  intrinsics.unflatten (n, m, xs) :> [n][m]t+  intrinsics.unflatten n m xs :> [n][m]t  -- | Like `unflatten`, but produces three dimensions. def unflatten_3d [p] 't (n: i64) (m: i64) (l: i64) (xs: [p]t): [n][m][l]t =

@@ -48,7 +48,7 @@ -- -- **Span:** *O(S(f))* def map 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x =-  intrinsics.map (f, as)+  intrinsics.map f as  -- | Apply the given function to each element of a single array. --@@ -104,7 +104,7 @@ -- Note that the complexity implies that parallelism in the combining -- operator will *not* be exploited. def reduce [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): a =-  intrinsics.reduce (op, ne, as)+  intrinsics.reduce op ne as  -- | As `reduce`, but the operator must also be commutative.  This is -- potentially faster than `reduce`.  For simple built-in operators,@@ -115,7 +115,7 @@ -- -- **Span:** *O(log(n) ✕ W(op))* def reduce_comm [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): a =-  intrinsics.reduce_comm (op, ne, as)+  intrinsics.reduce_comm op ne as  -- | `h = hist op ne k is as` computes a generalised `k`-bin histogram -- `h`, such that `h[i]` is the sum of those values `as[j]` for which@@ -130,7 +130,7 @@ -- -- In practice, linear span only occurs if *k* is also very large. def hist 'a [n] (op: a -> a -> a) (ne: a) (k: i64) (is: [n]i64) (as: [n]a) : *[k]a =-  intrinsics.hist_1d (1, map (\_ -> ne) (0..1..<k), op, ne, is, as)+  intrinsics.hist_1d 1 (map (\_ -> ne) (0..1..<k)) op ne is as  -- | Like `hist`, but with initial contents of the histogram, and the -- complexity is proportional only to the number of input elements,@@ -143,15 +143,15 @@ -- -- In practice, linear span only occurs if *k* is also very large. def reduce_by_index 'a [k] [n] (dest : *[k]a) (f : a -> a -> a) (ne : a) (is : [n]i64) (as : [n]a) : *[k]a =-  intrinsics.hist_1d (1, dest, f, ne, is, as)+  intrinsics.hist_1d 1 dest f ne is as  -- | As `reduce_by_index`, but with two-dimensional indexes. def reduce_by_index_2d 'a [k] [n] [m] (dest : *[k][m]a) (f : a -> a -> a) (ne : a) (is : [n](i64,i64)) (as : [n]a) : *[k][m]a =-  intrinsics.hist_2d (1, dest, f, ne, is, as)+  intrinsics.hist_2d 1 dest f ne is as  -- | As `reduce_by_index`, but with three-dimensional indexes. def reduce_by_index_3d 'a [k] [n] [m] [l] (dest : *[k][m][l]a) (f : a -> a -> a) (ne : a) (is : [n](i64,i64,i64)) (as : [n]a) : *[k][m][l]a =-  intrinsics.hist_3d (1, dest, f, ne, is, as)+  intrinsics.hist_3d 1 dest f ne is as  -- | Inclusive prefix scan.  Has the same caveats with respect to -- associativity and complexity as `reduce`.@@ -160,7 +160,7 @@ -- -- **Span:** *O(log(n) ✕ W(op))* def scan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): *[n]a =-  intrinsics.scan (op, ne, as)+  intrinsics.scan op ne as  -- | Remove all those elements of `as` that do not satisfy the -- predicate `p`.@@ -169,7 +169,7 @@ -- -- **Span:** *O(log(n) ✕ W(p))* def filter [n] 'a (p: a -> bool) (as: [n]a): *[]a =-  let (as', is) = intrinsics.partition (1, \x -> if p x then 0 else 1, as)+  let (as', is) = intrinsics.partition 1 (\x -> if p x then 0 else 1) as   in as'[:is[0]]  -- | Split an array into those elements that satisfy the given@@ -180,7 +180,7 @@ -- **Span:** *O(log(n) ✕ W(p))* def partition [n] 'a (p: a -> bool) (as: [n]a): ([]a, []a) =   let p' x = if p x then 0 else 1-  let (as', is) = intrinsics.partition (2, p', as)+  let (as', is) = intrinsics.partition 2 p' as   in (as'[0:is[0]], as'[is[0]:n])  -- | Split an array by two predicates, producing three arrays.@@ -190,7 +190,7 @@ -- **Span:** *O(log(n) ✕ (W(p1) + W(p2)))* def partition2 [n] 'a (p1: a -> bool) (p2: a -> bool) (as: [n]a): ([]a, []a, []a) =   let p' x = if p1 x then 0 else if p2 x then 1 else 2-  let (as', is) = intrinsics.partition (3, p', as)+  let (as', is) = intrinsics.partition 3 p' as   in (as'[0:is[0]], as'[is[0]:is[0]+is[1]], as'[is[0]+is[1]:n])  -- | Return `true` if the given function returns `true` for all@@ -223,7 +223,7 @@ -- -- **Span:** *O(1)* def spread 't [n] (k: i64) (x: t) (is: [n]i64) (vs: [n]t): *[k]t =-  intrinsics.scatter (map (\_ -> x) (0..1..<k), is, vs)+  intrinsics.scatter (map (\_ -> x) (0..1..<k)) is vs  -- | Like `spread`, but takes an array indicating the initial values, -- and has different work complexity.@@ -232,7 +232,7 @@ -- -- **Span:** *O(1)* def scatter 't [k] [n] (dest: *[k]t) (is: [n]i64) (vs: [n]t): *[k]t =-  intrinsics.scatter (dest, is, vs)+  intrinsics.scatter dest is vs  -- | `scatter_2d as is vs` is the equivalent of a `scatter` on a 2-dimensional -- array.@@ -241,7 +241,7 @@ -- -- **Span:** *O(1)* def scatter_2d 't [k] [n] [l] (dest: *[k][n]t) (is: [l](i64, i64)) (vs: [l]t): *[k][n]t =-  intrinsics.scatter_2d (dest, is, vs)+  intrinsics.scatter_2d dest is vs  -- | `scatter_3d as is vs` is the equivalent of a `scatter` on a 3-dimensional -- array.@@ -250,4 +250,4 @@ -- -- **Span:** *O(1)* def scatter_3d 't [k] [n] [o] [l] (dest: *[k][n][o]t) (is: [l](i64, i64, i64)) (vs: [l]t): *[k][n][o]t =-  intrinsics.scatter_3d (dest, is, vs)+  intrinsics.scatter_3d dest is vs

@@ -11,11 +11,11 @@ -- depended upon by soacs.fut.  So we just define a quick-and-dirty -- internal one here that uses the intrinsic version. local def internal_map 'a [n] 'x (f: a -> x) (as: [n]a): [n]x =-  intrinsics.map (f, as)+  intrinsics.map f as  -- | Construct an array of pairs from two arrays. def zip [n] 'a 'b (as: [n]a) (bs: [n]b): *[n](a,b) =-  intrinsics.zip (as, bs)+  intrinsics.zip as bs  -- | Construct an array of pairs from two arrays. def zip2 [n] 'a 'b (as: [n]a) (bs: [n]b): *[n](a,b) =

@@ -929,8 +929,6 @@   CUDA_SUCCEED_FATAL(cuMemcpyHtoD(ctx->global_failure, &no_error, sizeof(no_error)));   // The +1 is to avoid zero-byte allocations.   CUDA_SUCCEED_FATAL(cuMemAlloc(&ctx->global_failure_args, sizeof(int64_t)*(max_failure_args+1)));--  set_tuning_params(ctx);   return 0; } 

@@ -1209,8 +1209,6 @@                    CL_MEM_READ_WRITE,                    sizeof(int64_t)*(max_failure_args+1), NULL, &error);   OPENCL_SUCCEED_OR_RETURN(error);--  set_tuning_params(ctx);   return 0; } 

@@ -122,6 +122,7 @@   ctx->error = NULL;   ctx->log = stderr;   if (backend_context_setup(ctx) == 0) {+    set_tuning_params(ctx);     setup_program(ctx);     init_constants(ctx);     (void)futhark_context_clear_caches(ctx);

@@ -10,6 +10,9 @@ int futhark_context_config_set_tuning_param(struct futhark_context_config *cfg,                                             const char *param_name,                                             size_t new_value);+int futhark_get_tuning_param_count(void);+const char* futhark_get_tuning_param_name(int i);+const char* futhark_get_tuning_param_class(int i);  typedef int (*restore_fn)(const void*, FILE *, struct futhark_context*, void*); typedef void (*store_fn)(const void*, FILE *, struct futhark_context*, void*);@@ -154,6 +157,7 @@ struct entry_point {   const char *name;   entry_point_fn f;+  const char** tuning_params;   const struct type **out_types;   bool *out_unique;   const struct type **in_types;@@ -566,6 +570,39 @@   } } +void cmd_tuning_params(struct server_state *s, const char *args[]) {+  const char *name = get_arg(args, 0);+  struct entry_point *e = get_entry_point(s, name);++  if (e == NULL) {+    failure();+    printf("Unknown entry point: %s\n", name);+    return;+  }++  const char **params = e->tuning_params;+  for (int i = 0; params[i] != NULL; i++) {+    printf("%s\n", params[i]);+  }+}++void cmd_tuning_param_class(struct server_state *s, const char *args[]) {+  (void)s;+  const char *param = get_arg(args, 0);++  int n = futhark_get_tuning_param_count();++  for (int i = 0; i < n; i++) {+    if (strcmp(futhark_get_tuning_param_name(i), param) == 0) {+      printf("%s\n", futhark_get_tuning_param_class(i));+      return;+    }+  }++  failure();+  printf("Unknown tuning parameter: %s\n", param);+}+ void cmd_fields(struct server_state *s, const char *args[]) {   const char *type = get_arg(args, 0);   const struct type *t = get_type(s, type);@@ -784,6 +821,10 @@     cmd_report(s, tokens+1);   } else if (strcmp(command, "set_tuning_param") == 0) {     cmd_set_tuning_param(s, tokens+1);+  } else if (strcmp(command, "tuning_params") == 0) {+    cmd_tuning_params(s, tokens+1);+  } else if (strcmp(command, "tuning_param_class") == 0) {+    cmd_tuning_param_class(s, tokens+1);   } else if (strcmp(command, "fields") == 0) {     cmd_fields(s, tokens+1);   } else if (strcmp(command, "new") == 0) {

@@ -229,7 +229,11 @@           updateSubExpAdj v adj_i  vjpOps :: VjpOps-vjpOps = VjpOps diffLambda diffStm+vjpOps =+  VjpOps+    { vjpLambda = diffLambda,+      vjpStm = diffStm+    }  diffStm :: Stm SOACS -> ADM () -> ADM () diffStm (Let pat aux (BasicOp e)) m =

@@ -1,5 +1,7 @@ {-# LANGUAGE TypeFamilies #-} +-- | VJP transformation for Map SOACs.  This is a pretty complicated+-- case due to the possibility of free variables. module Futhark.AD.Rev.Map (vjpMap) where  import Control.Monad@@ -73,6 +75,8 @@     subAD $ mkLambda (cert_params ++ acc_params) $ m $ map paramName acc_params   letTupExp "withhacc_res" $ WithAcc inputs acc_lam +-- | Perform VJP on a Map.  The 'Adj' list is the adjoints of the+-- result of the map. vjpMap :: VjpOps -> [Adj] -> StmAux () -> SubExp -> Lambda SOACS -> [VName] -> ADM () vjpMap ops res_adjs _ w map_lam as   | Just res_ivs <- mapM isSparse res_adjs = returnSweepCode $ do@@ -80,7 +84,7 @@       -- (length res_ivs), there is no need for the return sweep code to       -- contain a Map at all. -      free <- filterM isActive $ namesToList $ freeIn map_lam+      free <- filterM isActive $ namesToList $ freeIn map_lam `namesSubtract` namesFromList as       free_ts <- mapM lookupType free       let adjs_for = map paramName (lambdaParams map_lam) ++ free           adjs_ts = map paramType (lambdaParams map_lam) ++ free_ts@@ -96,7 +100,13 @@             forM_ (zip as adjs_ts) $ \(a, t) -> do               scratch <- letSubExp "oo_scratch" =<< eBlank t               updateAdjIndex a (OutOfBounds, adj_i) scratch-            first subExpsRes . adjsReps <$> mapM lookupAdj as+            -- We must make sure that all free variables have the same+            -- representation in the oo-branch as in the ib-branch.+            -- In practice we do this by manifesting the adjoint.+            -- This is probably efficient, since the adjoint of a free+            -- variable is probably either a scalar or an accumulator.+            forM_ free $ \v -> insAdj v =<< adjVal =<< lookupAdj v+            first subExpsRes . adjsReps <$> mapM lookupAdj (as <> free)           inBounds res_i adj_i adj_v = subAD . buildRenamedBody $ do             forM_ (zip (lambdaParams map_lam) as) $ \(p, a) -> do               a_t <- lookupType a@@ -105,15 +115,19 @@             adj_elems <-               fmap (map resSubExp) . bodyBind . lambdaBody                 =<< vjpLambda ops (oneHot res_i (AdjVal adj_v)) adjs_for map_lam-            forM_ (zip as adj_elems) $ \(a, a_adj_elem) -> do+            let (as_adj_elems, free_adj_elems) = splitAt (length as) adj_elems+            forM_ (zip as as_adj_elems) $ \(a, a_adj_elem) ->               updateAdjIndex a (AssumeBounds, adj_i) a_adj_elem-            first subExpsRes . adjsReps <$> mapM lookupAdj as+            forM_ (zip free free_adj_elems) $ \(v, adj_se) -> do+              adj_se_v <- letExp "adj_v" (BasicOp $ SubExp adj_se)+              insAdj v adj_se_v+            first subExpsRes . adjsReps <$> mapM lookupAdj (as <> free)            -- Generate an iteration of the map function for every           -- position.  This is a bit inefficient - probably we could do           -- some deduplication.           forPos res_i (check, adj_i, adj_v) = do-            as_adj <-+            adjs <-               case check of                 CheckBounds b -> do                   (obbranch, mkadjs) <- ooBounds adj_i@@ -129,7 +143,7 @@                 OutOfBounds ->                   mapM lookupAdj as -            zipWithM setAdj as as_adj+            zipWithM setAdj (as <> free) adjs            -- Generate an iteration of the map function for every result.           forRes res_i = mapM_ (forPos res_i)

@@ -230,10 +230,12 @@ adjVal (AdjSparse sparse) = sparseArray sparse adjVal (AdjZero shape t) = zeroArray shape $ Prim t +-- | Set a specific adjoint. setAdj :: VName -> Adj -> ADM () setAdj v v_adj = modify $ \env ->   env {stateAdjs = M.insert v v_adj $ stateAdjs env} +-- | Set an 'AdjVal' adjoint.  Simple wrapper around 'setAdj'. insAdj :: VName -> VName -> ADM () insAdj v = setAdj v . AdjVal . Var @@ -471,6 +473,7 @@ isActive :: VName -> ADM Bool isActive = fmap (/= Prim Unit) . lookupType +-- | Ignore any changes to adjoints made while evaluating this action. subAD :: ADM a -> ADM a subAD m = do   old_state_adjs <- gets stateAdjs

@@ -130,9 +130,25 @@         redomapToMapAndReduce pat (w, reds, map_lam, as)       vjpStm ops mapstm $ vjpStm ops redstm m +-- Differentiating Scanomaps+vjpSOAC ops pat _aux (Screma w as form) m+  | Just (scans, map_lam) <-+      isScanomapSOAC form = do+      (mapstm, scanstm) <-+        scanomapToMapAndScan pat (w, scans, map_lam, as)+      vjpStm ops mapstm $ vjpStm ops scanstm m+ -- Differentiating Scatter-vjpSOAC ops pat aux (Scatter w lam ass written_info) m =-  vjpScatter ops pat aux (w, lam, ass, written_info) m+vjpSOAC ops pat aux (Scatter w ass lam written_info) m+  | isIdentityLambda lam =+      vjpScatter ops pat aux (w, ass, lam, written_info) m+  | otherwise = do+      map_idents <- mapM (\t -> newIdent "map_res" (arrayOfRow t w)) $ lambdaReturnType lam+      let map_stm = mkLet map_idents $ Op $ Screma w ass $ mapSOAC lam+      lam_id <- mkIdentityLambda $ lambdaReturnType lam+      let scatter_stm = Let pat aux $ Op $ Scatter w (map identName map_idents) lam_id written_info+      vjpStm ops map_stm $ vjpStm ops scatter_stm m+ -- Differentiating Histograms vjpSOAC ops pat aux (Hist n as histops f) m   | isIdentityLambda f,

@@ -272,7 +272,7 @@         pure $ Scan lam' $ scanNeutral s      splitScanRes sc res d =-      concat $ take (div d $ specialScans sc) <$> orderArgs sc res+      concatMap (take (div d $ specialScans sc)) (orderArgs sc res)  diffScanVec ::   VjpOps ->

@@ -108,7 +108,7 @@   traverse _ (ValueExp v) =     pure $ ValueExp v   traverse f (BinOpExp op x y) =-    constFoldPrimExp <$> (BinOpExp op <$> traverse f x <*> traverse f y)+    BinOpExp op <$> traverse f x <*> traverse f y   traverse f (CmpOpExp op x y) =     CmpOpExp op <$> traverse f x <*> traverse f y   traverse f (ConvOpExp op x) =@@ -230,6 +230,8 @@   | zeroIshExp y = x constFoldPrimExp (UnOpExp Abs {} x)   | not $ negativeIshExp x = x+constFoldPrimExp (UnOpExp Not {} (ValueExp (BoolValue x))) =+  ValueExp $ BoolValue $ not x constFoldPrimExp (BinOpExp UMod {} x y)   | sameIshExp x y,     IntType it <- primExpType x =@@ -244,6 +246,15 @@       ValueExp $ IntValue $ intValue it (0 :: Integer) constFoldPrimExp e = e +constFoldCmpExp :: Eq v => PrimExp v -> PrimExp v+constFoldCmpExp (CmpOpExp (CmpEq _) x y)+  | x == y =+      untyped true+constFoldCmpExp (CmpOpExp (CmpEq _) (ValueExp x) (ValueExp y))+  | x /= y =+      untyped false+constFoldCmpExp e = constFoldPrimExp e+ -- | The class of numeric types that can be used for constructing -- 'TPrimExp's. class NumExp t where@@ -465,44 +476,40 @@   sgn _ = Nothing  -- | Lifted logical conjunction.-(.&&.) :: TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v+(.&&.) :: Eq v => TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v TPrimExp x .&&. TPrimExp y = TPrimExp $ constFoldPrimExp $ BinOpExp LogAnd x y  -- | Lifted logical conjunction.-(.||.) :: TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v+(.||.) :: Eq v => TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v TPrimExp x .||. TPrimExp y = TPrimExp $ constFoldPrimExp $ BinOpExp LogOr x y  -- | Lifted relational operators; assuming signed numbers in case of -- integers.-(.<.), (.>.), (.<=.), (.>=.), (.==.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v+(.<.), (.>.), (.<=.), (.>=.), (.==.) :: Eq v => TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v TPrimExp x .<. TPrimExp y =-  TPrimExp $-    constFoldPrimExp $-      CmpOpExp cmp x y+  TPrimExp $ constFoldCmpExp $ CmpOpExp cmp x y   where     cmp = case primExpType x of       IntType t -> CmpSlt t       FloatType t -> FCmpLt t       _ -> CmpLlt TPrimExp x .<=. TPrimExp y =-  TPrimExp $-    constFoldPrimExp $-      CmpOpExp cmp x y+  TPrimExp $ constFoldCmpExp $ CmpOpExp cmp x y   where     cmp = case primExpType x of       IntType t -> CmpSle t       FloatType t -> FCmpLe t       _ -> CmpLle TPrimExp x .==. TPrimExp y =-  TPrimExp $-    constFoldPrimExp $-      CmpOpExp (CmpEq $ primExpType x `min` primExpType y) x y+  TPrimExp $ constFoldCmpExp $ CmpOpExp (CmpEq t) x y+  where+    t = primExpType x `min` primExpType y x .>. y = y .<. x x .>=. y = y .<=. x  -- | Lifted bitwise operators.  The right-shift is logical, *not* arithmetic.-(.&.), (.|.), (.^.), (.>>.), (.<<.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp t v-bitPrimExp :: (IntType -> BinOp) -> TPrimExp t v -> TPrimExp t v -> TPrimExp t v+(.&.), (.|.), (.^.), (.>>.), (.<<.) :: Eq v => TPrimExp t v -> TPrimExp t v -> TPrimExp t v+bitPrimExp :: Eq v => (IntType -> BinOp) -> TPrimExp t v -> TPrimExp t v -> TPrimExp t v bitPrimExp op (TPrimExp x) (TPrimExp y) =   TPrimExp $     constFoldPrimExp $

@@ -20,7 +20,6 @@ import System.Environment (getExecutablePath) import System.Exit import System.FilePath-import System.Process import Text.Read (readMaybe) import Text.Regex.TDFA @@ -102,10 +101,12 @@     : "-L"     : optExtraOptions opts +checkCmd :: Either CmdFailure a -> IO a+checkCmd = either (error . T.unpack . T.unlines . failureMsg) pure+ setTuningParam :: Server -> T.Text -> Int -> IO () setTuningParam server name val =-  either (error . T.unpack . T.unlines . failureMsg) (const $ pure ())-    =<< cmdSetTuningParam server name (showText val)+  void $ checkCmd =<< cmdSetTuningParam server name (showText val)  setTuningParams :: Server -> Path -> IO () setTuningParams server = mapM_ (uncurry $ setTuningParam server)@@ -225,25 +226,29 @@     t False = thresholdMax     t True = thresholdMin -thresholdForest :: FilePath -> IO ThresholdForest-thresholdForest prog = do-  thresholds <--    getThresholds . T.pack-      <$> readProcess ("." </> dropExtension prog) ["--print-params"] ""+allTuningParams :: Server -> IO [(T.Text, T.Text)]+allTuningParams server = do+  entry_points <- checkCmd =<< cmdEntryPoints server+  param_names <- concat <$> mapM (checkCmd <=< cmdTuningParams server) entry_points+  param_classes <- mapM (checkCmd <=< cmdTuningParamClass server) param_names+  pure $ zip param_names param_classes++thresholdForest :: Server -> IO ThresholdForest+thresholdForest server = do+  thresholds <- mapMaybe findThreshold <$> allTuningParams server   let root (v, _) = ((v, False), [])   pure $     unfoldForest (unfold thresholds) $       map root $         filter (null . snd) thresholds   where-    getThresholds = mapMaybe findThreshold . T.lines-    regex = makeRegex ("(.*) \\(threshold\\(([^ ]+,)(.*)\\)\\)" :: T.Text)+    regex = makeRegex ("threshold\\(([^ ]+,)(.*)\\)" :: T.Text) -    findThreshold :: T.Text -> Maybe (T.Text, [(T.Text, Bool)])-    findThreshold l = do-      [grp1, _, grp2] <- regexGroups regex l+    findThreshold :: (T.Text, T.Text) -> Maybe (T.Text, [(T.Text, Bool)])+    findThreshold (name, param_class) = do+      [_, grp] <- regexGroups regex param_class       pure-        ( grp1,+        ( name,           filter (not . T.null . fst)             $ map               ( \x ->@@ -251,7 +256,7 @@                     then (T.drop 1 x, False)                     else (x, True)               )-            $ T.words grp2+            $ T.words grp         )      unfold thresholds ((parent, parent_cmp), ancestors) =@@ -292,8 +297,9 @@       (Maybe (Int, Int), M.Map DatasetName Int) ->       (DatasetName, RunDataset, T.Text) ->       IO (Maybe (Int, Int), M.Map DatasetName Int)-    tuneDataset (thresholds, best_runtimes) (dataset_name, run, entry_point) =-      if not $ T.isPrefixOf (entry_point <> ".") v+    tuneDataset (thresholds, best_runtimes) (dataset_name, run, entry_point) = do+      relevant <- checkCmd =<< cmdTuningParams server entry_point+      if v `notElem` relevant         then do           when (optVerbose opts > 0) $             T.putStrLn $@@ -437,17 +443,15 @@   putStrLn $ "Compiling " ++ prog ++ "..."   datasets <- prepare opts futhark prog -  forest <- thresholdForest prog-  when (optVerbose opts > 0) $-    putStrLn $-      ("Threshold forest:\n" ++) $-        drawForest $-          map (fmap show) forest--  let progbin = "." </> dropExtension prog   putStrLn $ "Running with options: " ++ unwords (serverOptions opts)+  let progbin = "." </> dropExtension prog+  withServer (futharkServerCfg progbin (serverOptions opts)) $ \server -> do+    forest <- thresholdForest server+    when (optVerbose opts > 0) $+      putStrLn $+        ("Threshold forest:\n" <>) $+          drawForest (map (fmap show) forest) -  withServer (futharkServerCfg progbin (serverOptions opts)) $ \server ->     fmap fst . foldM (tuneThreshold opts server datasets) ([], mempty) $       tuningPaths forest 

@@ -198,8 +198,8 @@   V.ValueType ds t' <- toValueType t   pure $ V.ValueType (d' : ds) t'   where-    constantDim (SizeExpConst k _) = Right k-    constantDim _ = Left "Array has non-constant size."+    constantDim (SizeExp (IntLit k _ _) _) = Right $ fromInteger k+    constantDim _ = Left "Array has non-constant dimension declaration." toValueType (TEVar (QualName [] v) _)   | Just t <- lookup v m = Right $ V.ValueType [] t   where

@@ -432,7 +432,7 @@       "Disable type-checking.",     Option       []-      ["prettyString-print"]+      ["pretty-print"]       ( NoArg $           Right $ \opts ->             opts {futharkPipeline = PrettyPrint}

@@ -14,6 +14,7 @@ import Futhark.Pipeline (FutharkM, Verbosity (..), runFutharkM) import Futhark.Util (directoryContents, trim) import Futhark.Util.Options+import Language.Futhark.Semantic (mkInitialImport) import Language.Futhark.Syntax (DocComment (..), progDoc) import System.Directory (createDirectoryIfMissing) import System.Exit@@ -59,7 +60,7 @@  printDecs :: DocConfig -> FilePath -> [FilePath] -> Imports -> IO () printDecs cfg dir files imports = do-  let direct_imports = map (normalise . dropExtension) files+  let direct_imports = map (mkInitialImport . normalise . dropExtension) files       (file_htmls, _warnings) =         renderFiles direct_imports $           filter (not . ignored) imports

@@ -550,7 +550,7 @@  plottable :: CompoundValue -> Maybe [Value] plottable (ValueTuple vs) = do-  (vs', ns') <- unzip <$> mapM inspect vs+  (vs', ns') <- mapAndUnzipM inspect vs   guard $ length (nubOrd ns') == 1   Just vs'   where

@@ -23,6 +23,7 @@ import Futhark.Util.Pretty (prettyTextOneLine) import Language.Futhark.Parser.Lexer (scanTokens) import Language.Futhark.Prop (isBuiltin)+import Language.Futhark.Semantic (includeToString) import System.Environment (getExecutablePath) import System.Exit import System.FilePath@@ -36,7 +37,7 @@     [file] -> Just $ do       (_, prog_imports, _) <- readProgramOrDie file       liftIO . putStr . unlines . map (++ ".fut") . filter (not . isBuiltin) $-        map fst prog_imports+        map (includeToString . fst) prog_imports     _ -> Nothing  -- | @futhark hash@

@@ -139,7 +139,7 @@     futharkiLoaded :: Maybe FilePath   } -extendEnvs :: LoadedProg -> (T.Env, I.Ctx) -> [String] -> (T.Env, I.Ctx)+extendEnvs :: LoadedProg -> (T.Env, I.Ctx) -> [ImportName] -> (T.Env, I.Ctx) extendEnvs prog (tenv, ictx) opens = (tenv', ictx')   where     tenv' = T.envWithImports t_imports tenv@@ -245,8 +245,8 @@ onDec d = do   old_imports <- gets $ lpImports . futharkiProg   cur_import <- gets $ T.mkInitialImport . fromMaybe "." . futharkiLoaded-  let mkImport = uncurry $ T.mkImportFrom cur_import-      files = map (T.includeToFilePath . mkImport) $ decImports d+  let mkImport = T.mkImportFrom cur_import+      files = map (T.includeToFilePath . mkImport . fst) $ decImports d    cur_prog <- gets futharkiProg   imp_r <- liftIO $ extendProg cur_prog files M.empty@@ -254,13 +254,15 @@     Left e -> liftIO $ putDoc $ prettyProgErrors e     Right prog -> do       env <- gets futharkiEnv-      let (tenv, ienv) = extendEnvs prog env $ map fst $ decImports d+      let (tenv, ienv) =+            extendEnvs prog env $ map (T.mkInitialImport . fst) $ decImports d           imports = lpImports prog           src = lpNameSource prog       case T.checkDec imports src tenv cur_import d of         (_, Left e) -> liftIO $ putDoc $ T.prettyTypeErrorNoLoc e         (_, Right (tenv', d', src')) -> do-          let new_imports = filter ((`notElem` map fst old_imports) . fst) imports+          let new_imports =+                filter ((`notElem` map fst old_imports) . fst) imports           int_r <- runInterpreter $ do             let onImport ienv' (s, imp) =                   I.interpretImport ienv' (s, T.fileProg imp)

@@ -532,7 +532,7 @@         case testAction $ testCaseTest tc of           CompileTimeFailure _ -> 1           RunCases ios sts wts ->-            (length . concat) (iosTestRuns <$> ios)+            length (concatMap iosTestRuns ios)               + length sts               + length wts 

@@ -33,7 +33,7 @@ -- | Compile the program to C with calls to CUDA. compileProg :: MonadFreshNames m => T.Text -> Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg version prog = do-  (ws, Program cuda_code cuda_prelude kernels _ sizes failures prog') <-+  (ws, Program cuda_code cuda_prelude kernels _ params failures prog') <-     ImpGen.compileProg prog   let cost_centres =         [ copyDevToDev,@@ -48,12 +48,12 @@           cuda_prelude           cost_centres           kernels-          sizes           failures   (ws,)     <$> GC.compileProg       "cuda"       version+      params       operations       extra       cuda_includes@@ -277,7 +277,7 @@   args_arr <- newVName "kernel_args"   time_start <- newVName "time_start"   time_end <- newVName "time_end"-  (args', shared_vars) <- unzip <$> mapM mkArgs args+  (args', shared_vars) <- mapAndUnzipM mkArgs args   let (shared_sizes, shared_offsets) = unzip $ catMaybes shared_vars       shared_offsets_sc = mkOffsets shared_sizes       shared_args = zip shared_offsets shared_offsets_sc

@@ -19,7 +19,6 @@     copyScalarToDev,     costCentreReport,     failureMsgFunction,-    generateTuningParams,     kernelRuns,     kernelRuntime,   )@@ -89,10 +88,9 @@   T.Text ->   [Name] ->   M.Map KernelName KernelSafety ->-  M.Map Name SizeClass ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-generateBoilerplate cuda_program cuda_prelude cost_centres kernels sizes failures = do+generateBoilerplate cuda_program cuda_prelude cost_centres kernels failures = do   let cuda_program_fragments =         -- Some C compilers limit the size of literal strings, so         -- chunk the entire program into small bits here, and@@ -100,7 +98,6 @@         [[C.cinit|$string:s|] | s <- chunk 2000 $ T.unpack $ cuda_prelude <> cuda_program]       program_fragments = cuda_program_fragments ++ [[C.cinit|NULL|]]   let max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures-  generateTuningParams sizes   mapM_     GC.earlyDecl     [C.cunit|static const int max_failure_args = $int:max_failure_args;@@ -122,17 +119,6 @@   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_tile_size(struct futhark_context_config *cfg, int size);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_reg_tile_size(struct futhark_context_config *cfg, int size);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]--  let set_tuning_params =-        zipWith-          (\i k -> [C.cstm|ctx->tuning_params.$id:k = &ctx->cfg->tuning_params[$int:i];|])-          [(0 :: Int) ..]-          $ M.keys sizes--  GC.earlyDecl-    [C.cedecl|static void set_tuning_params(struct futhark_context* ctx) {-                $stms:set_tuning_params-              }|]    GC.generateProgramStruct 

@@ -38,7 +38,7 @@       opencl_prelude       kernels       types-      sizes+      params       failures       prog'     ) <-@@ -54,6 +54,7 @@     <$> GC.compileProg       "opencl"       version+      params       operations       ( generateBoilerplate           opencl_code@@ -61,7 +62,6 @@           cost_centres           kernels           types-          sizes           failures       )       include_opencl_h

@@ -14,13 +14,11 @@     kernelRuntime,     kernelRuns,     sizeLoggingCode,-    generateTuningParams,   ) where  import Control.Monad.State import Data.Map qualified as M-import Data.Maybe import Data.Text qualified as T import Futhark.CodeGen.Backends.GenericC qualified as GC import Futhark.CodeGen.Backends.GenericC.Options@@ -28,7 +26,7 @@ import Futhark.CodeGen.ImpCode.OpenCL import Futhark.CodeGen.OpenCL.Heuristics import Futhark.CodeGen.RTS.C (backendsOpenclH)-import Futhark.Util (chunk, zEncodeText)+import Futhark.Util (chunk) import Futhark.Util.Pretty (prettyTextOneLine) import Language.C.Quote.OpenCL qualified as C import Language.C.Syntax qualified as C@@ -70,23 +68,6 @@                              &ctx->program->$id:(kernelRuns name),                              &ctx->program->$id:(kernelRuntime name))|] -generateTuningParams :: M.Map Name SizeClass -> GC.CompilerM op a ()-generateTuningParams sizes = do-  let strinit s = [C.cinit|$string:(T.unpack s)|]-      intinit x = [C.cinit|$int:x|]-      size_name_inits = map (strinit . prettyText) $ M.keys sizes-      size_var_inits = map (strinit . zEncodeText . prettyText) $ M.keys sizes-      size_class_inits = map (strinit . prettyText) $ M.elems sizes-      size_default_inits = map (intinit . fromMaybe 0 . sizeDefault) $ M.elems sizes-      size_decls = map (\k -> [C.csdecl|typename int64_t *$id:k;|]) $ M.keys sizes-      num_sizes = length sizes-  GC.earlyDecl [C.cedecl|struct tuning_params { $sdecls:size_decls };|]-  GC.earlyDecl [C.cedecl|static const int num_tuning_params = $int:num_sizes;|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_names[] = { $inits:size_name_inits };|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_vars[] = { $inits:size_var_inits };|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_classes[] = { $inits:size_class_inits };|]-  GC.earlyDecl [C.cedecl|static typename int64_t tuning_param_defaults[] = { $inits:size_default_inits };|]- releaseKernel :: (KernelName, KernelSafety) -> C.Stm releaseKernel (name, _) = [C.cstm|OPENCL_SUCCEED_FATAL(clReleaseKernel(ctx->program->$id:name));|] @@ -148,10 +129,9 @@   [Name] ->   M.Map KernelName KernelSafety ->   [PrimType] ->-  M.Map Name SizeClass ->   [FailureMsg] ->   GC.CompilerM OpenCL () ()-generateBoilerplate opencl_program opencl_prelude cost_centres kernels types sizes failures = do+generateBoilerplate opencl_program opencl_prelude cost_centres kernels types failures = do   let opencl_program_fragments =         -- Some C compilers limit the size of literal strings, so         -- chunk the entire program into small bits here, and@@ -162,7 +142,6 @@         | FloatType Float64 `elem` types = [C.cexp|1|]         | otherwise = [C.cexp|0|]       max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures-  generateTuningParams sizes   mapM_     GC.earlyDecl     [C.cunit|static const int max_failure_args = $int:max_failure_args;@@ -190,17 +169,6 @@   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_default_threshold(struct futhark_context_config *cfg, int size);|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_command_queue(struct futhark_context_config *cfg, typename cl_command_queue);|]   GC.headerDecl GC.MiscDecl [C.cedecl|typename cl_command_queue futhark_context_get_command_queue(struct futhark_context* ctx);|]--  let set_tuning_params =-        zipWith-          (\i k -> [C.cstm|ctx->tuning_params.$id:k = &ctx->cfg->tuning_params[$int:i];|])-          [(0 :: Int) ..]-          $ M.keys sizes--  GC.earlyDecl-    [C.cedecl|static void set_tuning_params(struct futhark_context* ctx) {-                $stms:set_tuning_params-              }|]    GC.generateProgramStruct 

@@ -7,6 +7,7 @@   ( compileProg,     compileProg',     defaultOperations,+    ParamMap,     CParts (..),     asLibrary,     asExecutable,@@ -18,10 +19,12 @@  import Control.Monad.Reader import Control.Monad.State+import Data.Bifunctor (second) import Data.DList qualified as DL import Data.Loc import Data.Map.Strict qualified as M import Data.Maybe+import Data.Set qualified as S import Data.Text qualified as T import Futhark.CodeGen.Backends.GenericC.CLI (cliDefs) import Futhark.CodeGen.Backends.GenericC.Code@@ -34,8 +37,10 @@ import Futhark.CodeGen.Backends.GenericC.Types import Futhark.CodeGen.ImpCode import Futhark.CodeGen.RTS.C (cacheH, contextH, contextPrototypesH, errorsH, freeListH, halfH, lockH, timingH, utilH)+import Futhark.IR.GPU.Sizes import Futhark.Manifest qualified as Manifest import Futhark.MonadFreshNames+import Futhark.Util (zEncodeText) import Language.C.Quote.OpenCL qualified as C import Language.C.Syntax qualified as C import NeatInterpolation (untrimming)@@ -320,10 +325,15 @@ asServer parts =   gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cServer parts <> cLib parts +relevantParams :: Name -> ParamMap -> [Name]+relevantParams fname m =+  map fst $ filter ((fname `S.member`) . snd . snd) $ M.toList m+ compileProg' ::   MonadFreshNames m =>   T.Text ->   T.Text ->+  ParamMap ->   Operations op s ->   s ->   CompilerM op s () ->@@ -332,7 +342,7 @@   [Option] ->   Definitions op ->   m (CParts, CompilerState s)-compileProg' backend version ops def extra header_extra (arr_space, spaces) options prog = do+compileProg' backend version params ops def extra header_extra (arr_space, spaces) options prog = do   src <- getNameSource   let ((prototypes, definitions, entry_point_decls, manifest), endstate) =         runCompilerM ops src def compileProgAction@@ -451,17 +461,18 @@     Definitions types consts (Functions funs) = prog      compileProgAction = do-      (memfuns, memreport) <- unzip <$> mapM defineMemorySpace spaces+      (memfuns, memreport) <- mapAndUnzipM defineMemorySpace spaces        get_consts <- compileConstants consts        ctx_ty <- contextType        (prototypes, functions) <--        unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs+        mapAndUnzipM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs        (entry_points, entry_points_manifest) <--        unzip . catMaybes <$> mapM (uncurry (onEntryPoint get_consts)) funs+        fmap (unzip . catMaybes) $ forM funs $ \(fname, fun) ->+          onEntryPoint get_consts (relevantParams fname params) fname fun        headerDecl InitDecl [C.cedecl|struct futhark_context_config;|]       headerDecl InitDecl [C.cedecl|struct futhark_context_config* futhark_context_config_new(void);|]@@ -473,10 +484,20 @@       headerDecl InitDecl [C.cedecl|void futhark_context_free(struct futhark_context* cfg);|]       headerDecl MiscDecl [C.cedecl|int futhark_context_sync(struct futhark_context* ctx);|] +      generateTuningParams params       extra -      mapM_ earlyDecl $ concat memfuns+      let set_tuning_params =+            zipWith+              (\i k -> [C.cstm|ctx->tuning_params.$id:k = &ctx->cfg->tuning_params[$int:i];|])+              [(0 :: Int) ..]+              $ M.keys params+      earlyDecl+        [C.cedecl|static void set_tuning_params(struct futhark_context* ctx) {+                $stms:set_tuning_params+              }|] +      mapM_ earlyDecl $ concat memfuns       type_funs <- generateAPITypes arr_space types       generateCommonLibFuns memreport @@ -493,6 +514,7 @@   MonadFreshNames m =>   T.Text ->   T.Text ->+  ParamMap ->   Operations op () ->   CompilerM op () () ->   T.Text ->@@ -500,8 +522,27 @@   [Option] ->   Definitions op ->   m CParts-compileProg backend version ops extra header_extra (arr_space, spaces) options prog =-  fst <$> compileProg' backend version ops () extra header_extra (arr_space, spaces) options prog+compileProg backend version params ops extra header_extra (arr_space, spaces) options prog =+  fst <$> compileProg' backend version params ops () extra header_extra (arr_space, spaces) options prog++generateTuningParams :: ParamMap -> CompilerM op a ()+generateTuningParams params = do+  let (param_names, (param_classes, _param_users)) =+        second unzip $ unzip $ M.toList params+      strinit s = [C.cinit|$string:(T.unpack s)|]+      intinit x = [C.cinit|$int:x|]+      size_name_inits = map (strinit . prettyText) param_names+      size_var_inits = map (strinit . zEncodeText . prettyText) param_names+      size_class_inits = map (strinit . prettyText) param_classes+      size_default_inits = map (intinit . fromMaybe 0 . sizeDefault) param_classes+      size_decls = map (\k -> [C.csdecl|typename int64_t *$id:k;|]) param_names+      num_params = length params+  earlyDecl [C.cedecl|struct tuning_params { $sdecls:size_decls };|]+  earlyDecl [C.cedecl|static const int num_tuning_params = $int:num_params;|]+  earlyDecl [C.cedecl|static const char *tuning_param_names[] = { $inits:size_name_inits };|]+  earlyDecl [C.cedecl|static const char *tuning_param_vars[] = { $inits:size_var_inits };|]+  earlyDecl [C.cedecl|static const char *tuning_param_classes[] = { $inits:size_class_inits };|]+  earlyDecl [C.cedecl|static typename int64_t tuning_param_defaults[] = { $inits:size_default_inits };|]  generateCommonLibFuns :: [C.BlockItem] -> CompilerM op s () generateCommonLibFuns memreport = do

@@ -296,7 +296,7 @@  cliEntryPoint ::   Manifest -> T.Text -> EntryPoint -> (C.Definition, C.Initializer)-cliEntryPoint manifest entry_point_name (EntryPoint cfun outputs inputs) =+cliEntryPoint manifest entry_point_name (EntryPoint cfun _tuning_params outputs inputs) =   let (input_items, pack_input, free_input, free_parsed, input_args) =         unzip5 $ readInputs manifest $ map inputType inputs 

@@ -36,7 +36,7 @@       onPart (ErrorVal (FloatType Float16) x) = ("%f",) . asDouble <$> compileExp x       onPart (ErrorVal (FloatType Float32) x) = ("%f",) . asDouble <$> compileExp x       onPart (ErrorVal (FloatType Float64) x) = ("%f",) <$> compileExp x-  (formatstrs, formatargs) <- unzip <$> mapM onPart parts+  (formatstrs, formatargs) <- mapAndUnzipM onPart parts   pure (mconcat formatstrs, formatargs)  compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName

@@ -133,13 +133,17 @@ entryName :: Name -> T.Text entryName = ("entry_" <>) . escapeName . nameToText +tuningParamsName :: Name -> T.Text+tuningParamsName = ("tuning_params_for_" <>) . escapeName . nameToText+ onEntryPoint ::   [C.BlockItem] ->+  [Name] ->   Name ->   Function op ->   CompilerM op s (Maybe (C.Definition, (T.Text, Manifest.EntryPoint)))-onEntryPoint _ _ (Function Nothing _ _ _) = pure Nothing-onEntryPoint get_consts fname (Function (Just (EntryPoint ename results args)) outputs inputs _) = inNewFunction $ do+onEntryPoint _ _ _ (Function Nothing _ _ _) = pure Nothing+onEntryPoint get_consts relevant_params fname (Function (Just (EntryPoint ename results args)) outputs inputs _) = inNewFunction $ do   let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs       in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs @@ -148,6 +152,7 @@   decl_mem <- declAllocatedMem    entry_point_function_name <- publicName $ entryName ename+  tuning_params_name <- publicName $ tuningParamsName ename    (inputs', unpack_entry_inputs) <- prepareEntryInputs $ map snd args   let (entry_point_input_params, entry_point_input_checks) = unzip inputs'@@ -164,6 +169,10 @@                                       $params:entry_point_output_params,                                       $params:entry_point_input_params);|] +  headerDecl+    MiscDecl+    [C.cedecl|const int* $id:tuning_params_name(void);|]+   let checks = catMaybes entry_point_input_checks       check_input =         if null checks@@ -206,11 +215,13 @@          $items:(criticalSection ops critical)           return ret;-       }|]+       }+       |]        manifest =         Manifest.EntryPoint           { Manifest.entryPointCFun = entry_point_function_name,+            Manifest.entryPointTuningParams = map nameToText relevant_params,             -- Note that our convention about what is "input/output"             -- and what is "results/args" is different between the             -- manifest and ImpCode.@@ -224,7 +235,7 @@       declMem name space     stubParam (ScalarParam name ty) = do       let ty' = primTypeToCType ty-      decl [C.cdecl|$ty:ty' $id:name;|]+      decl [C.cdecl|$ty:ty' $id:name = $exp:(blankPrimValue ty);|]      vdType (TransparentValue (ScalarValue pt signed _)) =       prettySigned (signed == Unsigned) pt

@@ -56,7 +56,7 @@  compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) compileFun get_constants extra (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do-  (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs+  (outparams, out_ptrs) <- mapAndUnzipM compileOutput outputs   inparams <- mapM compileInput inputs    cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do@@ -91,7 +91,7 @@ -- memory non-lexxical or do anything fancy). compileVoidFun :: [C.BlockItem] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) compileVoidFun get_constants (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do-  (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs+  (outparams, out_ptrs) <- mapAndUnzipM compileOutput outputs   inparams <- mapM compileInput inputs    cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do

@@ -212,7 +212,7 @@     manifest  oneEntryBoilerplate :: Manifest -> (T.Text, EntryPoint) -> ([C.Definition], C.Initializer)-oneEntryBoilerplate manifest (name, EntryPoint cfun outputs inputs) =+oneEntryBoilerplate manifest (name, EntryPoint cfun tuning_params outputs inputs) =   let call_f = "call_" <> nameFromText name       out_types = map outputType outputs       in_types = map inputType inputs@@ -220,6 +220,7 @@       in_types_name = nameFromText name <> "_in_types"       out_unique_name = nameFromText name <> "_out_unique"       in_unique_name = nameFromText name <> "_in_unique"+      tuning_params_name = nameFromText name <> "_tuning_params"       (out_items, out_args)         | null out_types = ([C.citems|(void)outs;|], mempty)         | otherwise = unzip $ zipWith loadOut [0 ..] out_types@@ -241,6 +242,10 @@                 bool $id:in_unique_name[] = {                   $inits:(map inputUniqueInit inputs)                 };+                const char* $id:tuning_params_name[] = {+                  $inits:(map textInit tuning_params),+                  NULL+                };                 int $id:call_f(struct futhark_context *ctx, void **outs, void **ins) {                   $items:out_items                   $items:in_items@@ -250,6 +255,7 @@         [C.cinit|{             .name = $string:(T.unpack name),             .f = $id:call_f,+            .tuning_params = $id:tuning_params_name,             .in_types = $id:in_types_name,             .out_types = $id:out_types_name,             .in_unique = $id:in_unique_name,@@ -273,6 +279,8 @@        in ( [C.citem|$ty:(cType manifest tname) $id:v = *($ty:(cType manifest tname)*)ins[$int:i];|],             [C.cexp|$id:v|]           )++    textInit t = [C.cinit|$string:(T.unpack t)|]  entryBoilerplate :: Manifest -> ([C.Definition], [C.Initializer]) entryBoilerplate manifest =

@@ -1171,7 +1171,7 @@       onPart (Imp.ErrorVal FloatType {} x) = ("%f",) <$> compileExp x       onPart (Imp.ErrorVal Imp.Bool x) = ("%r",) <$> compileExp x       onPart (Imp.ErrorVal Unit {} x) = ("%r",) <$> compileExp x-  (formatstrs, formatargs) <- unzip <$> mapM onPart parts+  (formatstrs, formatargs) <- mapAndUnzipM onPart parts   pure (mconcat formatstrs, formatargs)  compileCode :: Imp.Code op -> CompilerM op s ()

@@ -51,6 +51,7 @@     ( GC.compileProg         "multicore"         version+        mempty         operations         generateBoilerplate         ""@@ -452,12 +453,10 @@       GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]       pure mempty -  free_all_mem <- GC.freeAllocatedMem   let ftask_err = fpar_task <> "_err"       code =         [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name);                   if ($id:ftask_err != 0) {-                    $items:free_all_mem;                     err = $id:ftask_err;                     goto cleanup;                   }|]

@@ -11,11 +11,6 @@ -- | Generate the necessary boilerplate. generateBoilerplate :: GC.CompilerM op s () generateBoilerplate = do-  GC.earlyDecl [C.cedecl|static const int num_tuning_params = 0;|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_names[1];|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_vars[1];|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_classes[1];|]-  GC.earlyDecl [C.cedecl|static typename int64_t *tuning_param_defaults[1];|]   mapM_ GC.earlyDecl [C.cunit|$esc:(T.unpack schedulerH)|]   mapM_ GC.earlyDecl [C.cunit|$esc:(T.unpack backendsMulticoreH)|]   GC.headerDecl GC.InitDecl [C.cedecl|void futhark_context_config_set_num_threads(struct futhark_context_config *cfg, int n);|]

@@ -72,6 +72,7 @@       ( GC.compileProg'           "ispc"           version+          mempty           operations           (ISPCState mempty mempty)           ( do@@ -229,22 +230,21 @@           extra_c = [[C.cparam|struct futhark_context * ctx|]]           extra_exp = [[C.cexp|$id:p|] | C.Param (Just p) _ _ _ <- extra] -      (inparams_c, in_args_c) <- unzip <$> mapM (compileInputsExtern []) inputs-      (outparams_c, out_args_c) <- unzip <$> mapM (compileOutputsExtern []) outputs+      (inparams_c, in_args_c) <- mapAndUnzipM (compileInputsExtern []) inputs+      (outparams_c, out_args_c) <- mapAndUnzipM (compileOutputsExtern []) outputs -      (inparams_extern, _) <- unzip <$> mapM (compileInputsExtern [C.ctyquals|$tyqual:uniform|]) inputs-      (outparams_extern, _) <- unzip <$> mapM (compileOutputsExtern [C.ctyquals|$tyqual:uniform|]) outputs+      (inparams_extern, _) <- mapAndUnzipM (compileInputsExtern [C.ctyquals|$tyqual:uniform|]) inputs+      (outparams_extern, _) <- mapAndUnzipM (compileOutputsExtern [C.ctyquals|$tyqual:uniform|]) outputs -      (inparams_uni, in_args_noderef) <- unzip <$> mapM compileInputsUniform inputs-      (outparams_uni, out_args_noderef) <- unzip <$> mapM compileOutputsUniform outputs+      (inparams_uni, in_args_noderef) <- mapAndUnzipM compileInputsUniform inputs+      (outparams_uni, out_args_noderef) <- mapAndUnzipM compileOutputsUniform outputs        (inparams_varying, in_args_vary, prebody_in') <- unzip3 <$> mapM compileInputsVarying inputs       (outparams_varying, out_args_vary, prebody_out', postbody_out') <- unzip4 <$> mapM compileOutputsVarying outputs       let (prebody_in, prebody_out, postbody_out) = over each concat (prebody_in', prebody_out', postbody_out')        GC.libDecl-        =<< pure-          [C.cedecl|int $id:(funName fname <> "_extern")($params:extra_c, $params:outparams_c, $params:inparams_c) {+        [C.cedecl|int $id:(funName fname <> "_extern")($params:extra_c, $params:outparams_c, $params:inparams_c) {                   return $id:(funName fname)($args:extra_exp, $args:out_args_c, $args:in_args_c);                 }|] 

@@ -48,6 +48,7 @@     GC.compileProg       "wasm_multicore"       version+      mempty       MC.operations       generateBoilerplate       ""

@@ -15,8 +15,8 @@   ( ErrorMsg (..),     ErrorMsgPart (..),     FailureMsg (..),+    ParamMap,     PrimType (..),-    SizeClass (..),     errorMsgArgTypes,     sizeDefault,     untyped,@@ -31,7 +31,7 @@ -- | Python code (as a string) that calls the -- @initiatialize_opencl_object@ procedure.  Should be put in the -- class constructor.-openClInit :: [PrimType] -> String -> M.Map Name SizeClass -> [FailureMsg] -> T.Text+openClInit :: [PrimType] -> String -> ParamMap -> [FailureMsg] -> T.Text openClInit types assign sizes failures =   [text| size_heuristics=$size_heuristics@@ -76,10 +76,10 @@     onPart (ErrorString s) = formatEscape $ T.unpack s     onPart ErrorVal {} = "{}" -sizeClassesToPython :: M.Map Name SizeClass -> PyExp+sizeClassesToPython :: ParamMap -> PyExp sizeClassesToPython = Dict . map f . M.toList   where-    f (size_name, size_class) =+    f (size_name, (size_class, _)) =       ( String $ prettyText size_name,         Dict           [ (String "class", String $ prettyText size_class),

@@ -23,7 +23,16 @@ compileProg :: MonadFreshNames m => T.Text -> Prog SeqMem -> m (ImpGen.Warnings, GC.CParts) compileProg version =   traverse-    (GC.compileProg "c" version operations generateBoilerplate mempty (DefaultSpace, [DefaultSpace]) [])+    ( GC.compileProg+        "c"+        version+        mempty+        operations+        generateBoilerplate+        mempty+        (DefaultSpace, [DefaultSpace])+        []+    )     <=< ImpGen.compileProg   where     operations :: GC.Operations Imp.Sequential ()

@@ -11,11 +11,6 @@ -- | Generate the necessary boilerplate. generateBoilerplate :: GC.CompilerM op s () generateBoilerplate = do-  GC.earlyDecl [C.cedecl|static const int num_tuning_params = 0;|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_names[1];|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_vars[1];|]-  GC.earlyDecl [C.cedecl|static const char *tuning_param_classes[1];|]-  GC.earlyDecl [C.cedecl|static typename int64_t *tuning_param_defaults[1];|]   GC.earlyDecl [C.cedecl|$esc:(T.unpack backendsCH)|]   GC.generateProgramStruct {-# NOINLINE generateBoilerplate #-}

@@ -43,6 +43,7 @@     GC.compileProg       "wasm"       version+      mempty       operations       generateBoilerplate       ""

@@ -75,6 +75,7 @@     lexicalMemoryUsage,     calledFuncs,     callGraph,+    ParamMap,      -- * Typed enumerations     Bytes,@@ -105,7 +106,7 @@ import Data.Traversable import Futhark.Analysis.PrimExp import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.GPU.Sizes (Count (..))+import Futhark.IR.GPU.Sizes (Count (..), SizeClass (..)) import Futhark.IR.Pretty () import Futhark.IR.Prop.Names import Futhark.IR.Syntax.Core@@ -405,6 +406,10 @@       let grow v = maybe (S.singleton v) (S.insert v) (M.lookup v cur)           next = M.map (foldMap grow) cur        in if next == cur then cur else loop next++-- | A mapping from names of tuning parameters to their class, as well+-- as which functions make use of them (including transitively).+type ParamMap = M.Map Name (SizeClass, S.Set Name)  -- | A side-effect free expression whose execution will produce a -- single primitive value.

@@ -39,7 +39,7 @@     -- | So we can detect whether the device is capable.     openClUsedTypes :: [PrimType],     -- | Runtime-configurable constants.-    openClSizes :: M.Map Name SizeClass,+    openClParams :: ParamMap,     -- | Assertion failure error messages.     openClFailures :: [FailureMsg],     hostDefinitions :: Definitions OpenCL

@@ -657,7 +657,7 @@   Maybe [EntryResult] ->   ImpM rep r op (Maybe [(Uniqueness, Imp.ExternalValue)], [Imp.Param], [ValueDestination]) compileOutParams types orig_rts maybe_orig_epts = do-  (maybe_params, dests) <- unzip <$> mapM compileOutParam orig_rts+  (maybe_params, dests) <- mapAndUnzipM compileOutParam orig_rts   evs <- case maybe_orig_epts of     Just orig_epts ->       Just <$> compileExternalValues types orig_rts orig_epts maybe_params

@@ -256,7 +256,7 @@           groupLoop iterations $ \i -> do             dIndexSpace (zip ltids dims') $ sExt64 i             m-        Just num_chunks -> do+        Just num_chunks -> localOps threadOperations $ do           let ltid = kernelLocalThreadId constants           sFor "chunk_i" num_chunks $ \chunk_i -> do             i <- dPrimVE "i" $ chunk_i * sExt32 group_size + ltid@@ -286,11 +286,13 @@   updateAcc acc is vs compileGroupExp (Pat [dest]) (BasicOp (Replicate ds se)) = do   flat <- newVName "rep_flat"-  is <- replicateM (shapeRank ds) (newVName "rep_i")+  is <- replicateM (arrayRank dest_t) (newVName "rep_i")   let is' = map le64 is-  groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ shapeDims ds) $-    copyDWIMFix (patElemName dest) is' se []+  groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ arrayDims dest_t) $+    copyDWIMFix (patElemName dest) is' se (drop (shapeRank ds) is')   sOp $ Imp.Barrier Imp.FenceLocal+  where+    dest_t = patElemType dest compileGroupExp (Pat [dest]) (BasicOp (Rotate rs arr)) = do   ds <- map pe64 . arrayDims <$> lookupType arr   groupCoverSpace ds $ \is -> do@@ -682,19 +684,25 @@ segOpSizes :: Stms GPUMem -> SegOpSizes segOpSizes = onStms   where-    onStms = foldMap (onExp . stmExp)-    onExp (Op (Inner (SegOp op))) =+    onStms = foldMap onStm+    onStm (Let _ _ (Op (Inner (SegOp op)))) =       case segVirt $ segLevel op of         SegNoVirtFull seq_dims ->           S.singleton $ map snd $ snd $ partitionSeqDims seq_dims $ segSpace op         _ -> S.singleton $ map snd $ unSegSpace $ segSpace op-    onExp (BasicOp (Replicate shape _)) =-      S.singleton $ shapeDims shape-    onExp (Match _ cases defbody _) =+    onStm (Let (Pat [pe]) _ (BasicOp (Replicate {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let (Pat [pe]) _ (BasicOp (Iota {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let (Pat [pe]) _ (BasicOp (Copy {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let (Pat [pe]) _ (BasicOp (Manifest {}))) =+      S.singleton $ arrayDims $ patElemType pe+    onStm (Let _ _ (Match _ cases defbody _)) =       foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody)-    onExp (DoLoop _ _ body) =+    onStm (Let _ _ (DoLoop _ _ body)) =       onStms (bodyStms body)-    onExp _ = mempty+    onStm _ = mempty  -- | Precompute various constants and useful information. precomputeConstants :: Count GroupSize (Imp.TExp Int64) -> Stms GPUMem -> CallKernelGen Precomputed

@@ -12,6 +12,7 @@ import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor (second)+import Data.Foldable (toList) import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S@@ -45,10 +46,11 @@   ImpGPU.Program ->   ImpOpenCL.Program translateGPU target prog =-  let ( prog',+  let env = envFromProg prog+      ( prog',         ToOpenCL kernels device_funs used_types sizes failures         ) =-          (`runState` initialOpenCL) . (`runReaderT` envFromProg prog) $ do+          (`runState` initialOpenCL) . (`runReaderT` env) $ do             let ImpGPU.Definitions                   types                   (ImpGPU.Constants ps consts)@@ -78,7 +80,7 @@         opencl_prelude         kernels'         (S.toList used_types)-        (cleanSizes sizes)+        (findParamUsers env prog' (cleanSizes sizes))         failures         prog'   where@@ -96,6 +98,31 @@       SizeThreshold (filter ((`elem` known) . fst) path) def     clean s = s +findParamUsers ::+  Env ->+  Definitions ImpOpenCL.OpenCL ->+  M.Map Name SizeClass ->+  ParamMap+findParamUsers env defs = M.mapWithKey onParam+  where+    cg = envCallGraph env++    getSize (ImpOpenCL.GetSize _ v) = Just v+    getSize (ImpOpenCL.CmpSizeLe _ v _) = Just v+    getSize (ImpOpenCL.GetSizeMax {}) = Nothing+    getSize (ImpOpenCL.LaunchKernel {}) = Nothing+    directUseInFun fun = mapMaybe getSize $ toList $ functionBody fun+    direct_uses = map (second directUseInFun) $ unFunctions $ defFuns defs++    calledBy fname = M.findWithDefault mempty fname cg+    indirectUseInFun fname =+      ( fname,+        foldMap snd $ filter ((`S.member` calledBy fname) . fst) direct_uses+      )+    indirect_uses = direct_uses <> map (indirectUseInFun . fst) direct_uses++    onParam k c = (c, S.fromList $ map fst $ filter ((k `elem`) . snd) indirect_uses)+ pointerQuals :: String -> [C.TypeQual] pointerQuals "global" = [C.ctyquals|__global|] pointerQuals "local" = [C.ctyquals|__local|]@@ -142,7 +169,8 @@  data Env = Env   { envFuns :: ImpGPU.Functions ImpGPU.HostOp,-    envFunsMayFail :: S.Set Name+    envFunsMayFail :: S.Set Name,+    envCallGraph :: M.Map Name (S.Set Name)   }  codeMayFail :: (a -> Bool) -> ImpGPU.Code a -> Bool@@ -172,7 +200,7 @@       any (`elem` base_mayfail) $ fname : S.toList (M.findWithDefault mempty fname cg)  envFromProg :: ImpGPU.Program -> Env-envFromProg prog = Env funs (funsMayFail cg funs)+envFromProg prog = Env funs (funsMayFail cg funs) cg   where     funs = defFuns prog     cg = ImpGPU.callGraph calledInHostOp funs

@@ -88,7 +88,7 @@  newtype UncheckedImport = UncheckedImport   { unChecked ::-      WithErrors (LoadedFile E.UncheckedProg, [(ImportName, MVar UncheckedImport)])+      WithErrors (LoadedFile E.UncheckedProg, [((ImportName, Loc), MVar UncheckedImport)])   }  -- | If mapped to Nothing, treat it as present.  This is used when@@ -99,14 +99,14 @@ newState known = newMVar $ M.fromList $ zip known $ repeat Nothing  orderedImports ::-  [(ImportName, MVar UncheckedImport)] ->+  [((ImportName, Loc), MVar UncheckedImport)] ->   IO [(ImportName, WithErrors (LoadedFile E.UncheckedProg))] orderedImports = fmap reverse . flip execStateT [] . mapM_ (spelunk [])   where-    spelunk steps (include, mvar)+    spelunk steps ((include, loc), mvar)       | include `elem` steps = do           let problem =-                ProgError (locOf include) . pretty $+                ProgError loc . pretty $                   "Import cycle: "                     <> intercalate                       " -> "@@ -150,8 +150,8 @@       now <- getCurrentTime       pure $ Just $ Right (file_contents, now) -readImportFile :: ImportName -> VFS -> IO (Either ProgError (LoadedFile T.Text))-readImportFile include vfs = do+readImportFile :: ImportName -> Loc -> VFS -> IO (Either ProgError (LoadedFile T.Text))+readImportFile include loc vfs = do   -- First we try to find a file of the given name in the search path,   -- then we look at the builtin library if we have to.  For the   -- builtins, we don't use the search path.@@ -161,11 +161,11 @@     (Just (Right (s, mod_time)), _) ->       pure $ Right $ loaded filepath s mod_time     (Just (Left e), _) ->-      pure $ Left $ ProgError (locOf include) $ pretty e+      pure $ Left $ ProgError loc $ pretty e     (Nothing, Just s) ->       pure $ Right $ loaded prelude_str s startupTime     (Nothing, Nothing) ->-      pure $ Left $ ProgError (locOf include) $ pretty not_found+      pure $ Left $ ProgError loc $ pretty not_found   where     prelude_str = "/" Posix.</> includeToString include Posix.<.> "fut" @@ -186,10 +186,10 @@     Left (SyntaxError loc err) ->       pure . UncheckedImport . Left . NE.singleton $ ProgError loc $ pretty err     Right prog -> do-      let imports = map (uncurry (mkImportFrom import_name)) $ E.progImports prog+      let imports = map (first $ mkImportFrom import_name) $ E.progImports prog       mvars <-         mapMaybe sequenceA . zip imports-          <$> mapM (readImport state_mvar vfs) imports+          <$> mapM (uncurry $ readImport state_mvar vfs) imports       let file =             LoadedFile               { lfPath = file_name,@@ -199,14 +199,14 @@               }       pure $ UncheckedImport $ Right (file, mvars) -readImport :: ReaderState -> VFS -> ImportName -> IO (Maybe (MVar UncheckedImport))-readImport state_mvar vfs include =+readImport :: ReaderState -> VFS -> ImportName -> Loc -> IO (Maybe (MVar UncheckedImport))+readImport state_mvar vfs include loc =   modifyMVar state_mvar $ \state ->     case M.lookup include state of       Just x -> pure (state, x)       Nothing -> do         prog_mvar <- newImportMVar $ do-          readImportFile include vfs >>= \case+          readImportFile include loc vfs >>= \case             Left e -> pure $ UncheckedImport $ Left $ NE.singleton e             Right file -> handleFile state_mvar vfs file         pure (M.insert include prog_mvar state, prog_mvar)@@ -219,16 +219,16 @@ readUntypedLibraryExceptKnown known vfs fps = do   state_mvar <- liftIO $ newState known   let prelude_import = mkInitialImport "/prelude/prelude"-  prelude_mvar <- liftIO $ readImport state_mvar vfs prelude_import+  prelude_mvar <- liftIO $ readImport state_mvar vfs prelude_import mempty   fps_mvars <- liftIO (mapM (onFile state_mvar) fps)-  let unknown_mvars = onlyUnknown ((prelude_import, prelude_mvar) : fps_mvars)+  let unknown_mvars = onlyUnknown (((prelude_import, mempty), prelude_mvar) : fps_mvars)   fmap (map snd) . errorsToTop <$> orderedImports unknown_mvars   where     onlyUnknown = mapMaybe sequenceA     onFile state_mvar fp =       modifyMVar state_mvar $ \state -> do         case M.lookup include state of-          Just prog_mvar -> pure (state, (include, prog_mvar))+          Just prog_mvar -> pure (state, ((include, mempty), prog_mvar))           Nothing -> do             prog_mvar <- newImportMVar $ do               if takeExtension fp /= ".fut"@@ -256,7 +256,7 @@                       pure . UncheckedImport . Left . NE.singleton $                         ProgError NoLoc $                           pretty fp <> ": file not found."-            pure (M.insert include prog_mvar state, (include, prog_mvar))+            pure (M.insert include prog_mvar state, ((include, mempty), prog_mvar))       where         include = mkInitialImport fp_name         (fp_name, _) = Posix.splitExtension fp@@ -274,7 +274,7 @@ asImports :: [LoadedFile CheckedFile] -> Imports asImports = map f   where-    f lf = (includeToString (lfImportName lf), cfMod $ lfMod lf)+    f lf = (lfImportName lf, cfMod $ lfMod lf)  typeCheckProg ::   [LoadedFile CheckedFile] ->@@ -353,7 +353,7 @@ lpImports :: LoadedProg -> Imports lpImports = map f . lpFiles   where-    f lf = (includeToString (lfImportName lf), cfMod $ lfMod lf)+    f lf = (lfImportName lf, cfMod $ lfMod lf)  -- | All warnings of a 'LoadedProg'. lpWarnings :: LoadedProg -> Warnings

@@ -117,7 +117,7 @@       mconcat (map (vname Type) (M.keys abs))         <> forEnv file_env       where-        file' = makeRelative "/" file+        file' = makeRelative "/" $ includeToFilePath file         vname ns v = M.singleton (qualLeaf v) (file', ns)         vname' ((ns, _), v) = vname ns v @@ -133,13 +133,13 @@ -- @important_imports@ considered most important.  The HTML files must -- be written to the specific locations indicated in the return value, -- or the relative links will be wrong.-renderFiles :: [FilePath] -> Imports -> ([(FilePath, Html)], Warnings)+renderFiles :: [ImportName] -> Imports -> ([(FilePath, Html)], Warnings) renderFiles important_imports imports = runWriter $ do   (import_pages, documented) <- runWriterT $     forM imports $ \(current, fm) ->       let ctx =             Context-              { ctxCurrent = makeRelative "/" current,+              { ctxCurrent = makeRelative "/" $ includeToFilePath current,                 ctxFileMod = fm,                 ctxImports = imports,                 ctxNoLink = mempty,@@ -155,15 +155,19 @@              pure               ( current,-                ( H.docTypeHtml ! A.lang "en" $-                    addBoilerplateWithNav important_imports imports ("doc" </> current) current $-                      H.main $-                        maybe_abstract-                          <> selfLink "synopsis" (H.h2 "Synopsis")-                          <> (H.div ! A.id "overview") synopsis-                          <> selfLink "description" (H.h2 "Description")-                          <> description-                          <> maybe_sections,+                ( H.docTypeHtml ! A.lang "en"+                    $ addBoilerplateWithNav+                      important_imports+                      imports+                      ("doc" </> includeToFilePath current)+                      (includeToString current)+                    $ H.main+                    $ maybe_abstract+                      <> selfLink "synopsis" (H.h2 "Synopsis")+                      <> (H.div ! A.id "overview") synopsis+                      <> selfLink "description" (H.h2 "Description")+                      <> description+                      <> maybe_sections,                   first_paragraph                 )               )@@ -175,7 +179,8 @@       ++ map (importHtml *** fst) import_pages   where     file_map = vnameToFileMap imports-    importHtml import_name = "doc" </> makeRelative "/" import_name <.> "html"+    importHtml import_name =+      "doc" </> makeRelative "/" (fromString (includeToString import_name)) <.> "html"  -- | The header documentation (which need not be present) can contain -- an abstract and further sections.@@ -201,7 +206,7 @@     firstParagraph = unlines . takeWhile (not . paragraphSeparator) . lines     paragraphSeparator = all isSpace -contentsPage :: [FilePath] -> [(String, Html)] -> Html+contentsPage :: [ImportName] -> [(ImportName, Html)] -> Html contentsPage important_imports pages =   H.docTypeHtml $     addBoilerplate "index.html" "Futhark Library Documentation" $@@ -229,12 +234,15 @@         (H.dt ! A.class_ "desc_header") (importLink "index.html" name)           <> (H.dd ! A.class_ "desc_doc") maybe_abstract -importLink :: FilePath -> String -> Html+importLink :: FilePath -> ImportName -> Html importLink current name =-  let file = relativise ("doc" </> makeRelative "/" name -<.> "html") current-   in (H.a ! A.href (fromString file) $ fromString name)+  let file =+        relativise+          ("doc" </> makeRelative "/" (includeToFilePath name) -<.> "html")+          current+   in (H.a ! A.href (fromString file) $ fromString (includeToString name)) -indexPage :: [FilePath] -> Imports -> Documented -> FileMap -> Html+indexPage :: [ImportName] -> Imports -> Documented -> FileMap -> Html indexPage important_imports imports documented fm =   H.docTypeHtml $     addBoilerplateWithNav important_imports imports "doc-index.html" "Index" $@@ -344,7 +352,7 @@     futhark_doc_url =       "https://futhark.readthedocs.io/en/latest/man/futhark-doc.html" -addBoilerplateWithNav :: [FilePath] -> Imports -> String -> String -> Html -> Html+addBoilerplateWithNav :: [ImportName] -> Imports -> String -> String -> Html -> Html addBoilerplateWithNav important_imports imports current titleText content =   addBoilerplate current titleText $     (H.nav ! A.id "filenav" $ files) <> content@@ -390,7 +398,7 @@   Just $ parens <$> me' synopsisOpened (ModImport _ (Info file) _) = Just $ do   current <- asks ctxCurrent-  let dest = fromString $ relativise file current <> ".html"+  let dest = fromString $ relativise (includeToFilePath file) current <> ".html"   pure $ keyword "import " <> (H.a ! A.href dest) (fromString $ show file) synopsisOpened (ModAscript _ se _ _) = Just $ do   se' <- synopsisSigExp se@@ -411,7 +419,7 @@           map typeParamName tparams             ++ map identName (S.toList $ mconcat $ map patIdents params)   rettype' <- noLink' $ maybe (retTypeHtml rettype) typeExpHtml retdecl-  params' <- noLink' $ mapM patternHtml params+  params' <- noLink' $ mapM paramHtml params   pure     ( keyword "val " <> (H.span ! A.class_ "decl_name") name,       tparams',@@ -485,6 +493,10 @@       <> ": "       <> t' +dietHtml :: Diet -> Html+dietHtml Consume = "*"+dietHtml Observe = ""+ typeHtml :: StructType -> DocM Html typeHtml t = case t of   Array _ u shape et -> do@@ -505,14 +517,14 @@     targs' <- mapM typeArgHtml targs     et' <- qualNameHtml et     pure $ prettyU u <> et' <> mconcat (map (" " <>) targs')-  Scalar (Arrow _ pname t1 t2) -> do+  Scalar (Arrow _ pname d t1 t2) -> do     t1' <- typeHtml t1     t2' <- retTypeHtml t2     pure $ case pname of       Named v ->-        parens (vnameHtml v <> ": " <> t1') <> " -> " <> t2'+        parens (vnameHtml v <> ": " <> dietHtml d <> t1') <> " -> " <> t2'       Unnamed ->-        t1' <> " -> " <> t2'+        dietHtml d <> t1' <> " -> " <> t2'   Scalar (Sum cs) -> pipes <$> mapM ppClause (sortConstrs cs)     where       ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeHtml ts@@ -615,7 +627,7 @@     specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig   IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e -typeExpHtml :: TypeExp VName -> DocM Html+typeExpHtml :: TypeExp Info VName -> DocM Html typeExpHtml e = case e of   TEUnique t _ -> ("*" <>) <$> typeExpHtml t   TEArray d at _ -> do@@ -680,12 +692,12 @@     then "#" ++ show tag     else relativise file current ++ ".html#" ++ show tag -patternHtml :: Pat -> DocM Html-patternHtml pat = do-  let (pat_param, t) = patternParam pat+paramHtml :: Pat -> DocM Html+paramHtml pat = do+  let (pat_param, d, t) = patternParam pat   t' <- typeHtml t   pure $ case pat_param of-    Named v -> parens (vnameHtml v <> ": " <> t')+    Named v -> parens (vnameHtml v <> ": " <> dietHtml d <> t')     Unnamed -> t'  relativise :: FilePath -> FilePath -> FilePath@@ -697,13 +709,12 @@ dimDeclHtml (ConstSize n) = pure $ brackets $ toHtml (show n) dimDeclHtml AnySize {} = pure $ brackets mempty -dimExpHtml :: SizeExp VName -> DocM Html-dimExpHtml SizeExpAny = pure $ brackets mempty-dimExpHtml (SizeExpNamed v _) = brackets <$> qualNameHtml v-dimExpHtml (SizeExpConst n _) = pure $ brackets $ toHtml (show n)+dimExpHtml :: SizeExp Info VName -> DocM Html+dimExpHtml (SizeExpAny _) = pure $ brackets mempty+dimExpHtml (SizeExp e _) = pure $ brackets $ toHtml $ prettyString e -typeArgExpHtml :: TypeArgExp VName -> DocM Html-typeArgExpHtml (TypeArgExpDim d _) = dimExpHtml d+typeArgExpHtml :: TypeArgExp Info VName -> DocM Html+typeArgExpHtml (TypeArgExpSize d) = dimExpHtml d typeArgExpHtml (TypeArgExpType d) = typeExpHtml d  typeParamHtml :: TypeParam -> Html@@ -762,13 +773,13 @@ lookupName :: (Namespace, String, Maybe FilePath) -> DocM (Maybe VName) lookupName (namespace, name, file) = do   current <- asks ctxCurrent-  let file' = includeToString . flip (mkImportFrom (mkInitialImport current)) mempty <$> file+  let file' = mkImportFrom (mkInitialImport current) <$> file   env <- lookupEnvForFile file'   case M.lookup (namespace, nameFromString name) . envNameMap =<< env of     Nothing -> pure Nothing     Just qn -> pure $ Just $ qualLeaf qn -lookupEnvForFile :: Maybe FilePath -> DocM (Maybe Env)+lookupEnvForFile :: Maybe ImportName -> DocM (Maybe Env) lookupEnvForFile Nothing = asks $ Just . fileEnv . ctxFileMod lookupEnvForFile (Just file) = asks $ fmap fileEnv . lookup file . ctxImports 

@@ -50,6 +50,7 @@ import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe (fromJust, isJust, isNothing)+import Data.Traversable import Futhark.Analysis.AlgSimplify qualified as AlgSimplify import Futhark.Analysis.PrimExp import Futhark.Analysis.PrimExp.Convert@@ -217,16 +218,16 @@   freeIn' (LMADDim s n _ _) = freeIn' s <> freeIn' n  instance Functor LMAD where-  fmap f = runIdentity . traverse (pure . f)+  fmap = fmapDefault  instance Functor IxFun where-  fmap f = runIdentity . traverse (pure . f)+  fmap = fmapDefault  instance Foldable LMAD where-  foldMap f = execWriter . traverse (tell . f)+  foldMap = foldMapDefault  instance Foldable IxFun where-  foldMap f = execWriter . traverse (tell . f)+  foldMap = foldMapDefault  instance Traversable LMAD where   traverse f (LMAD offset dims) =@@ -1043,9 +1044,11 @@             && isNothing               ( selfOverlap () () less_thans (map (flip LeafExp $ IntType Int64) $ namesToList non_negatives) interval2''               )-            && any-              (not . uncurry (intervalOverlap less_thans non_negatives))-              (zip interval1'' interval2'')+            && not+              ( all+                  (uncurry (intervalOverlap less_thans non_negatives))+                  (zip interval1'' interval2'')+              )         _ ->           False @@ -1079,9 +1082,10 @@                 (Nothing, Nothing) ->                   case namesFromList <$> mapM justLeafExp non_negatives of                     Just non_negatives' ->-                      any-                        (not . uncurry (intervalOverlap less_thans non_negatives'))-                        (zip is1 is2)+                      not $+                        all+                          (uncurry (intervalOverlap less_thans non_negatives'))+                          (zip is1 is2)                     _ -> False                 (Just overlapping_dim, _) ->                   let expanded_offset = AlgSimplify.simplifySofP' <$> expandOffset offset is1

@@ -110,8 +110,7 @@ callKernelRules =   standardRules     <> ruleBook-      [ RuleBasicOp copyCopyToCopy,-        RuleMatch unExistentialiseMemory,+      [ RuleMatch unExistentialiseMemory,         RuleOp decertifySafeAlloc       ]       []@@ -182,34 +181,6 @@       | otherwise =           fixable unExistentialiseMemory _ _ _ _ = Skip---- | If we are copying something that is itself a copy, just copy the--- original one instead.-copyCopyToCopy ::-  ( BuilderOps rep,-    LetDec rep ~ (VarWisdom, MemBound u)-  ) =>-  TopDownRuleBasicOp rep-copyCopyToCopy vtable pat@(Pat [pat_elem]) _ (Copy v1)-  | Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable,-    Just (_, MemArray _ _ _ (ArrayIn srcmem src_ixfun)) <--      ST.entryLetBoundDec =<< ST.lookup v1 vtable,-    Just (Mem src_space) <- ST.lookupType srcmem vtable,-    (_, MemArray _ _ _ (ArrayIn destmem dest_ixfun)) <- patElemDec pat_elem,-    Just (Mem dest_space) <- ST.lookupType destmem vtable,-    src_space == dest_space,-    dest_ixfun == src_ixfun =-      Simplify $ certifying v1_cs $ letBind pat $ BasicOp $ Copy v2-copyCopyToCopy vtable pat _ (Copy v0)-  | Just (BasicOp (Rearrange perm v1), v0_cs) <- ST.lookupExp v0 vtable,-    Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable = Simplify $ do-      v0' <--        certifying (v0_cs <> v1_cs) $-          letExp "rearrange_v0" $-            BasicOp $-              Rearrange perm v2-      letBind pat $ BasicOp $ Copy v0'-copyCopyToCopy _ _ _ _ = Skip  -- If an allocation is statically known to be safe, then we can remove -- the certificates on it.  This can help hoist things that would

@@ -273,7 +273,7 @@ instance PrettyRep rep => Pretty (Exp rep) where   pretty (Match [c] [Case [Just (BoolValue True)] t] f (MatchDec ret ifsort)) =     "if"-      <+> info'+      <> info'       <+> pretty c       </> "then"       <+> maybeNest t@@ -284,8 +284,8 @@     where       info' = case ifsort of         MatchNormal -> mempty-        MatchFallback -> "<fallback>"-        MatchEquiv -> "<equiv>"+        MatchFallback -> " <fallback>"+        MatchEquiv -> " <equiv>"   pretty (Match ses cs defb (MatchDec ret ifsort)) =     ("match" <+> info' <+> ppTuple' (map pretty ses))       </> stack (map pretty cs)@@ -297,8 +297,8 @@     where       info' = case ifsort of         MatchNormal -> mempty-        MatchFallback -> "<fallback>"-        MatchEquiv -> "<equiv>"+        MatchFallback -> " <fallback>"+        MatchEquiv -> " <equiv>"   pretty (BasicOp op) = pretty op   pretty (Apply fname args ret (safety, _, _)) =     applykw

@@ -749,6 +749,9 @@   S.Set (Pat (LetDec rep), ArrayOp) arrayOps cs = mconcat . map onStm . stmsToList . bodyStms   where+    -- It is not safe to move everything out of branches (#1874);+    -- probably we need to put some more intelligence in here somehow.+    onStm (Let _ _ Match {}) = mempty     onStm (Let pat aux e) =       case isArrayOp (cs <> stmAuxCerts aux) e of         Just op -> S.singleton (pat, op)

@@ -1123,7 +1123,7 @@   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)   (reds', reds_hoisted) <-     Engine.localVtable (<> scope_vtable) $-      unzip <$> mapM (simplifySegBinOp (segFlat space)) reds+      mapAndUnzipM (simplifySegBinOp (segFlat space)) reds   (kbody', body_hoisted) <- simplifyKernelBody space kbody    pure@@ -1137,7 +1137,7 @@   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)   (scans', scans_hoisted) <-     Engine.localVtable (<> scope_vtable) $-      unzip <$> mapM (simplifySegBinOp (segFlat space)) scans+      mapAndUnzipM (simplifySegBinOp (segFlat space)) scans   (kbody', body_hoisted) <- simplifyKernelBody space kbody    pure

@@ -7,7 +7,7 @@ import Data.Bifunctor import Data.Bitraversable import Data.Foldable-import Data.List (partition, sortOn, tails)+import Data.List (partition, sortOn) import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe@@ -17,6 +17,7 @@ import Futhark.Util (mapAccumLM) import Language.Futhark import Language.Futhark.Traversals+import Language.Futhark.TypeChecker.Types (Subst (..), applySubst)  -- | A static value stores additional information about the result of -- defunctionalization of an expression, aside from the residual expression.@@ -143,18 +144,11 @@         loc     onExp substs e = onAST substs e -    onTypeExpDim substs d@(SizeExpNamed v loc) =-      case M.lookup (qualLeaf v) substs of-        Just (SubstNamed v') ->-          SizeExpNamed v' loc-        Just (SubstConst x) ->-          SizeExpConst x loc-        Nothing ->-          d-    onTypeExpDim _ d = d+    onTypeExpDim substs (SizeExp e loc) = SizeExp (onExp substs e) loc+    onTypeExpDim _ (SizeExpAny loc) = SizeExpAny loc -    onTypeArgExp substs (TypeArgExpDim d loc) =-      TypeArgExpDim (onTypeExpDim substs d) loc+    onTypeArgExp substs (TypeArgExpSize d) =+      TypeArgExpSize (onTypeExpDim substs d)     onTypeArgExp substs (TypeArgExpType te) =       TypeArgExpType (onTypeExp substs te) @@ -287,7 +281,7 @@  data SizeSubst   = SubstNamed (QualName VName)-  | SubstConst Int+  | SubstConst Int64   deriving (Eq, Ord, Show)  dimMapping ::@@ -384,7 +378,7 @@         [pat'] -> (pat', ret, e0)         (pat' : pats') ->           ( pat',-            RetType [] $ foldFunType (map (toStruct . patternType) pats') ret,+            RetType [] $ funType pats' ret,             Lambda pats' e0 Nothing (Info (mempty, ret)) loc           ) @@ -447,10 +441,10 @@   (e', sv) <- defuncExp e   pure (QualParens qn e' loc, sv) defuncExp (TupLit es loc) = do-  (es', svs) <- unzip <$> mapM defuncExp es+  (es', svs) <- mapAndUnzipM defuncExp es   pure (TupLit es' loc, RecordSV $ zip tupleFieldNames svs) defuncExp (RecordLit fs loc) = do-  (fs', names_svs) <- unzip <$> mapM defuncField fs+  (fs', names_svs) <- mapAndUnzipM defuncField fs   pure (RecordLit fs' loc, RecordSV names_svs)   where     defuncField (RecordFieldExplicit vn e loc') = do@@ -494,7 +488,7 @@     -- Intrinsic functions used as variables are eta-expanded, so we     -- can get rid of them.     IntrinsicSV -> do-      (pats, body, tp) <- etaExpand (typeOf e) e+      (pats, body, tp) <- etaExpand (RetType [] (typeOf e)) e       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty     HoleSV _ hole_loc ->       pure (Hole (Info t) hole_loc, sv)@@ -533,16 +527,8 @@   (e2', sv) <- defuncExp e2   (e3', _) <- defuncExp e3   pure (AppExp (If e1' e2' e3' loc) res, sv)-defuncExp e@(AppExp (Apply f@(Var f' _ _) arg d loc) res)-  | baseTag (qualLeaf f') <= maxIntrinsicTag,-    TupLit es tuploc <- arg = do-      -- defuncSoacExp also works fine for non-SOACs.-      es' <- mapM defuncSoacExp es-      pure-        ( AppExp (Apply f (TupLit es' tuploc) d loc) res,-          Dynamic $ typeOf e-        )-defuncExp e@(AppExp Apply {} _) = defuncApply 0 e+defuncExp (AppExp (Apply f args loc) (Info appres)) =+  defuncApply f (fmap (first unInfo) args) appres loc defuncExp (Negate e0 loc) = do   (e0', sv) <- defuncExp e0   pure (Negate e0' loc, sv)@@ -633,7 +619,7 @@   (e2', sv) <- defuncExp e2   pure (Assert e1' e2' desc loc, sv) defuncExp (Constr name es (Info sum_t@(Scalar (Sum all_fs))) loc) = do-  (es', svs) <- unzip <$> mapM defuncExp es+  (es', svs) <- mapAndUnzipM defuncExp es   let sv =         SumSV name svs $           M.toList $@@ -701,34 +687,42 @@   pure $ Lambda params e0' decl tp loc defuncSoacExp e   | Scalar Arrow {} <- typeOf e = do-      (pats, body, tp) <- etaExpand (typeOf e) e+      (pats, body, tp) <- etaExpand (RetType [] (typeOf e)) e       let env = foldMap envFromPat pats       body' <- localEnv env $ defuncExp' body       pure $ Lambda pats body' Nothing (Info (mempty, tp)) mempty   | otherwise = defuncExp' e -etaExpand :: PatType -> Exp -> DefM ([Pat], Exp, StructRetType)+etaExpand :: PatRetType -> Exp -> DefM ([Pat], Exp, StructRetType) etaExpand e_t e = do-  let (ps, ret) = getType $ RetType [] e_t+  let (ps, ret) = getType e_t   -- Some careful hackery to avoid duplicate names.   (_, (pats, vars)) <- second unzip <$> mapAccumLM f [] ps-  let e' =-        foldl'-          ( \e1 (e2, t2, argtypes) ->-              AppExp-                (Apply e1 e2 (Info (diet t2, Nothing)) mempty)-                (Info (AppRes (foldFunType argtypes ret) []))-          )+  -- Important that we synthesize new existential names and substitute+  -- them into the (body) return type.+  ext' <- mapM newName $ retDims ret+  let extsubst =+        M.fromList . zip (retDims ret) $+          map (SizeSubst . NamedSize . qualName) ext'+      ret' = applySubst (`M.lookup` extsubst) ret+      e' =+        mkApply           e-          $ zip3 vars (map snd ps) (drop 1 $ tails $ map snd ps)+          (zip3 (map (diet . snd . snd) ps) (repeat Nothing) vars)+          (AppRes (retType ret') ext')   pure (pats, e', second (const ()) ret)   where-    getType (RetType _ (Scalar (Arrow _ p t1 t2))) =-      let (ps, r) = getType t2 in ((p, t1) : ps, r)+    getType (RetType _ (Scalar (Arrow _ p d t1 t2))) =+      let (ps, r) = getType t2+       in ((p, (d, t1)) : ps, r)     getType t = ([], t) -    f prev (p, t) = do-      let t' = fromStruct t+    f prev (p, (d, t)) = do+      let t' =+            fromStruct t+              `setUniqueness` case d of+                Consume -> Unique+                Observe -> Nonunique       x <- case p of         Named x | x `notElem` prev -> pure x         _ -> newNameFromString "x"@@ -812,156 +806,21 @@     onDim (NamedSize d) | qualLeaf d `elem` ext = AnySize Nothing     onDim d = d --- | Defunctionalize an application expression at a given depth of application.--- Calls to dynamic (first-order) functions are preserved at much as possible,--- but a new lifted function is created if a dynamic function is only partially--- applied.-defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)-defuncApply depth e@(AppExp (Apply e1 e2 d loc) t@(Info (AppRes ret ext))) = do-  let (argtypes, _) = unfoldFunType ret-  (e1', sv1) <- defuncApply (depth + 1) e1-  (e2', sv2) <- defuncExp e2-  let e' = AppExp (Apply e1' e2' d loc) t-  case sv1 of-    LambdaSV pat e0_t e0 closure_env -> do-      let env' = matchPatSV pat sv2-          dims = mempty-      (e0', sv) <--        localNewEnv (env' <> closure_env) $-          defuncExp e0--      let closure_pat = buildEnvPat dims closure_env-          pat' = updatePat pat sv2--      globals <- asks fst--      -- Lift lambda to top-level function definition.  We put in-      -- a lot of effort to try to infer the uniqueness attributes-      -- of the lifted function, but this is ultimately all a sham-      -- and a hack.  There is some piece we're missing.-      let params = [closure_pat, pat']-          params_for_rettype = params ++ svParams sv1 ++ svParams sv2-          svParams (LambdaSV sv_pat _ _ _) = [sv_pat]-          svParams _ = []-          lifted_rettype = buildRetType closure_env params_for_rettype (unRetType e0_t) $ typeOf e0'--          already_bound =-            globals-              <> S.fromList dims-              <> S.map identName (foldMap patIdents params)--          more_dims =-            S.toList $-              S.filter (`S.notMember` already_bound) $-                foldMap patternArraySizes params--          -- Embed some information about the original function-          -- into the name of the lifted function, to make the-          -- result slightly more human-readable.-          liftedName i (Var f _ _) =-            "defunc_" ++ show i ++ "_" ++ baseString (qualLeaf f)-          liftedName i (AppExp (Apply f _ _ _) _) =-            liftedName (i + 1) f-          liftedName _ _ = "defunc"--      -- Ensure that no parameter sizes are AnySize.  The internaliser-      -- expects this.  This is easy, because they are all-      -- first-order.-      let bound_sizes = S.fromList (dims <> more_dims) <> globals-      (missing_dims, params') <- sizesForAll bound_sizes params--      fname <- newNameFromString $ liftedName (0 :: Int) e1-      liftValDec-        fname-        (RetType [] $ toStruct lifted_rettype)-        (dims ++ more_dims ++ missing_dims)-        params'-        e0'--      let t1 = toStruct $ typeOf e1'-          t2 = toStruct $ typeOf e2'-          fname' = qualName fname-          fname'' =-            Var-              fname'-              ( Info-                  ( Scalar . Arrow mempty Unnamed t1 . RetType [] $-                      Scalar . Arrow mempty Unnamed t2 $-                        RetType [] lifted_rettype-                  )-              )-              loc--          callret = AppRes (combineTypeShapes ret lifted_rettype) ext--          innercallret =-            AppRes-              (Scalar $ Arrow mempty Unnamed t2 $ RetType [] lifted_rettype)-              []--      pure-        ( AppExp-            ( Apply-                ( AppExp-                    (Apply fname'' e1' (Info (Observe, Nothing)) loc)-                    (Info innercallret)-                )-                e2'-                d-                loc-            )-            (Info callret),-          sv-        )--    -- If e1 is a dynamic function, we just leave the application in place,-    -- but we update the types since it may be partially applied or return-    -- a higher-order term.-    DynamicFun _ sv -> do-      let (argtypes', rettype) = dynamicFunType sv argtypes-          restype = foldFunType argtypes' (RetType [] rettype) `setAliases` aliases ret-          callret = AppRes (combineTypeShapes ret restype) ext-          apply_e = AppExp (Apply e1' e2' d loc) (Info callret)-      pure (apply_e, sv)-    -- Propagate the 'IntrinsicsSV' until we reach the outermost application,-    -- where we construct a dynamic static value with the appropriate type.-    IntrinsicSV -> intrinsicOrHole argtypes e' sv1-    HoleSV {} -> intrinsicOrHole argtypes e' sv1-    _ ->-      error $-        "Application of an expression\n"-          ++ prettyString e1-          ++ "\nthat is neither a static lambda "-          ++ "nor a dynamic function, but has static value:\n"-          ++ show sv1-  where-    intrinsicOrHole argtypes e' sv-      | depth == 0 =-          -- If the intrinsic is fully applied, then we are done.-          -- Otherwise we need to eta-expand it and recursively-          -- defunctionalise. XXX: might it be better to simply-          -- eta-expand immediately any time we encounter a-          -- non-fully-applied intrinsic?-          if null argtypes-            then pure (e', Dynamic $ typeOf e)-            else do-              (pats, body, tp) <- etaExpand (typeOf e') e'-              defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty-      | otherwise = pure (e', sv)-defuncApply depth e@(Var qn (Info t) loc) = do+defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)+defuncApplyFunction e@(Var qn (Info t) loc) num_args = do   let (argtypes, _) = unfoldFunType t   sv <- lookupVar (toStruct t) (qualLeaf qn)    case sv of     DynamicFun _ _-      | fullyApplied sv depth -> do+      | fullyApplied sv num_args -> do           -- We still need to update the types in case the dynamic           -- function returns a higher-order term.           let (argtypes', rettype) = dynamicFunType sv argtypes           pure (Var qn (Info (foldFunType argtypes' $ RetType [] rettype)) loc, sv)       | otherwise -> do           fname <- newVName $ "dyn_" <> baseString (qualLeaf qn)-          let (pats, e0, sv') = liftDynFun (prettyString qn) sv depth+          let (pats, e0, sv') = liftDynFun (prettyString qn) sv num_args               (argtypes', rettype) = dynamicFunType sv' argtypes               dims' = mempty @@ -982,9 +841,132 @@             )     IntrinsicSV -> pure (e, IntrinsicSV)     _ -> pure (Var qn (Info (typeFromSV sv)) loc, sv)-defuncApply depth (Parens e _) = defuncApply depth e-defuncApply _ expr = defuncExp expr+defuncApplyFunction e _ = defuncExp e +-- Embed some information about the original function+-- into the name of the lifted function, to make the+-- result slightly more human-readable.+liftedName :: Int -> Exp -> String+liftedName i (Var f _ _) =+  "defunc_" ++ show i ++ "_" ++ baseString (qualLeaf f)+liftedName i (AppExp (Apply f _ _) _) =+  liftedName (i + 1) f+liftedName _ _ = "defunc"++defuncApplyArg ::+  String ->+  (Exp, StaticVal) ->+  (((Diet, Maybe VName), Exp), [(Diet, StructType)]) ->+  DefM (Exp, StaticVal)+defuncApplyArg fname_s (f', f_sv@(LambdaSV pat lam_e_t lam_e closure_env)) (((d, argext), arg), _) = do+  (arg', arg_sv) <- defuncExp arg+  let env' = matchPatSV pat arg_sv+      dims = mempty+  (lam_e', sv) <-+    localNewEnv (env' <> closure_env) $+      defuncExp lam_e++  let closure_pat = buildEnvPat dims closure_env+      pat' = updatePat pat arg_sv++  globals <- asks fst++  -- Lift lambda to top-level function definition.  We put in+  -- a lot of effort to try to infer the uniqueness attributes+  -- of the lifted function, but this is ultimately all a sham+  -- and a hack.  There is some piece we're missing.+  let params = [closure_pat, pat']+      params_for_rettype = params ++ svParams f_sv ++ svParams arg_sv+      svParams (LambdaSV sv_pat _ _ _) = [sv_pat]+      svParams _ = []+      lifted_rettype = buildRetType closure_env params_for_rettype (unRetType lam_e_t) $ typeOf lam_e'++      already_bound =+        globals+          <> S.fromList dims+          <> S.map identName (foldMap patIdents params)++      more_dims =+        S.toList $+          S.filter (`S.notMember` already_bound) $+            foldMap patternArraySizes params++  -- Ensure that no parameter sizes are AnySize.  The internaliser+  -- expects this.  This is easy, because they are all+  -- first-order.+  let bound_sizes = S.fromList (dims <> more_dims) <> globals+  (missing_dims, params') <- sizesForAll bound_sizes params++  fname <- newNameFromString fname_s+  liftValDec+    fname+    (RetType [] $ toStruct lifted_rettype)+    (dims ++ more_dims ++ missing_dims)+    params'+    lam_e'++  let f_t = toStruct $ typeOf f'+      arg_t = toStruct $ typeOf arg'+      d1 = Observe+      fname_t = foldFunType [(d1, f_t), (d, arg_t)] $ RetType [] lifted_rettype+      fname' = Var (qualName fname) (Info fname_t) (srclocOf arg)+      callret = AppRes lifted_rettype []++  pure+    ( mkApply fname' [(Observe, Nothing, f'), (Observe, argext, arg')] callret,+      sv+    )+-- If 'f' is a dynamic function, we just leave the application in+-- place, but we update the types since it may be partially+-- applied or return a higher-order value.+defuncApplyArg _ (f', DynamicFun _ sv) (((d, argext), arg), argtypes) = do+  (arg', _) <- defuncExp arg+  let (argtypes', rettype) = dynamicFunType sv argtypes+      restype = foldFunType argtypes' (RetType [] rettype)+      callret = AppRes restype []+      apply_e = mkApply f' [(d, argext, arg')] callret+  pure (apply_e, sv)+--+defuncApplyArg _ (_, sv) _ =+  error $ "defuncApplyArg: cannot apply StaticVal\n" <> show sv++updateReturn :: AppRes -> Exp -> Exp+updateReturn (AppRes ret1 ext1) (AppExp apply (Info (AppRes ret2 ext2))) =+  AppExp apply $ Info $ AppRes (combineTypeShapes ret1 ret2) (ext1 <> ext2)+updateReturn _ e = e++defuncApply :: Exp -> NE.NonEmpty ((Diet, Maybe VName), Exp) -> AppRes -> SrcLoc -> DefM (Exp, StaticVal)+defuncApply f args appres loc = do+  (f', f_sv) <- defuncApplyFunction f (length args)+  case f_sv of+    IntrinsicSV -> do+      args' <- fmap (first Info) <$> traverse (traverse defuncSoacExp) args+      let e' = AppExp (Apply f' args' loc) (Info appres)+      intrinsicOrHole e'+    HoleSV {} -> do+      args' <- fmap (first Info) <$> traverse (traverse $ fmap fst . defuncExp) args+      let e' = AppExp (Apply f' args' loc) (Info appres)+      intrinsicOrHole e'+    _ -> do+      let fname = liftedName 0 f+          (argtypes, _) = unfoldFunType $ typeOf f+      fmap (first $ updateReturn appres) $+        foldM (defuncApplyArg fname) (f', f_sv) $+          NE.zip args $+            NE.tails argtypes+  where+    intrinsicOrHole e' = do+      -- If the intrinsic is fully applied, then we are done.+      -- Otherwise we need to eta-expand it and recursively+      -- defunctionalise. XXX: might it be better to simply eta-expand+      -- immediately any time we encounter a non-fully-applied+      -- intrinsic?+      if null $ fst $ unfoldFunType $ appResType appres+        then pure (e', Dynamic $ appResType appres)+        else do+          (pats, body, tp) <- etaExpand (RetType [] (typeOf e')) e'+          defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty+ -- | Check if a 'StaticVal' and a given application depth corresponds -- to a fully applied dynamic function. fullyApplied :: StaticVal -> Int -> Bool@@ -1129,9 +1111,10 @@  -- | Construct the type for a fully-applied dynamic function from its -- static value and the original types of its arguments.-dynamicFunType :: StaticVal -> [StructType] -> ([PatType], PatType)+dynamicFunType :: StaticVal -> [(Diet, StructType)] -> ([(Diet, PatType)], PatType) dynamicFunType (DynamicFun _ sv) (p : ps) =-  let (ps', ret) = dynamicFunType sv ps in (fromStruct p : ps', ret)+  let (ps', ret) = dynamicFunType sv ps+   in (second fromStruct p : ps', ret) dynamicFunType sv _ = ([], typeFromSV sv)  -- | Match a pattern with its static value. Returns an environment with@@ -1177,11 +1160,10 @@ matchPatSV pat (HoleSV t _) = matchPatSV pat $ svFromType t matchPatSV pat sv =   error $-    "Tried to match pattern "+    "Tried to match pattern\n"       ++ prettyString pat-      ++ " with static value "+      ++ "\n with static value\n"       ++ show sv-      ++ "."  orderZeroSV :: StaticVal -> Bool orderZeroSV Dynamic {} = True@@ -1231,9 +1213,9 @@ updatePat pat (HoleSV t _) = updatePat pat (svFromType t) updatePat pat sv =   error $-    "Tried to update pattern "+    "Tried to update pattern\n"       ++ prettyString pat-      ++ "to reflect the static value "+      ++ "\nto reflect the static value\n"       ++ show sv  -- | Convert a record (or tuple) type to a record static value. This is used for@@ -1248,9 +1230,9 @@ -- boolean is true if the function is a 'DynamicFun'. defuncValBind :: ValBind -> DefM (ValBind, Env) -- Eta-expand entry points with a functional return type.-defuncValBind (ValBind entry name _ (Info (RetType _ rettype)) tparams params body _ attrs loc)-  | Scalar Arrow {} <- rettype = do-      (body_pats, body', rettype') <- etaExpand (fromStruct rettype) body+defuncValBind (ValBind entry name _ (Info rettype) tparams params body _ attrs loc)+  | Scalar Arrow {} <- retType rettype = do+      (body_pats, body', rettype') <- etaExpand (second (const mempty) rettype) body       defuncValBind $         ValBind           entry

@@ -10,7 +10,7 @@ import Data.Set qualified as S import Futhark.MonadFreshNames import Language.Futhark-import Language.Futhark.Semantic (FileModule (..), Imports)+import Language.Futhark.Semantic (FileModule (..), Imports, includeToString) import Language.Futhark.Traversals import Prelude hiding (abs, mod) @@ -61,7 +61,7 @@ data Env = Env   { envScope :: Scope,     envGenerating :: Bool,-    envImports :: M.Map String Scope,+    envImports :: M.Map ImportName Scope,     envAbs :: TySet   } @@ -110,7 +110,7 @@ generating :: TransformM a -> TransformM a generating = local $ \env -> env {envGenerating = True} -bindingImport :: String -> Scope -> TransformM a -> TransformM a+bindingImport :: ImportName -> Scope -> TransformM a -> TransformM a bindingImport name scope = local $ \env ->   env {envImports = M.insert name scope $ envImports env} @@ -118,10 +118,10 @@ bindingAbs abs = local $ \env ->   env {envAbs = abs <> envAbs env} -lookupImport :: String -> TransformM Scope+lookupImport :: ImportName -> TransformM Scope lookupImport name = maybe bad pure =<< asks (M.lookup name . envImports)   where-    bad = error $ "Defunctorise: unknown import: " ++ name+    bad = error $ "Defunctorise: unknown import: " ++ includeToString name  lookupMod' :: QualName VName -> Scope -> Either String Mod lookupMod' mname scope =@@ -264,7 +264,7 @@               astMap (substituter $ modScope mod <> scope) e'         _ -> astMap (substituter scope) e -transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName)+transformTypeExp :: TypeExp Info VName -> TransformM (TypeExp Info VName) transformTypeExp = transformNames  transformStructType :: StructType -> TransformM StructType

@@ -30,7 +30,7 @@     decTypes (E.TypeDec tb) = [tb]     decTypes _ = [] -findType :: VName -> VisibleTypes -> Maybe (E.TypeExp VName)+findType :: VName -> VisibleTypes -> Maybe (E.TypeExp E.Info VName) findType v (VisibleTypes ts) = E.typeExp <$> find ((== v) . E.typeAlias) ts  valueType :: I.TypeBase I.Rank Uniqueness -> I.ValueType@@ -39,18 +39,18 @@ valueType I.Acc {} = error "valueType Acc" valueType I.Mem {} = error "valueType Mem" -withoutDims :: E.TypeExp VName -> (Int, E.TypeExp VName)+withoutDims :: E.TypeExp E.Info VName -> (Int, E.TypeExp E.Info VName) withoutDims (E.TEArray _ te _) =   let (d, te') = withoutDims te    in (d + 1, te') withoutDims te = (0 :: Int, te) -rootType :: E.TypeExp VName -> E.TypeExp VName-rootType (E.TEApply te E.TypeArgExpDim {} _) = rootType te+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 te = te -typeExpOpaqueName :: E.TypeExp VName -> Name+typeExpOpaqueName :: E.TypeExp E.Info VName -> Name typeExpOpaqueName = f . rootType   where     f (E.TEArray _ te _) =@@ -66,7 +66,7 @@ addType :: Name -> I.OpaqueType -> GenOpaque () addType s t = modify (<> I.OpaqueTypes [(s, t)]) -isRecord :: VisibleTypes -> E.TypeExp VName -> Maybe (M.Map Name (E.TypeExp VName))+isRecord :: VisibleTypes -> E.TypeExp E.Info VName -> Maybe (M.Map Name (E.TypeExp E.Info VName)) isRecord _ (E.TERecord fs _) = Just $ M.fromList fs isRecord _ (E.TETuple fs _) = Just $ E.tupleFields fs isRecord types (E.TEVar v _) = isRecord types =<< findType (E.qualLeaf v) types@@ -75,7 +75,7 @@ recordFields ::   VisibleTypes ->   M.Map Name E.StructType ->-  Maybe (E.TypeExp VName) ->+  Maybe (E.TypeExp E.Info VName) ->   [(Name, E.EntryType)] recordFields types fs t =   case isRecord types . rootType =<< t of

@@ -23,6 +23,7 @@ import Futhark.Util (splitAt3) import Futhark.Util.Pretty (align, docText, pretty) import Language.Futhark as E hiding (TypeArg)+import Language.Futhark.TypeChecker.Types qualified as E  -- | Convert a program in source Futhark to a program in the Futhark -- core language.@@ -331,15 +332,19 @@             ++ dt'             ++ ["`."]     ensureExtShape (errorMsg parts) loc (I.fromDecl t') desc e'-internaliseAppExp desc _ e@E.Apply {} =+internaliseAppExp desc (E.AppRes et ext) e@E.Apply {} =   case findFuncall e of-    (FunctionHole t loc, _args) -> do+    (FunctionHole loc, _args) -> do       -- The function we are supposed to call doesn't exist, but we       -- have to synthesize some fake values of the right type.  The       -- easy way to do this is to just ignore the arguments and       -- create a hole whose type is the type of the entire-      -- application.-      internaliseExp desc (E.Hole (Info (fromStruct $ snd $ E.unfoldFunType t)) loc)+      -- application.  One caveat is that we need to replace any+      -- existential sizes, too (with zeroes, because they don't+      -- matter).+      let subst = zip ext $ repeat $ E.SizeSubst $ E.ConstSize 0+          et' = E.applySubst (`lookup` subst) et+      internaliseExp desc (E.Hole (Info et') loc)     (FunctionName qfname, args) -> do       -- Argument evaluation is outermost-in so that any existential sizes       -- created by function applications can be brought into scope.@@ -1439,18 +1444,17 @@  data Function   = FunctionName (E.QualName VName)-  | FunctionHole E.PatType SrcLoc+  | FunctionHole SrcLoc   deriving (Show)  findFuncall :: E.AppExp -> (Function, [(E.Exp, Maybe VName)])-findFuncall (E.Apply f arg (Info (_, argext)) _)-  | E.AppExp f_e _ <- f =-      let (f_e', args) = findFuncall f_e-       in (f_e', args ++ [(arg, argext)])+findFuncall (E.Apply f args _)   | E.Var fname _ _ <- f =-      (FunctionName fname, [(arg, argext)])-  | E.Hole (Info t) loc <- f =-      (FunctionHole t loc, [(arg, argext)])+      (FunctionName fname, map onArg $ NE.toList args)+  | E.Hole (Info _) loc <- f =+      (FunctionHole loc, map onArg $ NE.toList args)+  where+    onArg (Info (_, argext), e) = (e, argext) findFuncall e =   error $ "Invalid function expression in application:\n" ++ prettyString e @@ -1602,13 +1606,13 @@           fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'     handleOps _ _ = Nothing -    handleSOACs [TupLit [lam, arr] _] "map" = Just $ \desc -> do+    handleSOACs [lam, arr] "map" = Just $ \desc -> do       arr' <- internaliseExpToVars "map_arr" arr       arr_ts <- mapM lookupType arr'       lam' <- internaliseLambdaCoerce lam $ map rowType arr_ts       let w = arraysSize 0 arr_ts       letTupExp' desc $ I.Op $ I.Screma w arr' (I.mapSOAC lam')-    handleSOACs [TupLit [k, lam, arr] _] "partition" = do+    handleSOACs [k, lam, arr] "partition" = do       k' <- fromIntegral <$> fromInt32 k       Just $ \_desc -> do         arrs <- internaliseExpToVars "partition_input" arr@@ -1618,43 +1622,43 @@         fromInt32 (Literal (SignedValue (Int32Value k')) _) = Just k'         fromInt32 (IntLit k' (Info (E.Scalar (E.Prim (E.Signed Int32)))) _) = Just $ fromInteger k'         fromInt32 _ = Nothing-    handleSOACs [TupLit [lam, ne, arr] _] "reduce" = Just $ \desc ->+    handleSOACs [lam, ne, arr] "reduce" = Just $ \desc ->       internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc)       where         reduce w red_lam nes arrs =           I.Screma w arrs             <$> I.reduceSOAC [Reduce Noncommutative red_lam nes]-    handleSOACs [TupLit [lam, ne, arr] _] "reduce_comm" = Just $ \desc ->+    handleSOACs [lam, ne, arr] "reduce_comm" = Just $ \desc ->       internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc)       where         reduce w red_lam nes arrs =           I.Screma w arrs             <$> I.reduceSOAC [Reduce Commutative red_lam nes]-    handleSOACs [TupLit [lam, ne, arr] _] "scan" = Just $ \desc ->+    handleSOACs [lam, ne, arr] "scan" = Just $ \desc ->       internaliseScanOrReduce desc "scan" reduce (lam, ne, arr, loc)       where         reduce w scan_lam nes arrs =           I.Screma w arrs <$> I.scanSOAC [Scan scan_lam nes]-    handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist_1d" = Just $ \desc ->+    handleSOACs [rf, dest, op, ne, buckets, img] "hist_1d" = Just $ \desc ->       internaliseHist 1 desc rf dest op ne buckets img loc-    handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist_2d" = Just $ \desc ->+    handleSOACs [rf, dest, op, ne, buckets, img] "hist_2d" = Just $ \desc ->       internaliseHist 2 desc rf dest op ne buckets img loc-    handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist_3d" = Just $ \desc ->+    handleSOACs [rf, dest, op, ne, buckets, img] "hist_3d" = Just $ \desc ->       internaliseHist 3 desc rf dest op ne buckets img loc     handleSOACs _ _ = Nothing -    handleAccs [TupLit [dest, f, bs] _] "scatter_stream" = Just $ \desc ->+    handleAccs [dest, f, bs] "scatter_stream" = Just $ \desc ->       internaliseStreamAcc desc dest Nothing f bs-    handleAccs [TupLit [dest, op, ne, f, bs] _] "hist_stream" = Just $ \desc ->+    handleAccs [dest, op, ne, f, bs] "hist_stream" = Just $ \desc ->       internaliseStreamAcc desc dest (Just (op, ne)) f bs-    handleAccs [TupLit [acc, i, v] _] "acc_write" = Just $ \desc -> do+    handleAccs [acc, i, v] "acc_write" = Just $ \desc -> do       acc' <- head <$> internaliseExpToVars "acc" acc       i' <- internaliseExp1 "acc_i" i       vs <- internaliseExp "acc_v" v       fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc acc' [i'] vs     handleAccs _ _ = Nothing -    handleAD [TupLit [f, x, v] _] fname+    handleAD [f, x, v] fname       | fname `elem` ["jvp2", "vjp2"] = Just $ \desc -> do           x' <- internaliseExp "ad_x" x           v' <- internaliseExp "ad_v" v@@ -1665,10 +1669,10 @@               _ -> VJP lam x' v'     handleAD _ _ = Nothing -    handleRest [E.TupLit [a, si, v] _] "scatter" = Just $ scatterF 1 a si v-    handleRest [E.TupLit [a, si, v] _] "scatter_2d" = Just $ scatterF 2 a si v-    handleRest [E.TupLit [a, si, v] _] "scatter_3d" = Just $ scatterF 3 a si v-    handleRest [E.TupLit [n, m, arr] _] "unflatten" = Just $ \desc -> do+    handleRest [a, si, v] "scatter" = Just $ scatterF 1 a si v+    handleRest [a, si, v] "scatter_2d" = Just $ scatterF 2 a si v+    handleRest [a, si, v] "scatter_3d" = Just $ scatterF 3 a si v+    handleRest [n, m, arr] "unflatten" = Just $ \desc -> do       arrs <- internaliseExpToVars "unflatten_arr" arr       n' <- internaliseExp1 "n" n       m' <- internaliseExp1 "m" m@@ -1716,7 +1720,7 @@             I.ReshapeArbitrary             (reshapeOuter (I.Shape [k]) 2 $ I.arrayShape arr_t)             arr'-    handleRest [TupLit [x, y] _] "concat" = Just $ \desc -> do+    handleRest [x, y] "concat" = Just $ \desc -> do       xs <- internaliseExpToVars "concat_x" x       ys <- internaliseExpToVars "concat_y" y       outer_size <- arraysSize 0 <$> mapM lookupType xs@@ -1731,7 +1735,7 @@       let conc xarr yarr =             I.BasicOp $ I.Concat 0 (xarr :| [yarr]) ressize       mapM (letSubExp desc) $ zipWith conc xs ys-    handleRest [TupLit [offset, e] _] "rotate" = Just $ \desc -> do+    handleRest [offset, e] "rotate" = Just $ \desc -> do       offset' <- internaliseExp1 "rotation_offset" offset       internaliseOperation desc e $ \v -> do         r <- I.arrayRank <$> lookupType v@@ -1742,24 +1746,24 @@       internaliseOperation desc e $ \v -> do         r <- I.arrayRank <$> lookupType v         pure $ I.Rearrange ([1, 0] ++ [2 .. r - 1]) v-    handleRest [TupLit [x, y] _] "zip" = Just $ \desc ->+    handleRest [x, y] "zip" = Just $ \desc ->       mapM (letSubExp "zip_copy" . BasicOp . Copy)         =<< ( (++)                 <$> internaliseExpToVars (desc ++ "_zip_x") x                 <*> internaliseExpToVars (desc ++ "_zip_y") y             )     handleRest [x] "unzip" = Just $ flip internaliseExp x-    handleRest [TupLit [arr, offset, n1, s1, n2, s2] _] "flat_index_2d" = Just $ \desc -> do+    handleRest [arr, offset, n1, s1, n2, s2] "flat_index_2d" = Just $ \desc -> do       flatIndexHelper desc loc arr offset [(n1, s1), (n2, s2)]-    handleRest [TupLit [arr1, offset, s1, s2, arr2] _] "flat_update_2d" = Just $ \desc -> do+    handleRest [arr1, offset, s1, s2, arr2] "flat_update_2d" = Just $ \desc -> do       flatUpdateHelper desc loc arr1 offset [s1, s2] arr2-    handleRest [TupLit [arr, offset, n1, s1, n2, s2, n3, s3] _] "flat_index_3d" = Just $ \desc -> do+    handleRest [arr, offset, n1, s1, n2, s2, n3, s3] "flat_index_3d" = Just $ \desc -> do       flatIndexHelper desc loc arr offset [(n1, s1), (n2, s2), (n3, s3)]-    handleRest [TupLit [arr1, offset, s1, s2, s3, arr2] _] "flat_update_3d" = Just $ \desc -> do+    handleRest [arr1, offset, s1, s2, s3, arr2] "flat_update_3d" = Just $ \desc -> do       flatUpdateHelper desc loc arr1 offset [s1, s2, s3] arr2-    handleRest [TupLit [arr, offset, n1, s1, n2, s2, n3, s3, n4, s4] _] "flat_index_4d" = Just $ \desc -> do+    handleRest [arr, offset, n1, s1, n2, s2, n3, s3, n4, s4] "flat_index_4d" = Just $ \desc -> do       flatIndexHelper desc loc arr offset [(n1, s1), (n2, s2), (n3, s3), (n4, s4)]-    handleRest [TupLit [arr1, offset, s1, s2, s3, s4, arr2] _] "flat_update_4d" = Just $ \desc -> do+    handleRest [arr1, offset, s1, s2, s3, s4, arr2] "flat_update_4d" = Just $ \desc -> do       flatUpdateHelper desc loc arr1 offset [s1, s2, s3, s4] arr2     handleRest _ _ = Nothing @@ -2143,7 +2147,7 @@           (resultBodyM [this_one])           (resultBodyM [next_one]) -typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp]+typeExpForError :: E.TypeExp Info VName -> InternaliseM [ErrorMsgPart SubExp] typeExpForError (E.TEVar qn _) =   pure [ErrorString $ prettyText qn] typeExpForError (E.TEUnique te _) =@@ -2153,10 +2157,8 @@   where     dims' = mconcat (map onDim dims)     onDim d = "[" <> prettyText d <> "]"-typeExpForError (E.TEArray d te _) = do-  d' <- dimExpForError d-  te' <- typeExpForError te-  pure $ ["[", d', "]"] ++ te'+typeExpForError (E.TEArray d te _) =+  (<>) <$> sizeExpForError d <*> typeExpForError te typeExpForError (E.TETuple tes _) = do   tes' <- mapM typeExpForError tes   pure $ ["("] ++ intercalate [", "] tes' ++ [")"]@@ -2174,7 +2176,7 @@   t' <- typeExpForError t   arg' <- case arg of     TypeArgExpType argt -> typeExpForError argt-    TypeArgExpDim d _ -> pure <$> dimExpForError d+    TypeArgExpSize d -> sizeExpForError d   pure $ t' ++ [" "] ++ arg' typeExpForError (E.TESum cs _) = do   cs' <- mapM (onClause . snd) cs@@ -2184,16 +2186,11 @@       c' <- mapM typeExpForError c       pure $ intercalate [" "] c' -dimExpForError :: E.SizeExp VName -> InternaliseM (ErrorMsgPart SubExp)-dimExpForError (SizeExpNamed d _) = do-  substs <- lookupSubst $ E.qualLeaf d-  d' <- case substs of-    Just [v] -> pure v-    _ -> pure $ I.Var $ E.qualLeaf d-  pure $ ErrorVal int64 d'-dimExpForError (SizeExpConst d _) =-  pure $ ErrorString $ prettyText d-dimExpForError SizeExpAny = pure ""+sizeExpForError :: E.SizeExp Info VName -> InternaliseM [ErrorMsgPart SubExp]+sizeExpForError (SizeExp e _) = do+  e' <- internaliseExp1 "size" e+  pure ["[", ErrorVal int64 e', "]"]+sizeExpForError SizeExpAny {} = pure ["[]"]  -- A smart constructor that compacts neighbouring literals for easier -- reading in the IR.

@@ -55,9 +55,11 @@  existentials :: Exp -> S.Set VName existentials e =-  let here = case e of-        AppExp (Apply _ _ (Info (_, pdim)) _) (Info res) ->-          S.fromList (maybeToList pdim ++ appResExt res)+  let onArg (Info (_, pdim), _) =+        maybeToList pdim+      here = case e of+        AppExp (Apply _ args _) (Info res) ->+          S.fromList (foldMap onArg args <> appResExt res)         AppExp _ (Info res) ->           S.fromList (appResExt res)         _ ->@@ -129,7 +131,7 @@     apply f [] = f     apply f (p : rem_ps) =       let inner_ret = AppRes (fromStruct (augType rem_ps)) mempty-          inner = AppExp (Apply f (freeVar p) (Info (Observe, Nothing)) mempty) (Info inner_ret)+          inner = mkApply f [(Observe, Nothing, freeVar p)] inner_ret        in apply inner rem_ps  transformExp :: Exp -> LiftM Exp

@@ -32,6 +32,7 @@ import Data.Bitraversable import Data.Foldable import Data.List (partition)+import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe import Data.Sequence qualified as Seq@@ -219,9 +220,10 @@      applySizeArg (i, f) size_arg =       ( i - 1,-        AppExp-          (Apply f size_arg (Info (Observe, Nothing)) loc)-          (Info $ AppRes (foldFunType (replicate i i64) (RetType [] (fromStruct t))) [])+        mkApply+          f+          [(Observe, Nothing, size_arg)]+          (AppRes (foldFunType (replicate i (Observe, i64)) (RetType [] (fromStruct t))) [])       )      applySizeArgs fname' t' size_args =@@ -233,7 +235,7 @@               (qualName fname')               ( Info                   ( foldFunType-                      (map (const i64) size_args)+                      (map (const (Observe, i64)) size_args)                       (RetType [] $ fromStruct t')                   )               )@@ -310,8 +312,13 @@         <*> pure (Info res) transformAppExp (If e1 e2 e3 loc) res =   AppExp <$> (If <$> transformExp e1 <*> transformExp e2 <*> transformExp e3 <*> pure loc) <*> pure (Info res)-transformAppExp (Apply e1 e2 d loc) res =-  AppExp <$> (Apply <$> transformExp e1 <*> transformExp e2 <*> pure d <*> pure loc) <*> pure (Info res)+transformAppExp (Apply fe args _) res =+  mkApply+    <$> transformExp fe+    <*> mapM onArg (NE.toList args)+    <*> pure res+  where+    onArg (Info (d, ext), e) = (d,ext,) <$> transformExp e transformAppExp (DoLoop sparams pat e1 form e3 loc) res = do   e1' <- transformExp e1   form' <- case form of@@ -357,17 +364,10 @@           (Info (AppRes ret mempty))   where     applyOp fname' x y =-      AppExp-        ( Apply-            ( AppExp-                (Apply fname' x (Info (Observe, snd (unInfo d1))) loc)-                (Info $ AppRes ret mempty)-            )-            y-            (Info (Observe, snd (unInfo d2)))-            loc-        )-        (Info (AppRes ret ext))+      mkApply+        (mkApply fname' [(Observe, snd (unInfo d1), x)] (AppRes ret mempty))+        [(Observe, snd (unInfo d2), y)]+        (AppRes ret ext)      makeVarParam arg = do       let argtype = typeOf arg@@ -533,14 +533,10 @@   (v1, wrap_left, e1, p1) <- makeVarParam e_left $ fromStruct xtype   (v2, wrap_right, e2, p2) <- makeVarParam e_right $ fromStruct ytype   let apply_left =-        AppExp-          ( Apply-              op-              e1-              (Info (Observe, xext))-              loc-          )-          (Info $ AppRes (Scalar $ Arrow mempty yp ytype (RetType [] t)) [])+        mkApply+          op+          [(Observe, xext, e1)]+          (AppRes (Scalar $ Arrow mempty yp Observe ytype (RetType [] t)) [])       rettype' =         let onDim (NamedSize d)               | Named p <- xp, qualLeaf d == p = NamedSize $ qualName v1@@ -548,14 +544,7 @@             onDim d = d          in first onDim rettype       body =-        AppExp-          ( Apply-              apply_left-              e2-              (Info (Observe, yext))-              loc-          )-          (Info $ AppRes rettype' retext)+        mkApply apply_left [(Observe, yext, e2)] (AppRes rettype' retext)       rettype'' = toStruct rettype'   pure $     wrap_left $@@ -580,7 +569,7 @@       pure (v, id, var_e, [pat])  desugarProjectSection :: [Name] -> PatType -> SrcLoc -> MonoM Exp-desugarProjectSection fields (Scalar (Arrow _ _ t1 (RetType dims t2))) loc = do+desugarProjectSection fields (Scalar (Arrow _ _ _ t1 (RetType dims t2))) loc = do   p <- newVName "project_p"   let body = foldl project (Var (qualName p) (Info t1') mempty) fields   pure $@@ -606,7 +595,7 @@ desugarProjectSection _ t _ = error $ "desugarOpSection: not a function type: " ++ prettyString t  desugarIndexSection :: [DimIndex] -> PatType -> SrcLoc -> MonoM Exp-desugarIndexSection idxs (Scalar (Arrow _ _ t1 (RetType dims t2))) loc = do+desugarIndexSection idxs (Scalar (Arrow _ _ _ t1 (RetType dims t2))) loc = do   p <- newVName "index_i"   let body = AppExp (Index (Var (qualName p) (Info t1') loc) idxs loc) (Info (AppRes t2 []))   pure $@@ -649,11 +638,11 @@     ) transformPat (Id v t loc) = pure (Id v t loc, mempty) transformPat (TuplePat pats loc) = do-  (pats', rrs) <- unzip <$> mapM transformPat pats+  (pats', rrs) <- mapAndUnzipM transformPat pats   pure (TuplePat pats' loc, mconcat rrs) transformPat (RecordPat fields loc) = do   let (field_names, field_pats) = unzip fields-  (field_pats', rrs) <- unzip <$> mapM transformPat field_pats+  (field_pats', rrs) <- mapAndUnzipM transformPat field_pats   pure (RecordPat (zip field_names field_pats') loc, mconcat rrs) transformPat (PatParens pat loc) = do   (pat', rr) <- transformPat pat@@ -669,7 +658,7 @@   pure (PatAscription pat' td loc, rr) transformPat (PatLit e t loc) = pure (PatLit e t loc, mempty) transformPat (PatConstr name t all_ps loc) = do-  (all_ps', rrs) <- unzip <$> mapM transformPat all_ps+  (all_ps', rrs) <- mapAndUnzipM transformPat all_ps   pure (PatConstr name t all_ps' loc, mconcat rrs)  wildcard :: PatType -> SrcLoc -> Pat@@ -719,8 +708,8 @@   where     f (Array () u shape t) = Array () u shape (f' t)     f (Scalar t) = Scalar $ f' t-    f' (Arrow () _ t1 (RetType dims t2)) =-      Arrow () Unnamed (f t1) (RetType dims (f t2))+    f' (Arrow () _ d1 t1 (RetType dims t2)) =+      Arrow () Unnamed d1 (f t1) (RetType dims (f t2))     f' (Record fs) =       Record $ fmap f fs     f' (Sum cs) =@@ -737,7 +726,7 @@   MonoM (VName, InferSizeArgs, ValBind) monomorphiseBinding entry (PolyBinding rr (name, tparams, params, rettype, body, attrs, loc)) inst_t =   replaceRecordReplacements rr $ do-    let bind_t = foldFunType (map patternStructType params) rettype+    let bind_t = funType params rettype     (substs, t_shape_params) <- typeSubstsM loc (noSizes bind_t) $ noNamedParams inst_t     let substs' = M.map (Subst []) substs         rettype' = applySubst (`M.lookup` substs') rettype@@ -749,7 +738,7 @@           partition ((`S.member` mustBeExplicitInBinding bind_t') . typeParamName) $             shape_params ++ t_shape_params -    (params'', rrs) <- unzip <$> mapM transformPat params'+    (params'', rrs) <- mapAndUnzipM transformPat params'      mapM_ noticeDims $ retType rettype : map patternStructType params'' @@ -840,7 +829,7 @@         (map snd $ sortFields fields1)         (map snd $ sortFields fields2)     sub (Scalar Prim {}) (Scalar Prim {}) = pure ()-    sub (Scalar (Arrow _ _ t1a (RetType _ t1b))) (Scalar (Arrow _ _ t2a t2b)) = do+    sub (Scalar (Arrow _ _ _ t1a (RetType _ t1b))) (Scalar (Arrow _ _ _ t2a t2b)) = do       sub t1a t2a       subRet t1b t2b     sub (Scalar (Sum cs1)) (Scalar (Sum cs2)) =@@ -940,11 +929,10 @@     Nothing -> pure ()     Just (Info entry) -> do       t <--        removeTypeVariablesInType-          $ foldFunType-            (map patternStructType (valBindParams valbind))-          $ unInfo-          $ valBindRetType valbind+        removeTypeVariablesInType $+          funType (valBindParams valbind) $+            unInfo $+              valBindRetType valbind       (name, infer, valbind'') <- monomorphiseBinding True valbind' $ monoType t       entry' <- transformEntryPoint entry       tell $ Seq.singleton (name, valbind'' {valBindEntryPoint = Just $ Info entry'})

@@ -52,7 +52,7 @@     freeVarSubstitutions' subs [] = Just subs     freeVarSubstitutions' subs0 fvs =       let fvs_not_in_scope = filter (`M.notMember` scope0) fvs-       in case unzip <$> mapM getSubstitution fvs_not_in_scope of+       in case mapAndUnzipM getSubstitution fvs_not_in_scope of             -- We require that all free variables can be substituted             Just (subs, new_fvs) ->               freeVarSubstitutions' (subs0 <> mconcat subs) $ concat new_fvs

@@ -4,7 +4,7 @@ -- row-major order.  "Futhark.Pass.ExplicitAllocations" is -- conservative and inserts copies to ensure this is the case.  After -- simplification, it may turn out that those copies are redundant.--- This pass removes them.  It's a prettyString simple pass, as it only has+-- This pass removes them.  It's a pretty simple pass, as it only has -- to look at the top level of entry points. module Futhark.Optimise.EntryPointMem   ( entryPointMemGPU,

@@ -177,7 +177,7 @@     Match ses cases defbody dec -> do       let transformCase (Case vs body) =             first (Case vs) <$> transformBody aliases body-      (cases', cases_deps) <- unzip <$> mapM transformCase cases+      (cases', cases_deps) <- mapAndUnzipM transformCase cases       (defbody', defbody_deps) <- transformBody aliases defbody       let deps = depsOf ses <> mconcat cases_deps <> defbody_deps <> depsOf dec       pure (Match ses cases' defbody' dec, deps)

@@ -1216,9 +1216,7 @@   | isArray t =       do         body_depth <- outermostCopyableArray n-        case body_depth of-          Just bd -> recordCopyableMemory i bd-          Nothing -> pure ()+        forM_ body_depth (recordCopyableMemory i)   | otherwise =       pure () 

@@ -81,7 +81,7 @@   | Just r <- arrayRank <$> ST.lookupType x vtable,     let perm = [i] ++ [0 .. i - 1] ++ [i + 1 .. r - 1],     Just (x', x_cs) <- transposedBy perm x,-    Just (xs', xs_cs) <- unzip <$> mapM (transposedBy perm) xs = Simplify $ do+    Just (xs', xs_cs) <- mapAndUnzipM (transposedBy perm) xs = Simplify $ do       concat_rearrange <-         certifying (x_cs <> mconcat xs_cs) $           letExp "concat_rearrange" $@@ -98,12 +98,8 @@ -- Removing a concatenation that involves only a single array.  This -- may be produced as a result of other simplification rules. simplifyConcat _ pat aux (Concat _ (x :| []) _) =-  Simplify $-    -- Still need a copy because Concat produces a fresh array.-    auxing aux $-      letBind pat $-        BasicOp $-          Copy x+  -- Still need a copy because Concat produces a fresh array.+  Simplify $ auxing aux $ letBind pat $ BasicOp $ Copy x -- concat xs (concat ys zs) == concat xs ys zs simplifyConcat (vtable, _) pat (StmAux cs attrs _) (Concat i (x :| xs) new_d)   | x' /= x || concat xs' /= xs =

@@ -144,7 +144,7 @@     onOp op lam = op {histOp = lam} transformExp (WithAcc inputs lam) = do   lam' <- transformLambda lam-  (input_alloc_stms, inputs') <- unzip <$> mapM onInput inputs+  (input_alloc_stms, inputs') <- mapAndUnzipM onInput inputs   pure     ( mconcat input_alloc_stms,       WithAcc inputs' lam'@@ -533,7 +533,7 @@    -- We expand the invariant allocations by adding an inner dimension   -- equal to the sum of the sizes required by different threads.-  (alloc_stms, rebases) <- unzip <$> mapM expand variant_allocs'+  (alloc_stms, rebases) <- mapAndUnzipM expand variant_allocs'    pure (slice_stms' <> stmsFromList alloc_stms, mconcat rebases)   where

@@ -906,7 +906,7 @@     num_arrays = length $ filter ((> 0) . arrayRank . declExtTypeOf) rettype allocInExp (Match ses cases defbody (MatchDec rets ifsort)) = do   (defbody', def_reqs) <- allocInMatchBody rets defbody-  (cases', cases_reqs) <- unzip <$> mapM onCase cases+  (cases', cases_reqs) <- mapAndUnzipM onCase cases   let reqs = zipWith (foldl combMemReqTypes) def_reqs (transpose cases_reqs)   defbody'' <- addCtxToMatchBody reqs defbody'   cases'' <- mapM (traverse $ addCtxToMatchBody reqs) cases'

@@ -191,7 +191,7 @@ transformRedomap rename onBody w reds map_lam arrs = do   (gtid, space) <- mkSegSpace w   kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs-  (reds_stms, reds') <- unzip <$> mapM reduceToSegBinOp reds+  (reds_stms, reds') <- mapAndUnzipM reduceToSegBinOp reds   op' <-     renameIfNeeded rename $       SegRed () space reds' (lambdaReturnType map_lam) kbody@@ -208,7 +208,7 @@ transformHist rename onBody w hists map_lam arrs = do   (gtid, space) <- mkSegSpace w   kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs-  (hists_stms, hists') <- unzip <$> mapM histToSegBinOp hists+  (hists_stms, hists') <- mapAndUnzipM histToSegBinOp hists   op' <-     renameIfNeeded rename $       SegHist () space hists' (lambdaReturnType map_lam) kbody@@ -244,7 +244,7 @@   | Just (scans, map_lam) <- isScanomapSOAC form = do       (gtid, space) <- mkSegSpace w       kbody <- mapLambdaToKernelBody transformBody gtid map_lam arrs-      (scans_stms, scans') <- unzip <$> mapM scanToSegBinOp scans+      (scans_stms, scans') <- mapAndUnzipM scanToSegBinOp scans       pure $         mconcat scans_stms           <> oneStm

@@ -5,6 +5,7 @@ module Futhark.Tools   ( module Futhark.Construct,     redomapToMapAndReduce,+    scanomapToMapAndScan,     dissectScrema,     sequentialStreamWholeArray,     partitionChunkedFoldParameters,@@ -47,6 +48,28 @@     Let red_pat (defAux ()) . Op       <$> (Screma w red_arrs <$> reduceSOAC reds)   pure (map_stm, red_stm)++scanomapToMapAndScan ::+  ( MonadFreshNames m,+    Buildable rep,+    ExpDec rep ~ (),+    Op rep ~ SOAC rep+  ) =>+  Pat (LetDec rep) ->+  ( SubExp,+    [Scan rep],+    Lambda rep,+    [VName]+  ) ->+  m (Stm rep, Stm rep)+scanomapToMapAndScan (Pat pes) (w, scans, map_lam, arrs) = do+  (map_pat, scan_pat, scan_arrs) <-+    splitScanOrRedomap pes w map_lam $ map scanNeutral scans+  let map_stm = mkLet map_pat $ Op $ Screma w arrs (mapSOAC map_lam)+  scan_stm <-+    Let scan_pat (defAux ()) . Op+      <$> (Screma w scan_arrs <$> scanSOAC scans)+  pure (map_stm, scan_stm)  splitScanOrRedomap ::   (Typed dec, MonadFreshNames m) =>

@@ -57,6 +57,7 @@ import Control.Concurrent import Control.Exception import Control.Monad+import Control.Monad.State import Crypto.Hash.MD5 as MD5 import Data.ByteString qualified as BS import Data.ByteString.Base16 qualified as Base16@@ -96,18 +97,23 @@   where     eq x y = cmp x y == EQ --- | Like 'Data.Traversable.mapAccumL', but monadic.+-- | Like 'Data.Traversable.mapAccumL', but monadic and generalised to+-- any 'Traversable'. mapAccumLM ::-  Monad m =>+  (Monad m, Traversable t) =>   (acc -> x -> m (acc, y)) ->   acc ->-  [x] ->-  m (acc, [y])-mapAccumLM _ acc [] = pure (acc, [])-mapAccumLM f acc (x : xs) = do-  (acc', x') <- f acc x-  (acc'', xs') <- mapAccumLM f acc' xs-  pure (acc'', x' : xs')+  t x ->+  m (acc, t y)+mapAccumLM op initial l = do+  (l', acc) <- runStateT (traverse f l) initial+  pure (acc, l')+  where+    f x = do+      acc <- get+      (acc', y) <- lift $ op acc x+      put acc'+      pure y  -- | @chunk n a@ splits @a@ into @n@-size-chunks.  If the length of -- @a@ is not divisible by @n@, the last chunk will have fewer than

@@ -76,7 +76,7 @@     )       <> (freeInExp e2 `freeWithout` S.singleton vn)   AppExp (If e1 e2 e3 _) _ -> freeInExp e1 <> freeInExp e2 <> freeInExp e3-  AppExp (Apply e1 e2 _ _) _ -> freeInExp e1 <> freeInExp e2+  AppExp (Apply f args _) _ -> freeInExp f <> foldMap (freeInExp . snd) args   Negate e _ -> freeInExp e   Not e _ -> freeInExp e   Lambda pats e0 _ (Info (_, RetType dims t)) _ ->@@ -148,7 +148,7 @@       mempty     Scalar (Sum cs) ->       foldMap (foldMap freeInType) cs-    Scalar (Arrow _ v t1 (RetType dims t2)) ->+    Scalar (Arrow _ v _ t1 (RetType dims t2)) ->       S.filter (notV v) $ S.filter (`notElem` dims) $ freeInType t1 <> freeInType t2     Scalar (TypeVar _ _ _ targs) ->       foldMap typeArgDims targs

@@ -93,7 +93,7 @@ newtype EvalM a   = EvalM       ( ReaderT-          (Stack, M.Map FilePath Env)+          (Stack, M.Map ImportName Env)           (StateT Sizes (F ExtOp))           a       )@@ -102,11 +102,11 @@       Applicative,       Functor,       MonadFree ExtOp,-      MonadReader (Stack, M.Map FilePath Env),+      MonadReader (Stack, M.Map ImportName Env),       MonadState Sizes     ) -runEvalM :: M.Map FilePath Env -> EvalM a -> F ExtOp a+runEvalM :: M.Map ImportName Env -> EvalM a -> F ExtOp a runEvalM imports (EvalM m) = evalStateT (runReaderT m (mempty, imports)) mempty  stacking :: SrcLoc -> Env -> EvalM a -> EvalM a@@ -123,7 +123,7 @@ stacktrace :: EvalM [Loc] stacktrace = asks $ map stackFrameLoc . fst -lookupImport :: FilePath -> EvalM (Maybe Env)+lookupImport :: ImportName -> EvalM (Maybe Env) lookupImport f = asks $ M.lookup f . snd  putExtSize :: VName -> Int64 -> EvalM ()@@ -154,8 +154,8 @@           M.elems $             M.intersectionWith (zipWith match) poly_fields fields     match-      (Scalar (Arrow _ _ poly_t1 (RetType _ poly_t2)))-      (Scalar (Arrow _ _ t1 (RetType _ t2))) =+      (Scalar (Arrow _ _ _ poly_t1 (RetType _ poly_t2)))+      (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =         match poly_t1 t1 <> match poly_t2 t2     match poly_t t       | d1 : _ <- shapeDims (arrayShape poly_t),@@ -547,8 +547,8 @@ evalType :: Env -> StructType -> StructType evalType _ (Scalar (Prim pt)) = Scalar $ Prim pt evalType env (Scalar (Record fs)) = Scalar $ Record $ fmap (evalType env) fs-evalType env (Scalar (Arrow () p t1 (RetType dims t2))) =-  Scalar $ Arrow () p (evalType env t1) (RetType dims (evalType env t2))+evalType env (Scalar (Arrow () p d t1 (RetType dims t2))) =+  Scalar $ Arrow () p d (evalType env t1) (RetType dims (evalType env t2)) evalType env t@(Array _ u shape _) =   let et = stripArray (shapeRank shape) t       et' = evalType env et@@ -611,7 +611,7 @@   -- Eta-expand the rest to make any sizes visible.   etaExpand [] env rettype   where-    etaExpand vs env' (Scalar (Arrow _ _ pt (RetType _ rt))) =+    etaExpand vs env' (Scalar (Arrow _ _ _ pt (RetType _ rt))) =       pure $         ValueFun $ \v -> do           env'' <- matchPat env' (Wildcard (Info $ fromStruct pt) noLoc) v@@ -640,7 +640,7 @@   EvalM TermBinding evalFunctionBinding env tparams ps ret fbody = do   let ret' = evalType env $ retType ret-      arrow (xp, xt) yt = Scalar $ Arrow () xp xt $ RetType [] yt+      arrow (xp, d, xt) yt = Scalar $ Arrow () xp d xt $ RetType [] yt       ftype = foldr (arrow . patternParam) ret' ps       retext = case ps of         [] -> retDims ret@@ -778,12 +778,15 @@ evalAppExp env _ (If cond e1 e2 _) = do   cond' <- asBool <$> eval env cond   if cond' then eval env e1 else eval env e2-evalAppExp env _ (Apply f x (Info (_, ext)) loc) = do-  -- It is important that 'x' is evaluated first in order to bring any-  -- sizes into scope that may be used in the type of 'f'.-  x' <- evalArg env x ext+evalAppExp env _ (Apply f args loc) = do+  -- It is important that 'arguments' are evaluated in reverse order+  -- in order to bring any sizes into scope that may be used in the+  -- type of the functions.+  args' <- reverse <$> mapM evalArg' (reverse $ NE.toList args)   f' <- eval env f-  apply loc env f' x'+  foldM (apply loc env) f' args'+  where+    evalArg' (Info (_, ext), x) = evalArg env x ext evalAppExp env _ (Index e is loc) = do   is' <- mapM (evalDimIndex env) is   arr <- eval env e@@ -1119,7 +1122,7 @@ -- is how the REPL works. data Ctx = Ctx   { ctxEnv :: Env,-    ctxImports :: M.Map FilePath Env+    ctxImports :: M.Map ImportName Env   }  nanValue :: PrimValue -> Bool@@ -1231,58 +1234,62 @@      fun1 f =       TermValue Nothing $ ValueFun $ \x -> f x+     fun2 f =-      TermValue Nothing $-        ValueFun $ \x ->-          pure $ ValueFun $ \y -> f x y-    fun2t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y] -> f x y-            _ -> error $ "Expected pair; got: " <> show v-    fun3t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z] -> f x y z-            _ -> error $ "Expected triple; got: " <> show v+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y -> f x y -    fun5t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z, a, b] -> f x y z a b-            _ -> error $ "Expected pentuple; got: " <> show v+    fun3 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z -> f x y z -    fun6t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z, a, b, c] -> f x y z a b c-            _ -> error $ "Expected sextuple; got: " <> show v+    fun5 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z ->+            pure . ValueFun $ \a ->+              pure . ValueFun $ \b -> f x y z a b -    fun7t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z, a, b, c, d] -> f x y z a b c d-            _ -> error $ "Expected septuple; got: " <> show v+    fun6 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z ->+            pure . ValueFun $ \a ->+              pure . ValueFun $ \b ->+                pure . ValueFun $ \c -> f x y z a b c -    fun8t f =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z, a, b, c, d, e] -> f x y z a b c d e-            _ -> error $ "Expected sextuple; got: " <> show v+    fun7 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z ->+            pure . ValueFun $ \a ->+              pure . ValueFun $ \b ->+                pure . ValueFun $ \c ->+                  pure . ValueFun $ \d -> f x y z a b c d -    fun10t fun =-      TermValue Nothing $-        ValueFun $ \v ->-          case fromTuple v of-            Just [x, y, z, a, b, c, d, e, f, g] -> fun x y z a b c d e f g-            _ -> error $ "Expected octuple; got: " <> show v+    fun8 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z ->+            pure . ValueFun $ \a ->+              pure . ValueFun $ \b ->+                pure . ValueFun $ \c ->+                  pure . ValueFun $ \d ->+                    pure . ValueFun $ \e -> f x y z a b c d e +    fun10 f =+      TermValue Nothing . ValueFun $ \x ->+        pure . ValueFun $ \y ->+          pure . ValueFun $ \z ->+            pure . ValueFun $ \a ->+              pure . ValueFun $ \b ->+                pure . ValueFun $ \c ->+                  pure . ValueFun $ \d ->+                    pure . ValueFun $ \e ->+                      pure . ValueFun $ \g ->+                        pure . ValueFun $ \h -> f x y z a b c d e g h+     bopDef fs = fun2 $ \x y ->       case (x, y) of         (ValuePrim x', ValuePrim y')@@ -1459,14 +1466,14 @@                 _ -> error $ "Cannot unsign: " <> show x     def s       | "map_stream" `isPrefixOf` s =-          Just $ fun2t stream+          Just $ fun2 stream     def s | "reduce_stream" `isPrefixOf` s =-      Just $ fun3t $ \_ f arg -> stream f arg+      Just $ fun3 $ \_ f arg -> stream f arg     def "map" = Just $       TermPoly Nothing $ \t -> pure $-        ValueFun $ \v ->-          case (fromTuple v, unfoldFunType t) of-            (Just [f, xs], ([_], ret_t))+        ValueFun $ \f -> pure . ValueFun $ \xs ->+          case unfoldFunType t of+            ([_, _], ret_t)               | Just rowshape <- typeRowShape ret_t ->                   toArray' rowshape <$> mapM (apply noLoc mempty f) (snd $ fromArray xs)               | otherwise ->@@ -1474,21 +1481,21 @@             _ ->               error $                 "Invalid arguments to map intrinsic:\n"-                  ++ unlines [prettyString t, show v]+                  ++ unlines [prettyString t, show f, show xs]       where         typeRowShape = sequenceA . structTypeShape mempty . stripArray 1     def s | "reduce" `isPrefixOf` s = Just $-      fun3t $ \f ne xs ->+      fun3 $ \f ne xs ->         foldM (apply2 noLoc mempty f) ne $ snd $ fromArray xs     def "scan" = Just $-      fun3t $ \f ne xs -> do+      fun3 $ \f ne xs -> do         let next (out, acc) x = do               x' <- apply2 noLoc mempty f acc x               pure (x' : out, x')         toArray' (valueShape ne) . reverse . fst           <$> foldM next ([], ne) (snd $ fromArray xs)     def "scatter" = Just $-      fun3t $ \arr is vs ->+      fun3 $ \arr is vs ->         case arr of           ValueArray shape arr' ->             pure $@@ -1503,7 +1510,7 @@             then arr' // [(i, v)]             else arr'     def "scatter_2d" = Just $-      fun3t $ \arr is vs ->+      fun3 $ \arr is vs ->         case arr of           ValueArray _ _ ->             pure $@@ -1518,7 +1525,7 @@         update _ _ =           error "scatter_2d expects 2-dimensional indices"     def "scatter_3d" = Just $-      fun3t $ \arr is vs ->+      fun3 $ \arr is vs ->         case arr of           ValueArray _ _ ->             pure $@@ -1532,7 +1539,7 @@           fromMaybe arr $ writeArray (map (IndexingFix . asInt64) idxs) arr v         update _ _ =           error "scatter_3d expects 3-dimensional indices"-    def "hist_1d" = Just . fun6t $ \_ arr fun _ is vs ->+    def "hist_1d" = Just . fun6 $ \_ arr fun _ is vs ->       foldM         (update fun)         arr@@ -1541,7 +1548,7 @@         op = apply2 mempty mempty         update fun arr (i, v) =           fromMaybe arr <$> updateArray (op fun) [IndexingFix i] arr v-    def "hist_2d" = Just . fun6t $ \_ arr fun _ is vs ->+    def "hist_2d" = Just . fun6 $ \_ arr fun _ is vs ->       foldM         (update fun)         arr@@ -1553,7 +1560,7 @@             <$> updateArray (op fun) (map (IndexingFix . asInt64) idxs) arr v         update _ _ _ =           error "hist_2d: bad index value"-    def "hist_3d" = Just . fun6t $ \_ arr fun _ is vs ->+    def "hist_3d" = Just . fun6 $ \_ arr fun _ is vs ->       foldM         (update fun)         arr@@ -1566,7 +1573,7 @@         update _ _ _ =           error "hist_2d: bad index value"     def "partition" = Just $-      fun3t $ \k f xs -> do+      fun3 $ \k f xs -> do         let (ShapeDim _ rowshape, xs') = fromArray xs              next outs x = do@@ -1586,7 +1593,7 @@         insertAt i x (l : ls) = l : insertAt (i - 1) x ls         insertAt _ _ ls = ls     def "scatter_stream" = Just $-      fun3t $ \dest f vs ->+      fun3 $ \dest f vs ->         case (dest, vs) of           ( ValueArray dest_shape dest_arr,             ValueArray _ vs_arr@@ -1601,7 +1608,7 @@           _ ->             error $ "scatter_stream expects array, but got: " <> prettyString (show vs, show vs)     def "hist_stream" = Just $-      fun5t $ \dest op _ne f vs ->+      fun5 $ \dest op _ne f vs ->         case (dest, vs) of           ( ValueArray dest_shape dest_arr,             ValueArray _ vs_arr@@ -1616,7 +1623,7 @@           _ ->             error $ "hist_stream expects array, but got: " <> prettyString (show dest, show vs)     def "acc_write" = Just $-      fun3t $ \acc i v ->+      fun3 $ \acc i v ->         case (acc, i) of           ( ValueAcc op acc_arr,             ValuePrim (SignedValue (Int64Value i'))@@ -1630,7 +1637,7 @@           _ ->             error $ "acc_write invalid arguments: " <> prettyString (show acc, show i, show v)     ---    def "flat_index_2d" = Just . fun6t $ \arr offset n1 s1 n2 s2 -> do+    def "flat_index_2d" = Just . fun6 $ \arr offset n1 s1 n2 s2 -> do       let offset' = asInt64 offset           n1' = asInt64 n1           n2' = asInt64 n2@@ -1649,7 +1656,7 @@           bad mempty mempty $             "Index out of bounds: " <> prettyText [((n1', s1'), (n2', s2'))]     ---    def "flat_update_2d" = Just . fun5t $ \arr offset s1 s2 v -> do+    def "flat_update_2d" = Just . fun5 $ \arr offset s1 s2 v -> do       let offset' = asInt64 offset           s1' = asInt64 s1           s2' = asInt64 s2@@ -1666,7 +1673,7 @@                 "Index out of bounds: " <> prettyText [((n1, s1'), (n2, s2'))]         s -> error $ "flat_update_2d: invalid arg shape: " ++ show s     ---    def "flat_index_3d" = Just . fun8t $ \arr offset n1 s1 n2 s2 n3 s3 -> do+    def "flat_index_3d" = Just . fun8 $ \arr offset n1 s1 n2 s2 n3 s3 -> do       let offset' = asInt64 offset           n1' = asInt64 n1           n2' = asInt64 n2@@ -1688,7 +1695,7 @@           bad mempty mempty $             "Index out of bounds: " <> prettyText [((n1', s1'), (n2', s2'), (n3', s3'))]     ---    def "flat_update_3d" = Just . fun6t $ \arr offset s1 s2 s3 v -> do+    def "flat_update_3d" = Just . fun6 $ \arr offset s1 s2 s3 v -> do       let offset' = asInt64 offset           s1' = asInt64 s1           s2' = asInt64 s2@@ -1706,7 +1713,7 @@                 "Index out of bounds: " <> prettyText [((n1, s1'), (n2, s2'), (n3, s3'))]         s -> error $ "flat_update_3d: invalid arg shape: " ++ show s     ---    def "flat_index_4d" = Just . fun10t $ \arr offset n1 s1 n2 s2 n3 s3 n4 s4 -> do+    def "flat_index_4d" = Just . fun10 $ \arr offset n1 s1 n2 s2 n3 s3 n4 s4 -> do       let offset' = asInt64 offset           n1' = asInt64 n1           n2' = asInt64 n2@@ -1731,7 +1738,7 @@           bad mempty mempty $             "Index out of bounds: " <> prettyText [(((n1', s1'), (n2', s2')), ((n3', s3'), (n4', s4')))]     ---    def "flat_update_4d" = Just . fun7t $ \arr offset s1 s2 s3 s4 v -> do+    def "flat_update_4d" = Just . fun7 $ \arr offset s1 s2 s3 s4 v -> do       let offset' = asInt64 offset           s1' = asInt64 s1           s2' = asInt64 s2@@ -1762,7 +1769,7 @@         fromPair (Just [x, y]) = (x, y)         fromPair _ = error "Not a pair"     def "zip" = Just $-      fun2t $ \xs ys -> do+      fun2 $ \xs ys -> do         let ShapeDim _ xs_rowshape = valueShape xs             ShapeDim _ ys_rowshape = valueShape ys         pure $@@ -1770,7 +1777,7 @@             map toTuple $               transpose [snd $ fromArray xs, snd $ fromArray ys]     def "concat" = Just $-      fun2t $ \xs ys -> do+      fun2 $ \xs ys -> do         let (ShapeDim _ rowshape, xs') = fromArray xs             (_, ys') = fromArray ys         pure $ toArray' rowshape $ xs' ++ ys'@@ -1784,7 +1791,7 @@               genericTake m $                 transpose (map (snd . fromArray) xs') ++ repeat []     def "rotate" = Just $-      fun2t $ \i xs -> do+      fun2 $ \i xs -> do         let (shape, xs') = fromArray xs         pure $           let idx = if null xs' then 0 else rem (asInt i) (length xs')@@ -1800,7 +1807,7 @@         let (ShapeDim n (ShapeDim m shape), xs') = fromArray xs         pure $ toArray (ShapeDim (n * m) shape) $ concatMap (snd . fromArray) xs'     def "unflatten" = Just $-      fun3t $ \n m xs -> do+      fun3 $ \n m xs -> do         let (ShapeDim xs_size innershape, xs') = fromArray xs             rowshape = ShapeDim (asInt64 m) innershape             shape = ShapeDim (asInt64 n) rowshape@@ -1816,10 +1823,10 @@                 <> "]"           else pure $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) xs'     def "vjp2" = Just $-      fun3t $+      fun3 $         \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff."     def "jvp2" = Just $-      fun3t $+      fun3 $         \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff."     def "acc" = Nothing     def s | nameFromString s `M.member` namesToPrimTypes = Nothing@@ -1848,7 +1855,7 @@     pure $ env <> sizes   pure ctx {ctxEnv = env} -interpretImport :: Ctx -> (FilePath, Prog) -> F ExtOp Ctx+interpretImport :: Ctx -> (ImportName, Prog) -> F ExtOp Ctx interpretImport ctx (fp, prog) = do   env <- runEvalM (ctxImports ctx) $ foldM evalDec (ctxEnv ctx) $ progDecs prog   pure ctx {ctxImports = M.insert fp env $ ctxImports ctx}@@ -1878,7 +1885,7 @@  checkEntryArgs :: VName -> [V.Value] -> StructType -> Either T.Text () checkEntryArgs entry args entry_t-  | args_ts == param_ts =+  | args_ts == map snd param_ts =       pure ()   | otherwise =       Left . docText $@@ -1916,9 +1923,9 @@       f <- evalTermVar (ctxEnv ctx) (qualName fname) ft       foldM (apply noLoc mempty) f vs'   where-    updateType (vt : vts) (Scalar (Arrow als u pt (RetType dims rt))) = do+    updateType (vt : vts) (Scalar (Arrow als pn d pt (RetType dims rt))) = do       checkInput vt pt-      Scalar . Arrow als u (valueStructType vt) . RetType dims <$> updateType vts rt+      Scalar . Arrow als pn d (valueStructType vt) . RetType dims <$> updateType vts rt     updateType _ t =       Right t 

@@ -140,8 +140,7 @@           <+> align (pretty index)       where         index = AppExp (Index e (is ++ map DimFix xs) xloc) NoInfo-    op f x =-      pure $ AppExp (Apply f x NoInfo (srcspan f x)) NoInfo+    op f x = pure $ mkApplyUT f x  patternExp :: UncheckedPat -> ParserMonad UncheckedExp patternExp (Id v _ loc) = pure $ Var (qualName v) NoInfo loc

@@ -237,7 +237,7 @@                                 in SigArrow (Just name) $4 $7 (srcspan $1 $>) }         | SigExp '->' SigExp  { SigArrow Nothing $1 $3 (srcspan $1 $>) } -TypeRef :: { TypeRefBase Name }+TypeRef :: { TypeRefBase NoInfo Name }          : QualName TypeParams '=' TypeExpTerm            { TypeRef (fst $1) $2 $4 (srcspan (snd $1) $>) } @@ -466,18 +466,6 @@          | SumClauses %prec sumprec            { let (cs, loc) = $1 in TESum cs (srclocOf loc) } -         -- Errors-         | '[' SizeExp ']' %prec bottom-           {% parseErrorAt (srcspan $1 $>) $ Just $-                T.unlines ["Missing array row type.",-                           "Did you mean []"  <> prettyText $2 <> "?"]-           }-         | '...[' SizeExp ']' %prec bottom-           {% parseErrorAt (srcspan $1 $>) $ Just $-                T.unlines ["Missing array row type.",-                           "Did you mean []"  <> prettyText $2 <> "?"]-           }- SumClauses :: { ([(Name, [UncheckedTypeExp])], Loc) }             : SumClauses '|' SumClause %prec sumprec               { let (cs, loc1) = $1; (c, ts, loc2) = $3@@ -505,20 +493,16 @@              | '(' TypeExp ',' TupleTypes ')' { TETuple ($2:$4) (srcspan $1 $>) }              | '{' '}'                        { TERecord [] (srcspan $1 $>) }              | '{' FieldTypes1 '}'            { TERecord $2 (srcspan $1 $>) }-             | '[' SizeExp ']' TypeExpTerm-               { TEArray $2 $4 (srcspan $1 $>) }-             | '...[' SizeExp ']' TypeExpTerm-               { TEArray $2 $4 (srcspan $1 $>) }+             | SizeExp TypeExpTerm+               { TEArray $1 $2 (srcspan $1 $>) }              | QualName                       { TEVar (fst $1) (srclocOf (snd $1)) }  Constr :: { (Name, Loc) }         : constructor { let L _ (CONSTRUCTOR c) = $1 in (c, locOf $1) } -TypeArg :: { TypeArgExp Name }-         : '[' SizeExp ']' %prec top-           { TypeArgExpDim $2 (srcspan $1 $>) }-         | '...[' SizeExp ']' %prec top-           { TypeArgExpDim $2 (srcspan $1 $>) }+TypeArg :: { TypeArgExp NoInfo Name }+         : SizeExp %prec top+           { TypeArgExpSize $1 }          | TypeExpAtom            { TypeArgExpType $1 } @@ -533,17 +517,12 @@             : TypeExp                { [$1] }             | TypeExp ',' TupleTypes { $1 : $3 } -SizeExp :: { SizeExp Name }-        : QualName-          { SizeExpNamed (fst $1) (srclocOf (snd $1)) }-        | intlit-          { let L loc (INTLIT n) = $1-            in SizeExpConst (fromIntegral n) (srclocOf loc) }-        | natlit-          { let L loc (NATLIT _ n) = $1-            in SizeExpConst (fromIntegral n) (srclocOf loc) }-        |-          { SizeExpAny }++SizeExp :: { SizeExp NoInfo Name }+         : '[' Exp ']'    { SizeExp $2 (srcspan $1 $>) }+         | '['     ']'    { SizeExpAny (srcspan $1 $>) }+         | '...[' Exp ']' { SizeExp $2 (srcspan $1 $>) }+         | '...['     ']' { SizeExpAny (srcspan $1 $>) }  FunParam :: { PatBase NoInfo Name } FunParam : InnerPat { $1 }

@@ -86,10 +86,9 @@   pretty (NamedSize v) = pretty v   pretty (ConstSize n) = pretty n -instance IsName vn => Pretty (SizeExp vn) where-  pretty SizeExpAny = mempty-  pretty (SizeExpNamed v _) = pretty v-  pretty (SizeExpConst n _) = pretty n+instance (Eq vn, IsName vn, Annot f) => Pretty (SizeExp f vn) where+  pretty SizeExpAny {} = brackets mempty+  pretty (SizeExp e _) = brackets $ pretty e  instance Pretty (Shape Size) where   pretty (Shape ds) = mconcat (map (brackets . pretty) ds)@@ -115,6 +114,10 @@ instance Pretty (Shape dim) => Pretty (RetTypeBase dim as) where   pretty = prettyRetType 0 +instance Pretty Diet where+  pretty Consume = "*"+  pretty Observe = ""+ prettyScalarType :: Pretty (Shape dim) => Int -> ScalarTypeBase dim as -> Doc a prettyScalarType _ (Prim et) = pretty et prettyScalarType p (TypeVar _ u v targs) =@@ -128,11 +131,16 @@   where     ppField (name, t) = pretty (nameToString name) <> colon <+> align (pretty t)     fs' = map ppField $ M.toList fs-prettyScalarType p (Arrow _ (Named v) t1 t2) =+prettyScalarType p (Arrow _ (Named v) d t1 t2) =   parensIf (p > 1) $-    parens (prettyName v <> colon <+> align (pretty t1)) <+> "->" <+> prettyRetType 1 t2-prettyScalarType p (Arrow _ Unnamed t1 t2) =-  parensIf (p > 1) $ prettyType 2 t1 <+> "->" <+> prettyRetType 1 t2+    parens (prettyName v <> colon <+> pretty d <> align (pretty t1))+      <+> "->"+      <+> prettyRetType 1 t2+prettyScalarType p (Arrow _ Unnamed d t1 t2) =+  parensIf (p > 1) $+    (pretty d <> prettyType 2 t1)+      <+> "->"+      <+> prettyRetType 1 t2 prettyScalarType p (Sum cs) =   parensIf (p > 0) $     group (align (mconcat $ punctuate (" |" <> line) cs'))@@ -159,9 +167,9 @@ instance Pretty (TypeArg Size) where   pretty = prettyTypeArg 0 -instance (Eq vn, IsName vn) => Pretty (TypeExp vn) where+instance (Eq vn, IsName vn, Annot f) => Pretty (TypeExp f vn) where   pretty (TEUnique t _) = "*" <> pretty t-  pretty (TEArray d at _) = brackets (pretty d) <> pretty at+  pretty (TEArray d at _) = pretty d <> pretty at   pretty (TETuple ts _) = parens $ commasep $ map pretty ts   pretty (TERecord fs _) = braces $ commasep $ map ppField fs     where@@ -179,8 +187,8 @@   pretty (TEDim dims te _) =     "?" <> mconcat (map (brackets . prettyName) dims) <> "." <> pretty te -instance (Eq vn, IsName vn) => Pretty (TypeArgExp vn) where-  pretty (TypeArgExpDim d _) = brackets $ pretty d+instance (Eq vn, IsName vn, Annot f) => Pretty (TypeArgExp f vn) where+  pretty (TypeArgExpSize d) = pretty d   pretty (TypeArgExpType t) = pretty t  instance IsName vn => Pretty (QualName vn) where@@ -293,8 +301,10 @@     <+> align (pretty t)     </> "else"     <+> align (pretty f)-prettyAppExp p (Apply f arg _ _) =-  parensIf (p >= 10) $ prettyExp 0 f <+> prettyExp 10 arg+prettyAppExp p (Apply f args _) =+  parensIf (p >= 10) $+    prettyExp 0 f+      <+> hsep (map (prettyExp 10 . snd) $ NE.toList args)  instance (Eq vn, IsName vn, Annot f) => Pretty (AppExpBase f vn) where   pretty = prettyAppExp (-1)@@ -353,7 +363,7 @@   parensIf (p /= -1) $     "\\" <> hsep (map pretty params) <> ppAscription rettype       <+> "->"-      </> indent 2 (pretty body)+      </> indent 2 (align (pretty body)) prettyExp _ (OpSection binop _ _) =   parens $ pretty binop prettyExp _ (OpSectionLeft binop _ x _ _ _) =

@@ -48,6 +48,7 @@     orderZero,     unfoldFunType,     foldFunType,+    foldFunTypeFromParams,     typeVars,      -- * Operations on types@@ -174,8 +175,8 @@       Scalar . Sum <$> traverse (traverse (go bound b)) cs     go _ _ (Scalar (Prim t)) =       pure $ Scalar $ Prim t-    go bound _ (Scalar (Arrow als p t1 (RetType dims t2))) =-      Scalar <$> (Arrow als p <$> go bound' PosParam t1 <*> (RetType dims <$> go bound' PosReturn t2))+    go bound _ (Scalar (Arrow als p u t1 (RetType dims t2))) =+      Scalar <$> (Arrow als p u <$> go bound' PosParam t1 <*> (RetType dims <$> go bound' PosReturn t2))       where         bound' =           S.fromList dims@@ -208,16 +209,16 @@ aliases = bifoldMap (const mempty) id  -- | @diet t@ returns a description of how a function parameter of--- type @t@ might consume its argument.+-- type @t@ consumes its argument. diet :: TypeBase shape as -> Diet-diet (Scalar (Record ets)) = RecordDiet $ fmap diet ets+diet (Scalar (Record ets)) = foldl max Observe $ fmap diet ets diet (Scalar (Prim _)) = Observe-diet (Scalar (Arrow _ _ t1 (RetType _ t2))) = FuncDiet (diet t1) (diet t2)+diet (Scalar (Arrow {})) = Observe diet (Array _ Unique _ _) = Consume diet (Array _ Nonunique _ _) = Observe diet (Scalar (TypeVar _ Unique _ _)) = Consume diet (Scalar (TypeVar _ Nonunique _ _)) = Observe-diet (Scalar (Sum cs)) = SumDiet $ M.map (map diet) cs+diet (Scalar (Sum cs)) = foldl max Observe $ foldMap (map diet) cs  -- | Convert any type to one that has rank information, no alias -- information, and no embedded names.@@ -261,7 +262,7 @@   Shape dim ->   TypeBase dim as ->   TypeBase dim as-arrayOf u s t = arrayOfWithAliases mempty u s (t `setUniqueness` Nonunique)+arrayOf = arrayOfWithAliases mempty  arrayOfWithAliases ::   Monoid as =>@@ -334,8 +335,14 @@           M.map             (uncurry $ zipWith combineTypeShapes)             (M.intersectionWith (,) cs1 cs2)-combineTypeShapes (Scalar (Arrow als1 p1 a1 (RetType dims1 b1))) (Scalar (Arrow als2 _p2 a2 (RetType _ b2))) =-  Scalar $ Arrow (als1 <> als2) p1 (combineTypeShapes a1 a2) (RetType dims1 (combineTypeShapes b1 b2))+combineTypeShapes (Scalar (Arrow als1 p1 d1 a1 (RetType dims1 b1))) (Scalar (Arrow als2 _p2 _d2 a2 (RetType _ b2))) =+  Scalar $+    Arrow+      (als1 <> als2)+      p1+      d1+      (combineTypeShapes a1 a2)+      (RetType dims1 (combineTypeShapes b1 b2)) combineTypeShapes (Scalar (TypeVar als1 u1 v targs1)) (Scalar (TypeVar als2 _ _ targs2)) =   Scalar $ TypeVar (als1 <> als2) u1 v $ zipWith f targs1 targs2   where@@ -385,12 +392,12 @@             <$> traverse               (traverse (uncurry (matchDims' bound)))               (M.intersectionWith zip cs1 cs2)-        ( Scalar (Arrow als1 p1 a1 (RetType dims1 b1)),-          Scalar (Arrow als2 p2 a2 (RetType dims2 b2))+        ( Scalar (Arrow als1 p1 d1 a1 (RetType dims1 b1)),+          Scalar (Arrow als2 p2 _d2 a2 (RetType dims2 b2))           ) ->             let bound' = mapMaybe paramName [p1, p2] <> dims1 <> dims2 <> bound              in Scalar-                  <$> ( Arrow (als1 <> als2) p1+                  <$> ( Arrow (als1 <> als2) p1 d1                           <$> matchDims' bound' a1 a2                           <*> (RetType dims1 <$> matchDims' bound' b1 b2)                       )@@ -475,7 +482,7 @@ typeOf (StringLit vs _) =   Array     mempty-    Unique+    Nonunique     (Shape [ConstSize $ genericLength vs])     (Prim (Unsigned Int8)) typeOf (Project _ _ (Info t) _) = t@@ -488,45 +495,62 @@ typeOf (RecordUpdate _ _ _ (Info t) _) = t typeOf (Assert _ e _ _) = typeOf e typeOf (Lambda params _ _ (Info (als, t)) _) =-  let RetType [] t' = foldr (arrow . patternParam) t params-   in t' `setAliases` als-  where-    arrow (Named v, x) (RetType dims y) =-      RetType [] $ Scalar $ Arrow () (Named v) x $ RetType (v : dims) y-    arrow (pn, tx) y =-      RetType [] $ Scalar $ Arrow () pn tx y+  funType params t `setAliases` als typeOf (OpSection _ (Info t) _) =   t typeOf (OpSectionLeft _ _ _ (_, Info (pn, pt2)) (Info ret, _) _) =-  Scalar $ Arrow mempty pn pt2 ret+  Scalar $ Arrow mempty pn Observe pt2 ret typeOf (OpSectionRight _ _ _ (Info (pn, pt1), _) (Info ret) _) =-  Scalar $ Arrow mempty pn pt1 ret+  Scalar $ Arrow mempty pn Observe pt1 ret typeOf (ProjectSection _ (Info t) _) = t typeOf (IndexSection _ (Info t) _) = t typeOf (Constr _ _ (Info t) _) = t typeOf (Attr _ e _) = typeOf e typeOf (AppExp _ (Info res)) = appResType res +-- | The type of a function with the given parameters and return type.+funType :: [PatBase Info VName] -> StructRetType -> StructType+funType params ret =+  let RetType _ t = foldr (arrow . patternParam) ret params+   in t+  where+    arrow (xp, d, xt) yt =+      RetType [] $ Scalar $ Arrow () xp d xt' yt+      where+        xt' = xt `setUniqueness` Nonunique+ -- | @foldFunType ts ret@ creates a function type ('Arrow') that takes -- @ts@ as parameters and returns @ret@. foldFunType ::   Monoid as =>-  [TypeBase dim pas] ->+  [(Diet, TypeBase dim pas)] ->   RetTypeBase dim as ->   TypeBase dim as foldFunType ps ret =   let RetType _ t = foldr arrow ret ps    in t   where-    arrow t1 t2 =-      RetType [] $ Scalar $ Arrow mempty Unnamed (toStruct t1) t2+    arrow (d, t1) t2 =+      RetType [] $ Scalar $ Arrow mempty Unnamed d t1' t2+      where+        t1' = toStruct t1 `setUniqueness` Nonunique +foldFunTypeFromParams ::+  Monoid as =>+  [PatBase Info VName] ->+  RetTypeBase Size as ->+  TypeBase Size as+foldFunTypeFromParams params =+  foldFunType (zip (map diet params_ts) params_ts)+  where+    params_ts = map patternStructType params+ -- | Extract the parameter types and return type from a type. -- If the type is not an arrow type, the list of parameter types is empty.-unfoldFunType :: TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())-unfoldFunType (Scalar (Arrow _ _ t1 (RetType _ t2))) =+unfoldFunType :: TypeBase dim as -> ([(Diet, TypeBase dim ())], TypeBase dim ())+unfoldFunType (Scalar (Arrow _ _ d t1 (RetType _ t2))) =   let (ps, r) = unfoldFunType t2-   in (t1 : ps, r)+   in ((d, t1) : ps, r) unfoldFunType t = ([], toStruct t)  -- | The type scheme of a value binding, comprising the type@@ -547,14 +571,6 @@       [] -> retDims (unInfo (valBindRetType vb))       _ -> [] --- | The type of a function with the given parameters and return type.-funType :: [PatBase Info VName] -> StructRetType -> StructType-funType params ret =-  let RetType _ t = foldr (arrow . patternParam) ret params-   in t-  where-    arrow (xp, xt) yt = RetType [] $ Scalar $ Arrow () xp xt yt- -- | The type names mentioned in a type. typeVars :: Monoid as => TypeBase dim as -> S.Set VName typeVars t =@@ -562,7 +578,7 @@     Scalar Prim {} -> mempty     Scalar (TypeVar _ _ tn targs) ->       mconcat $ S.singleton (qualLeaf tn) : map typeArgFree targs-    Scalar (Arrow _ _ t1 (RetType _ t2)) -> typeVars t1 <> typeVars t2+    Scalar (Arrow _ _ _ t1 (RetType _ t2)) -> typeVars t1 <> typeVars t2     Scalar (Record fields) -> foldMap typeVars fields     Scalar (Sum cs) -> mconcat $ (foldMap . fmap) typeVars cs     Array _ _ _ rt -> typeVars $ Scalar rt@@ -644,17 +660,19 @@  -- | When viewed as a function parameter, does this pattern correspond -- to a named parameter of some type?-patternParam :: PatBase Info VName -> (PName, StructType)+patternParam :: PatBase Info VName -> (PName, Diet, StructType) patternParam (PatParens p _) =   patternParam p patternParam (PatAttr _ p _) =   patternParam p patternParam (PatAscription (Id v (Info t) _) _ _) =-  (Named v, toStruct t)+  (Named v, diet t, toStruct t) patternParam (Id v (Info t) _) =-  (Named v, toStruct t)+  (Named v, diet t, toStruct t) patternParam p =-  (Unnamed, patternStructType p)+  (Unnamed, diet p_t, p_t)+  where+    p_t = patternStructType p  -- | Names of primitive types to types.  This is only valid if no -- shadowing is going on, but useful for tools.@@ -677,7 +695,7 @@ data Intrinsic   = IntrinsicMonoFun [PrimType] PrimType   | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType)-  | IntrinsicPolyFun [TypeParamBase VName] [StructType] (RetTypeBase Size ())+  | IntrinsicPolyFun [TypeParamBase VName] [(Diet, StructType)] (RetTypeBase Size ())   | IntrinsicType Liftedness [TypeParamBase VName] StructType   | IntrinsicEquality -- Special cased. @@ -732,16 +750,16 @@           ++ [ ( "flatten",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m]-                   [Array () Nonunique (shape [n, m]) t_a]+                   [(Observe, Array () Nonunique (shape [n, m]) t_a)]                    $ RetType [k]                    $ Array () Nonunique (shape [k]) t_a                ),                ( "unflatten",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [ Scalar $ Prim $ Signed Int64,-                     Scalar $ Prim $ Signed Int64,-                     Array () Nonunique (shape [n]) t_a+                   [ (Observe, Scalar $ Prim $ Signed Int64),+                     (Observe, Scalar $ Prim $ Signed Int64),+                     (Observe, Array () Nonunique (shape [n]) t_a)                    ]                    $ RetType [k, m]                    $ Array () Nonunique (shape [k, m]) t_a@@ -749,33 +767,37 @@                ( "concat",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m]-                   [arr_a $ shape [n], arr_a $ shape [m]]+                   [ (Observe, array_a $ shape [n]),+                     (Observe, array_a $ shape [m])+                   ]                    $ RetType [k]-                   $ uarr_a+                   $ uarray_a                    $ shape [k]                ),                ( "rotate",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [Scalar $ Prim $ Signed Int64, arr_a $ shape [n]]+                   [ (Observe, Scalar $ Prim $ Signed Int64),+                     (Observe, array_a $ shape [n])+                   ]                    $ RetType []-                   $ arr_a+                   $ array_a                    $ shape [n]                ),                ( "transpose",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m]-                   [arr_a $ shape [n, m]]+                   [(Observe, array_a $ shape [n, m])]                    $ RetType []-                   $ arr_a+                   $ array_a                    $ shape [m, n]                ),                ( "scatter",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_l]-                   [ Array () Unique (shape [n]) t_a,-                     Array () Nonunique (shape [l]) (Prim $ Signed Int64),-                     Array () Nonunique (shape [l]) t_a+                   [ (Consume, Array () Unique (shape [n]) t_a),+                     (Observe, Array () Nonunique (shape [l]) (Prim $ Signed Int64)),+                     (Observe, Array () Nonunique (shape [l]) t_a)                    ]                    $ RetType []                    $ Array () Unique (shape [n]) t_a@@ -783,98 +805,100 @@                ( "scatter_2d",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m, sp_l]-                   [ uarr_a $ shape [n, m],-                     Array () Nonunique (shape [l]) (tupInt64 2),-                     Array () Nonunique (shape [l]) t_a+                   [ (Consume, uarray_a $ shape [n, m]),+                     (Observe, Array () Nonunique (shape [l]) (tupInt64 2)),+                     (Observe, Array () Nonunique (shape [l]) t_a)                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [n, m]                ),                ( "scatter_3d",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m, sp_k, sp_l]-                   [ uarr_a $ shape [n, m, k],-                     Array () Nonunique (shape [l]) (tupInt64 3),-                     Array () Nonunique (shape [l]) t_a+                   [ (Consume, uarray_a $ shape [n, m, k]),+                     (Observe, Array () Nonunique (shape [l]) (tupInt64 3)),+                     (Observe, Array () Nonunique (shape [l]) t_a)                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [n, m, k]                ),                ( "zip",                  IntrinsicPolyFun                    [tp_a, tp_b, sp_n]-                   [arr_a (shape [n]), arr_b (shape [n])]+                   [ (Observe, array_a (shape [n])),+                     (Observe, array_b (shape [n]))+                   ]                    $ RetType []-                   $ tuple_uarr (Scalar t_a) (Scalar t_b)+                   $ tuple_uarray (Scalar t_a) (Scalar t_b)                    $ shape [n]                ),                ( "unzip",                  IntrinsicPolyFun                    [tp_a, tp_b, sp_n]-                   [tuple_arr (Scalar t_a) (Scalar t_b) $ shape [n]]+                   [(Observe, tuple_arr (Scalar t_a) (Scalar t_b) $ shape [n])]                    $ RetType [] . Scalar . Record . M.fromList-                   $ zip tupleFieldNames [arr_a $ shape [n], arr_b $ shape [n]]+                   $ zip tupleFieldNames [array_a $ shape [n], array_b $ shape [n]]                ),                ( "hist_1d",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m]-                   [ Scalar $ Prim $ Signed Int64,-                     uarr_a $ shape [m],-                     Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     Array () Nonunique (shape [n]) (tupInt64 1),-                     arr_a (shape [n])+                   [ (Consume, Scalar $ Prim $ Signed Int64),+                     (Observe, uarray_a $ shape [m]),+                     (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, Array () Nonunique (shape [n]) (tupInt64 1)),+                     (Observe, array_a (shape [n]))                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [m]                ),                ( "hist_2d",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m, sp_k]-                   [ Scalar $ Prim $ Signed Int64,-                     uarr_a $ shape [m, k],-                     Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     Array () Nonunique (shape [n]) (tupInt64 2),-                     arr_a (shape [n])+                   [ (Observe, Scalar $ Prim $ Signed Int64),+                     (Consume, uarray_a $ shape [m, k]),+                     (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, Array () Nonunique (shape [n]) (tupInt64 2)),+                     (Observe, array_a (shape [n]))                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [m, k]                ),                ( "hist_3d",                  IntrinsicPolyFun                    [tp_a, sp_n, sp_m, sp_k, sp_l]-                   [ Scalar $ Prim $ Signed Int64,-                     uarr_a $ shape [m, k, l],-                     Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     Array () Nonunique (shape [n]) (tupInt64 3),-                     arr_a (shape [n])+                   [ (Observe, Scalar $ Prim $ Signed Int64),+                     (Consume, uarray_a $ shape [m, k, l]),+                     (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, Array () Nonunique (shape [n]) (tupInt64 3)),+                     (Observe, array_a (shape [n]))                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [m, k, l]                ),                ( "map",                  IntrinsicPolyFun                    [tp_a, tp_b, sp_n]-                   [ Scalar t_a `arr` Scalar t_b,-                     arr_a $ shape [n]+                   [ (Observe, Scalar t_a `arr` Scalar t_b),+                     (Observe, array_a $ shape [n])                    ]                    $ RetType []-                   $ uarr_b+                   $ uarray_b                    $ shape [n]                ),                ( "reduce",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [ Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     arr_a $ shape [n]+                   [ (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, array_a $ shape [n])                    ]                    $ RetType []                    $ Scalar t_a@@ -882,9 +906,9 @@                ( "reduce_comm",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [ Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     arr_a $ shape [n]+                   [ (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, array_a $ shape [n])                    ]                    $ RetType []                    $ Scalar t_a@@ -892,24 +916,24 @@                ( "scan",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [ Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     arr_a $ shape [n]+                   [ (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     (Observe, array_a $ shape [n])                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [n]                ),                ( "partition",                  IntrinsicPolyFun                    [tp_a, sp_n]-                   [ Scalar (Prim $ Signed Int32),-                     Scalar t_a `arr` Scalar (Prim $ Signed Int64),-                     arr_a $ shape [n]+                   [ (Observe, Scalar (Prim $ Signed Int32)),+                     (Observe, Scalar t_a `arr` Scalar (Prim $ Signed Int64)),+                     (Observe, array_a $ shape [n])                    ]                    ( RetType [m] . Scalar $                        tupleRecord-                         [ uarr_a $ shape [k],+                         [ uarray_a $ shape [k],                            Array () Unique (shape [n]) (Prim $ Signed Int64)                          ]                    )@@ -917,48 +941,52 @@                ( "acc_write",                  IntrinsicPolyFun                    [sp_k, tp_a]-                   [ Scalar $ accType arr_ka,-                     Scalar (Prim $ Signed Int64),-                     Scalar t_a+                   [ (Consume, Scalar $ accType array_ka),+                     (Observe, Scalar (Prim $ Signed Int64)),+                     (Observe, Scalar t_a)                    ]                    $ RetType []                    $ Scalar-                   $ accType arr_ka+                   $ accType array_ka                ),                ( "scatter_stream",                  IntrinsicPolyFun                    [tp_a, tp_b, sp_k, sp_n]-                   [ uarr_ka,-                     Scalar (accType arr_ka)-                       `arr` ( Scalar t_b-                                 `arr` Scalar (accType $ arr_a $ shape [k])-                             ),-                     arr_b $ shape [n]+                   [ (Consume, uarray_ka),+                     ( Observe,+                       Scalar (accType array_ka)+                         `carr` ( Scalar t_b+                                    `arr` Scalar (accType $ array_a $ shape [k])+                                )+                     ),+                     (Observe, array_b $ shape [n])                    ]-                   $ RetType [] uarr_ka+                   $ RetType [] uarray_ka                ),                ( "hist_stream",                  IntrinsicPolyFun                    [tp_a, tp_b, sp_k, sp_n]-                   [ uarr_a $ shape [k],-                     Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a),-                     Scalar t_a,-                     Scalar (accType arr_ka)-                       `arr` ( Scalar t_b-                                 `arr` Scalar (accType $ arr_a $ shape [k])-                             ),-                     arr_b $ shape [n]+                   [ (Consume, uarray_a $ shape [k]),+                     (Observe, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a)),+                     (Observe, Scalar t_a),+                     ( Observe,+                       Scalar (accType array_ka)+                         `carr` ( Scalar t_b+                                    `arr` Scalar (accType $ array_a $ shape [k])+                                )+                     ),+                     (Observe, array_b $ shape [n])                    ]                    $ RetType []-                   $ uarr_a+                   $ uarray_a                    $ shape [k]                ),                ( "jvp2",                  IntrinsicPolyFun                    [tp_a, tp_b]-                   [ Scalar t_a `arr` Scalar t_b,-                     Scalar t_a,-                     Scalar t_a+                   [ (Observe, Scalar t_a `arr` Scalar t_b),+                     (Observe, Scalar t_a),+                     (Observe, Scalar t_a)                    ]                    $ RetType []                    $ Scalar@@ -967,9 +995,9 @@                ( "vjp2",                  IntrinsicPolyFun                    [tp_a, tp_b]-                   [ Scalar t_a `arr` Scalar t_b,-                     Scalar t_a,-                     Scalar t_b+                   [ (Observe, Scalar t_a `arr` Scalar t_b),+                     (Observe, Scalar t_a),+                     (Observe, Scalar t_b)                    ]                    $ RetType []                    $ Scalar@@ -981,91 +1009,91 @@           [ ( "flat_index_2d",               IntrinsicPolyFun                 [tp_a, sp_n]-                [ arr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64)+                [ (Observe, array_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64))                 ]                 $ RetType [m, k]-                $ arr_a+                $ array_a                 $ shape [m, k]             ),             ( "flat_update_2d",               IntrinsicPolyFun                 [tp_a, sp_n, sp_k, sp_l]-                [ uarr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  arr_a $ shape [k, l]+                [ (Consume, uarray_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, array_a $ shape [k, l])                 ]                 $ RetType []-                $ uarr_a+                $ uarray_a                 $ shape [n]             ),             ( "flat_index_3d",               IntrinsicPolyFun                 [tp_a, sp_n]-                [ arr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64)+                [ (Observe, array_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64))                 ]                 $ RetType [m, k, l]-                $ arr_a+                $ array_a                 $ shape [m, k, l]             ),             ( "flat_update_3d",               IntrinsicPolyFun                 [tp_a, sp_n, sp_k, sp_l, sp_p]-                [ uarr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  arr_a $ shape [k, l, p]+                [ (Consume, uarray_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, array_a $ shape [k, l, p])                 ]                 $ RetType []-                $ uarr_a+                $ uarray_a                 $ shape [n]             ),             ( "flat_index_4d",               IntrinsicPolyFun                 [tp_a, sp_n]-                [ arr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64)+                [ (Observe, array_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64))                 ]                 $ RetType [m, k, l, p]-                $ arr_a+                $ array_a                 $ shape [m, k, l, p]             ),             ( "flat_update_4d",               IntrinsicPolyFun                 [tp_a, sp_n, sp_k, sp_l, sp_p, sp_q]-                [ uarr_a $ shape [n],-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  Scalar (Prim $ Signed Int64),-                  arr_a $ shape [k, l, p, q]+                [ (Consume, uarray_a $ shape [n]),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, Scalar (Prim $ Signed Int64)),+                  (Observe, array_a $ shape [k, l, p, q])                 ]                 $ RetType []-                $ uarr_a+                $ uarray_a                 $ shape [n]             )           ]@@ -1073,13 +1101,13 @@     [a, b, n, m, k, l, p, q] = zipWith VName (map nameFromString ["a", "b", "n", "m", "k", "l", "p", "q"]) [0 ..]      t_a = TypeVar () Nonunique (qualName a) []-    arr_a s = Array () Nonunique s t_a-    uarr_a s = Array () Unique s t_a+    array_a s = Array () Nonunique s t_a+    uarray_a s = Array () Unique s t_a     tp_a = TypeParamType Unlifted a mempty      t_b = TypeVar () Nonunique (qualName b) []-    arr_b s = Array () Nonunique s t_b-    uarr_b s = Array () Unique s t_b+    array_b s = Array () Nonunique s t_b+    uarray_b s = Array () Unique s t_b     tp_b = TypeParamType Unlifted b mempty      [sp_n, sp_m, sp_k, sp_l, sp_p, sp_q] = map (`TypeParamDim` mempty) [n, m, k, l, p, q]@@ -1092,12 +1120,13 @@         Nonunique         s         (Record (M.fromList $ zip tupleFieldNames [x, y]))-    tuple_uarr x y s = tuple_arr x y s `setUniqueness` Unique+    tuple_uarray x y s = tuple_arr x y s `setUniqueness` Unique -    arr x y = Scalar $ Arrow mempty Unnamed x (RetType [] y)+    arr x y = Scalar $ Arrow mempty Unnamed Observe x (RetType [] y)+    carr x y = Scalar $ Arrow mempty Unnamed Consume x (RetType [] y) -    arr_ka = Array () Nonunique (Shape [NamedSize $ qualName k]) t_a-    uarr_ka = Array () Unique (Shape [NamedSize $ qualName k]) t_a+    array_ka = Array () Nonunique (Shape [NamedSize $ qualName k]) t_a+    uarray_ka = Array () Unique (Shape [NamedSize $ qualName k]) t_a      accType t =       TypeVar () Unique (qualName (fst intrinsicAcc)) [TypeArgType t mempty]@@ -1204,23 +1233,23 @@ qualify k (QualName ks v) = QualName (k : ks) v  -- | The modules imported by a Futhark program.-progImports :: ProgBase f vn -> [(String, SrcLoc)]+progImports :: ProgBase f vn -> [(String, Loc)] progImports = concatMap decImports . progDecs  -- | The modules imported by a single declaration.-decImports :: DecBase f vn -> [(String, SrcLoc)]+decImports :: DecBase f vn -> [(String, Loc)] decImports (OpenDec x _) = modExpImports x decImports (ModDec md) = modExpImports $ modExp md decImports SigDec {} = [] decImports TypeDec {} = [] decImports ValDec {} = [] decImports (LocalDec d _) = decImports d-decImports (ImportDec x _ loc) = [(x, loc)]+decImports (ImportDec x _ loc) = [(x, locOf loc)] -modExpImports :: ModExpBase f vn -> [(String, SrcLoc)]+modExpImports :: ModExpBase f vn -> [(String, Loc)] modExpImports ModVar {} = [] modExpImports (ModParens p _) = modExpImports p-modExpImports (ModImport f _ loc) = [(f, loc)]+modExpImports (ModImport f _ loc) = [(f, locOf loc)] modExpImports (ModDecs ds _) = concatMap decImports ds modExpImports (ModApply _ me _ _ _) = modExpImports me modExpImports (ModAscript me _ _ _) = modExpImports me@@ -1345,7 +1374,7 @@ type UncheckedType = TypeBase (Shape Name) ()  -- | An expression with no type annotations.-type UncheckedTypeExp = TypeExp Name+type UncheckedTypeExp = TypeExp NoInfo Name  -- | An identifier with no type annotations. type UncheckedIdent = IdentBase NoInfo Name

@@ -88,7 +88,7 @@         Lambda params _ _ _ _ ->           mconcat (map patternDefs params)         AppExp (LetFun name (tparams, params, _, Info ret, _) _ loc) _ ->-          let name_t = foldFunType (map patternStructType params) ret+          let name_t = foldFunType (map (undefined . patternStructType) params) ret            in M.singleton name (DefBound $ BoundTerm name_t (locOf loc))                 <> mconcat (map typeParamDefs tparams)                 <> mconcat (map patternDefs params)@@ -111,7 +111,7 @@       <> expDefs (valBindBody vbind)   where     vbind_t =-      foldFunType (map patternStructType (valBindParams vbind)) $+      foldFunType (map (undefined . patternStructType) (valBindParams vbind)) $         unInfo $           valBindRetType vbind @@ -206,7 +206,7 @@     Loc start end -> pos >= start && pos <= end     NoLoc -> False -atPosInTypeExp :: TypeExp VName -> Pos -> Maybe RawAtPos+atPosInTypeExp :: TypeExp Info VName -> Pos -> Maybe RawAtPos atPosInTypeExp te pos =   case te of     TEVar qn loc -> do@@ -229,12 +229,10 @@     TEDim _ t _ ->       atPosInTypeExp t pos   where-    inArg (TypeArgExpDim dim _) = inDim dim+    inArg (TypeArgExpSize dim) = inDim dim     inArg (TypeArgExpType e2) = atPosInTypeExp e2 pos-    inDim (SizeExpNamed qn loc) = do-      guard $ loc `contains` pos-      Just $ RawAtName qn $ locOf loc-    inDim _ = Nothing+    inDim (SizeExp e _) = atPosInExp e pos+    inDim SizeExpAny {} = Nothing  atPosInPat :: Pat -> Pos -> Maybe RawAtPos atPosInPat (Id vn _ loc) pos = do@@ -370,7 +368,7 @@ containingModule imports (Pos file _ _ _) =   snd <$> find ((== file') . fst) imports   where-    file' = includeToString $ mkInitialImport $ fst $ Posix.splitExtension file+    file' = mkInitialImport $ fst $ Posix.splitExtension file  -- | Information about what is at the given source location. data AtPos = AtName (QualName VName) (Maybe BoundTo) Loc

@@ -24,32 +24,22 @@ import Data.Map.Strict qualified as M import Data.Text qualified as T import Futhark.Util (dropLast, fromPOSIX, toPOSIX)-import Futhark.Util.Loc import Futhark.Util.Pretty import Language.Futhark import System.FilePath qualified as Native import System.FilePath.Posix qualified as Posix import Prelude hiding (mod) --- | Canonical reference to a Futhark code file.  Does not include the--- @.fut@ extension.  This is most often a path relative to the--- current working directory of the compiler.-data ImportName = ImportName Posix.FilePath SrcLoc-  deriving (Eq, Ord, Show)--instance Located ImportName where-  locOf (ImportName _ loc) = locOf loc- -- | Create an import name immediately from a file path specified by -- the user. mkInitialImport :: Native.FilePath -> ImportName-mkInitialImport s = ImportName (Posix.normalise $ toPOSIX s) noLoc+mkInitialImport = ImportName . Posix.normalise . toPOSIX  -- | We resolve '..' paths here and assume that no shenanigans are -- going on with symbolic links.  If there is, too bad.  Don't do -- that.-mkImportFrom :: ImportName -> String -> SrcLoc -> ImportName-mkImportFrom (ImportName includer _) includee+mkImportFrom :: ImportName -> String -> ImportName+mkImportFrom (ImportName includer) includee   | Posix.isAbsolute includee = ImportName includee   | otherwise = ImportName $ Posix.normalise $ Posix.joinPath $ includer' ++ includee'   where@@ -63,17 +53,17 @@  -- | Create a @.fut@ file corresponding to an 'ImportName'. includeToFilePath :: ImportName -> Native.FilePath-includeToFilePath (ImportName s _) = fromPOSIX $ Posix.normalise s Posix.<.> "fut"+includeToFilePath (ImportName s) = fromPOSIX $ Posix.normalise s Posix.<.> "fut"  -- | Produce a human-readable canonicalized string from an -- 'ImportName'. includeToString :: ImportName -> String-includeToString (ImportName s _) = Posix.normalise s+includeToString (ImportName s) = Posix.normalise s  -- | Produce a human-readable canonicalized text from an -- 'ImportName'. includeToText :: ImportName -> T.Text-includeToText (ImportName s _) = T.pack $ Posix.normalise s+includeToText (ImportName s) = T.pack $ Posix.normalise s  -- | The result of type checking some file.  Can be passed to further -- invocations of the type checker.@@ -89,7 +79,7 @@   }  -- | A mapping from import names to imports.  The ordering is significant.-type Imports = [(String, FileModule)]+type Imports = [(ImportName, FileModule)]  -- | The space inhabited by a name. data Namespace

@@ -63,6 +63,7 @@     PatBase (..),      -- * Module language+    ImportName (..),     SpecBase (..),     SigExpBase (..),     TypeRefBase (..),@@ -90,6 +91,8 @@     Alias (..),     Aliasing,     QualName (..),+    mkApply,+    mkApplyUT,   ) where @@ -115,6 +118,7 @@     IntType (..),     IntValue (..),   )+import System.FilePath.Posix qualified as Posix import Prelude  -- | No information functor.  Usually used for placeholder type- or@@ -220,7 +224,7 @@     -- give rise to an assertion.     NamedSize (QualName VName)   | -- | The size is a constant.-    ConstSize Int+    ConstSize Int64   | -- | No known size.  If @Nothing@, then this is a name distinct     -- from any other.  The type checker should _never_ produce these     -- - they are a (hopefully temporary) thing introduced by@@ -298,7 +302,7 @@   | Sum (M.Map Name [TypeBase dim as])   | -- | The aliasing corresponds to the lexical     -- closure of the function.-    Arrow as PName (TypeBase dim ()) (RetTypeBase dim as)+    Arrow as PName Diet (TypeBase dim ()) (RetTypeBase dim as)   deriving (Eq, Ord, Show)  instance Bitraversable ScalarTypeBase where@@ -306,8 +310,8 @@   bitraverse f g (Record fs) = Record <$> traverse (bitraverse f g) fs   bitraverse f g (TypeVar als u t args) =     TypeVar <$> g als <*> pure u <*> pure t <*> traverse (traverse f) args-  bitraverse f g (Arrow als v t1 t2) =-    Arrow <$> g als <*> pure v <*> bitraverse f pure t1 <*> bitraverse f g t2+  bitraverse f g (Arrow als v d t1 t2) =+    Arrow <$> g als <*> pure v <*> pure d <*> bitraverse f pure t1 <*> bitraverse f g t2   bitraverse f g (Sum cs) = Sum <$> (traverse . traverse) (bitraverse f g) cs  instance Bifunctor ScalarTypeBase where@@ -383,49 +387,58 @@ -- | The return type version of 'PatType'. type PatRetType = RetTypeBase Size Aliasing --- | A size expression for use in a 'TypeExp'.-data SizeExp vn-  = -- | The size of the dimension is this name, which-    -- must be in scope.-    SizeExpNamed (QualName vn) SrcLoc-  | -- | The size is a constant.-    SizeExpConst Int SrcLoc+-- | A dimension declaration expression for use in a 'TypeExp'.+-- Syntactically includes the brackets.+data SizeExp f vn+  = -- | The size of the dimension is this expression, all of which+    -- free variables must be in scope.+    SizeExp (ExpBase f vn) SrcLoc   | -- | No dimension declaration.-    SizeExpAny-  deriving (Show)+    SizeExpAny SrcLoc -deriving instance Eq (SizeExp Name)+instance Located (SizeExp f vn) where+  locOf (SizeExp _ loc) = locOf loc+  locOf (SizeExpAny loc) = locOf loc -deriving instance Eq (SizeExp VName)+deriving instance Show (SizeExp Info VName) -deriving instance Ord (SizeExp Name)+deriving instance Show vn => Show (SizeExp NoInfo vn) -deriving instance Ord (SizeExp VName)+deriving instance Eq (SizeExp NoInfo VName) +deriving instance Eq (SizeExp Info VName)++deriving instance Ord (SizeExp NoInfo VName)++deriving instance Ord (SizeExp Info VName)+ -- | An unstructured type with type variables and possibly shape -- declarations - this is what the user types in the source program. -- These are used to construct 'TypeBase's in the type checker.-data TypeExp vn+data TypeExp f vn   = TEVar (QualName vn) SrcLoc-  | TETuple [TypeExp vn] SrcLoc-  | TERecord [(Name, TypeExp vn)] SrcLoc-  | TEArray (SizeExp vn) (TypeExp vn) SrcLoc-  | TEUnique (TypeExp vn) SrcLoc-  | TEApply (TypeExp vn) (TypeArgExp vn) SrcLoc-  | TEArrow (Maybe vn) (TypeExp vn) (TypeExp vn) SrcLoc-  | TESum [(Name, [TypeExp vn])] SrcLoc-  | TEDim [vn] (TypeExp vn) SrcLoc-  deriving (Show)+  | TETuple [TypeExp f vn] SrcLoc+  | TERecord [(Name, TypeExp f vn)] SrcLoc+  | TEArray (SizeExp f vn) (TypeExp f vn) SrcLoc+  | TEUnique (TypeExp f vn) SrcLoc+  | TEApply (TypeExp f vn) (TypeArgExp f vn) SrcLoc+  | TEArrow (Maybe vn) (TypeExp f vn) (TypeExp f vn) SrcLoc+  | TESum [(Name, [TypeExp f vn])] SrcLoc+  | TEDim [vn] (TypeExp f vn) SrcLoc -deriving instance Eq (TypeExp Name)+deriving instance Show (TypeExp Info VName) -deriving instance Eq (TypeExp VName)+deriving instance Show vn => Show (TypeExp NoInfo vn) -deriving instance Ord (TypeExp Name)+deriving instance Eq (TypeExp NoInfo VName) -deriving instance Ord (TypeExp VName)+deriving instance Eq (TypeExp Info VName) -instance Located (TypeExp vn) where+deriving instance Ord (TypeExp NoInfo VName)++deriving instance Ord (TypeExp Info VName)++instance Located (TypeExp f vn) where   locOf (TEArray _ _ loc) = locOf loc   locOf (TETuple _ loc) = locOf loc   locOf (TERecord _ loc) = locOf loc@@ -437,38 +450,33 @@   locOf (TEDim _ _ loc) = locOf loc  -- | A type argument expression passed to a type constructor.-data TypeArgExp vn-  = TypeArgExpDim (SizeExp vn) SrcLoc-  | TypeArgExpType (TypeExp vn)-  deriving (Show)+data TypeArgExp f vn+  = TypeArgExpSize (SizeExp f vn)+  | TypeArgExpType (TypeExp f vn) -deriving instance Eq (TypeArgExp Name)+deriving instance Show (TypeArgExp Info VName) -deriving instance Eq (TypeArgExp VName)+deriving instance Show vn => Show (TypeArgExp NoInfo vn) -deriving instance Ord (TypeArgExp Name)+deriving instance Eq (TypeArgExp NoInfo VName) -deriving instance Ord (TypeArgExp VName)+deriving instance Eq (TypeArgExp Info VName) -instance Located (TypeArgExp vn) where-  locOf (TypeArgExpDim _ loc) = locOf loc+deriving instance Ord (TypeArgExp NoInfo VName)++deriving instance Ord (TypeArgExp Info VName)++instance Located (TypeArgExp f vn) where+  locOf (TypeArgExpSize e) = locOf e   locOf (TypeArgExpType t) = locOf t --- | Information about which parts of a value/type are consumed.+-- | Information about which parts of a parameter are consumed.  This+-- can be considered kind of an effect on the function. data Diet-  = -- | Consumes these fields in the record.-    RecordDiet (M.Map Name Diet)-  | -- | Consume these parts of the constructors.-    SumDiet (M.Map Name [Diet])-  | -- | A function that consumes its argument(s) like this.-    -- The final 'Diet' should always be 'Observe', as there-    -- is no way for a function to consume its return value.-    FuncDiet Diet Diet-  | -- | Consumes this value.-    Consume-  | -- | Only observes value in this position, does-    -- not consume.+  = -- | Does not consume the parameter.     Observe+  | -- | Consumes the parameter.+    Consume   deriving (Eq, Ord, Show)  -- | An identifier consists of its name and the type of the value@@ -628,17 +636,21 @@ -- need, so we can pretend that an application expression was really -- bound to a name. data AppExpBase f vn-  = -- | The @Maybe VName@ is a possible existential size that is-    -- instantiated by this argument.  May have duplicates across the-    -- program, but they will all produce the same value (the-    -- expressions will be identical).+  = -- | Function application.  Parts of the compiler expects that the+    -- function expression is never itself an 'Apply'.  Use the+    -- 'mkApply' function to maintain this invariant, rather than+    -- constructing 'Apply' directly.+    --+    -- The @Maybe VNames@ are existential sizes generated by this+    -- argumnet.  May have duplicates across the program, but they+    -- will all produce the same value (the expressions will be+    -- identical).     Apply       (ExpBase f vn)-      (ExpBase f vn)-      (f (Diet, Maybe VName))+      (NE.NonEmpty (f (Diet, Maybe VName), ExpBase f vn))       SrcLoc   | -- | Size coercion: @e :> t@.-    Coerce (ExpBase f vn) (TypeExp vn) SrcLoc+    Coerce (ExpBase f vn) (TypeExp f vn) SrcLoc   | Range       (ExpBase f vn)       (Maybe (ExpBase f vn))@@ -654,7 +666,7 @@       vn       ( [TypeParamBase vn],         [PatBase f vn],-        Maybe (TypeExp vn),+        Maybe (TypeExp f vn),         f StructRetType,         ExpBase f vn       )@@ -702,7 +714,7 @@   locOf (BinOp _ _ _ _ loc) = locOf loc   locOf (If _ _ _ loc) = locOf loc   locOf (Coerce _ _ loc) = locOf loc-  locOf (Apply _ _ _ loc) = locOf loc+  locOf (Apply _ _ loc) = locOf loc   locOf (LetPat _ _ _ _ loc) = locOf loc   locOf (LetFun _ _ _ loc) = locOf loc   locOf (LetWith _ _ _ _ _ loc) = locOf loc@@ -766,7 +778,7 @@   | Lambda       [PatBase f vn]       (ExpBase f vn)-      (Maybe (TypeExp vn))+      (Maybe (TypeExp f vn))       (f (Aliasing, StructRetType))       SrcLoc   | -- | @+@; first two types are operands, third is result.@@ -792,7 +804,7 @@   | -- | Array indexing as a section: @(.[i,j])@.     IndexSection (SliceBase f vn) (f PatType) SrcLoc   | -- | Type ascription: @e : t@.-    Ascript (ExpBase f vn) (TypeExp vn) SrcLoc+    Ascript (ExpBase f vn) (TypeExp f vn) SrcLoc   | AppExp (AppExpBase f vn) (f AppRes)  deriving instance Show (ExpBase Info VName)@@ -908,7 +920,7 @@   | PatParens (PatBase f vn) SrcLoc   | Id vn (f PatType) SrcLoc   | Wildcard (f PatType) SrcLoc -- Nothing, i.e. underscore.-  | PatAscription (PatBase f vn) (TypeExp vn) SrcLoc+  | PatAscription (PatBase f vn) (TypeExp f vn) SrcLoc   | PatLit PatLit (f PatType) SrcLoc   | PatConstr Name (f PatType) [PatBase f vn] SrcLoc   | PatAttr (AttrInfo vn) (PatBase f vn) SrcLoc@@ -947,7 +959,7 @@ -- maybe also an ascribed type expression. data EntryType = EntryType   { entryType :: StructType,-    entryAscribed :: Maybe (TypeExp VName)+    entryAscribed :: Maybe (TypeExp Info VName)   }   deriving (Show) @@ -977,7 +989,7 @@     -- may refer to abstract types that are no longer in scope.     valBindEntryPoint :: Maybe (f EntryPoint),     valBindName :: vn,-    valBindRetDecl :: Maybe (TypeExp vn),+    valBindRetDecl :: Maybe (TypeExp f vn),     -- | If 'valBindParams' is null, then the 'retDims' are brought     -- into scope at this point.     valBindRetType :: f StructRetType,@@ -1001,7 +1013,7 @@   { typeAlias :: vn,     typeLiftedness :: Liftedness,     typeParams :: [TypeParamBase vn],-    typeExp :: TypeExp vn,+    typeExp :: TypeExp f vn,     typeElab :: f StructRetType,     typeDoc :: Maybe DocComment,     typeBindLocation :: SrcLoc@@ -1059,7 +1071,7 @@   = ValSpec       { specName :: vn,         specTypeParams :: [TypeParamBase vn],-        specTypeExp :: TypeExp vn,+        specTypeExp :: TypeExp f vn,         specType :: f StructType,         specDoc :: Maybe DocComment,         specLocation :: SrcLoc@@ -1086,7 +1098,7 @@   = SigVar (QualName vn) (f (M.Map VName VName)) SrcLoc   | SigParens (SigExpBase f vn) SrcLoc   | SigSpecs [SpecBase f vn] SrcLoc-  | SigWith (SigExpBase f vn) (TypeRefBase vn) SrcLoc+  | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc   | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc  deriving instance Show (SigExpBase Info VName)@@ -1094,13 +1106,13 @@ deriving instance Show (SigExpBase NoInfo Name)  -- | A type refinement.-data TypeRefBase vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeExp vn) SrcLoc+data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeExp f vn) SrcLoc -deriving instance Show (TypeRefBase VName)+deriving instance Show (TypeRefBase Info VName) -deriving instance Show (TypeRefBase Name)+deriving instance Show (TypeRefBase NoInfo Name) -instance Located (TypeRefBase vn) where+instance Located (TypeRefBase f vn) where   locOf (TypeRef _ _ _ loc) = locOf loc  instance Located (SigExpBase f vn) where@@ -1125,12 +1137,20 @@ instance Located (SigBindBase f vn) where   locOf = locOf . sigLoc +-- | Canonical reference to a Futhark code file.  Does not include the+-- @.fut@ extension.  This is most often a path relative to the+-- working directory of the compiler.  In a multi-file program, a file+-- is known by exactly one import name, even if it is referenced+-- relatively by different names by files in different subdirectories.+newtype ImportName = ImportName Posix.FilePath+  deriving (Eq, Ord, Show)+ -- | Module expression. data ModExpBase f vn   = ModVar (QualName vn) SrcLoc   | ModParens (ModExpBase f vn) SrcLoc   | -- | The contents of another file as a module.-    ModImport FilePath (f FilePath) SrcLoc+    ModImport FilePath (f ImportName) SrcLoc   | ModDecs [DecBase f vn] SrcLoc   | -- | Functor application.  The first mapping is from parameter     -- names to argument names, while the second maps names in the@@ -1201,7 +1221,7 @@   | ModDec (ModBindBase f vn)   | OpenDec (ModExpBase f vn) SrcLoc   | LocalDec (DecBase f vn) SrcLoc-  | ImportDec FilePath (f FilePath) SrcLoc+  | ImportDec FilePath (f ImportName) SrcLoc  deriving instance Show (DecBase Info VName) @@ -1226,6 +1246,28 @@ deriving instance Show (ProgBase Info VName)  deriving instance Show (ProgBase NoInfo Name)++-- | Construct an 'Apply' node, with type information.+mkApply :: ExpBase Info vn -> [(Diet, Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn+mkApply f args (AppRes t ext)+  | Just args' <- NE.nonEmpty $ map onArg args =+      case f of+        (AppExp (Apply f' f_args loc) (Info (AppRes _ f_ext))) ->+          AppExp+            (Apply f' (f_args <> args') (srcspan loc $ snd $ NE.last args'))+            (Info $ AppRes t $ f_ext <> ext)+        _ ->+          AppExp (Apply f args' (srcspan f $ snd $ NE.last args')) (Info (AppRes t ext))+  | otherwise = f+  where+    onArg (d, v, x) = (Info (d, v), x)++-- | Construct an 'Apply' node, without type information.+mkApplyUT :: ExpBase NoInfo vn -> ExpBase NoInfo vn -> ExpBase NoInfo vn+mkApplyUT (AppExp (Apply f args loc) _) x =+  AppExp (Apply f (args <> NE.singleton (NoInfo, x)) (srcspan loc x)) NoInfo+mkApplyUT f x =+  AppExp (Apply f (NE.singleton (NoInfo, x)) (srcspan f x)) NoInfo  --- Some prettyprinting definitions are here because we need them in --- the Attributes module.

@@ -76,8 +76,16 @@     If <$> mapOnExp tv c <*> mapOnExp tv texp <*> mapOnExp tv fexp <*> pure loc   astMap tv (Match e cases loc) =     Match <$> mapOnExp tv e <*> astMap tv cases <*> pure loc-  astMap tv (Apply f arg d loc) =-    Apply <$> mapOnExp tv f <*> mapOnExp tv arg <*> pure d <*> pure loc+  astMap tv (Apply f args loc) = do+    f' <- mapOnExp tv f+    args' <- traverse (traverse $ mapOnExp tv) args+    -- Safe to disregard return type because existentials cannot be+    -- instantiated here, as the return is necessarily a function.+    pure $ case f' of+      AppExp (Apply f_inner args_inner _) _ ->+        Apply f_inner (args_inner <> args') loc+      _ ->+        Apply f' args' loc   astMap tv (LetPat sizes pat e body loc) =     LetPat <$> astMap tv sizes <*> astMap tv pat <*> mapOnExp tv e <*> mapOnExp tv body <*> pure loc   astMap tv (LetFun name (fparams, params, ret, t, e) body loc) =@@ -231,14 +239,17 @@   astMap tv (ForIn pat e) = ForIn <$> astMap tv pat <*> mapOnExp tv e   astMap tv (While e) = While <$> mapOnExp tv e -instance ASTMappable (TypeExp VName) where-  astMap tv (TEVar qn loc) = TEVar <$> astMap tv qn <*> pure loc-  astMap tv (TETuple ts loc) = TETuple <$> traverse (astMap tv) ts <*> pure loc+instance ASTMappable (TypeExp Info VName) where+  astMap tv (TEVar qn loc) =+    TEVar <$> astMap tv qn <*> pure loc+  astMap tv (TETuple ts loc) =+    TETuple <$> traverse (astMap tv) ts <*> pure loc   astMap tv (TERecord ts loc) =     TERecord <$> traverse (traverse $ astMap tv) ts <*> pure loc   astMap tv (TEArray te dim loc) =     TEArray <$> astMap tv te <*> astMap tv dim <*> pure loc-  astMap tv (TEUnique t loc) = TEUnique <$> astMap tv t <*> pure loc+  astMap tv (TEUnique t loc) =+    TEUnique <$> astMap tv t <*> pure loc   astMap tv (TEApply t1 t2 loc) =     TEApply <$> astMap tv t1 <*> astMap tv t2 <*> pure loc   astMap tv (TEArrow v t1 t2 loc) =@@ -248,17 +259,13 @@   astMap tv (TEDim dims t loc) =     TEDim dims <$> astMap tv t <*> pure loc -instance ASTMappable (TypeArgExp VName) where-  astMap tv (TypeArgExpDim dim loc) =-    TypeArgExpDim <$> astMap tv dim <*> pure loc-  astMap tv (TypeArgExpType te) =-    TypeArgExpType <$> astMap tv te+instance ASTMappable (TypeArgExp Info VName) where+  astMap tv (TypeArgExpSize dim) = TypeArgExpSize <$> astMap tv dim+  astMap tv (TypeArgExpType te) = TypeArgExpType <$> astMap tv te -instance ASTMappable (SizeExp VName) where-  astMap tv (SizeExpNamed vn loc) =-    SizeExpNamed <$> astMap tv vn <*> pure loc-  astMap _ (SizeExpConst k loc) = pure $ SizeExpConst k loc-  astMap _ SizeExpAny = pure SizeExpAny+instance ASTMappable (SizeExp Info VName) where+  astMap tv (SizeExp e loc) = SizeExp <$> mapOnExp tv e <*> pure loc+  astMap _ (SizeExpAny loc) = pure $ SizeExpAny loc  instance ASTMappable Size where   astMap tv (NamedSize vn) = NamedSize <$> astMap tv vn@@ -301,10 +308,11 @@ traverseScalarType f g h (Record fs) = Record <$> traverse (traverseType f g h) fs traverseScalarType f g h (TypeVar als u t args) =   TypeVar <$> h als <*> pure u <*> f t <*> traverse (traverseTypeArg f g) args-traverseScalarType f g h (Arrow als v t1 (RetType dims t2)) =+traverseScalarType f g h (Arrow als v u t1 (RetType dims t2)) =   Arrow     <$> h als     <*> pure v+    <*> pure u     <*> traverseType f g pure t1     <*> (RetType dims <$> traverseType f g h t2) traverseScalarType f g h (Sum cs) =@@ -414,7 +422,7 @@ barePat (PatParens p loc) = PatParens (barePat p) loc barePat (Id v _ loc) = Id v NoInfo loc barePat (Wildcard _ loc) = Wildcard NoInfo loc-barePat (PatAscription pat t loc) = PatAscription (barePat pat) t loc+barePat (PatAscription pat t loc) = PatAscription (barePat pat) (bareTypeExp t) loc barePat (PatLit v _ loc) = PatLit v NoInfo loc barePat (PatConstr c _ ps loc) = PatConstr c NoInfo (map barePat ps) loc barePat (PatAttr attr p loc) = PatAttr attr (barePat p) loc@@ -433,6 +441,9 @@ bareCase :: CaseBase Info VName -> CaseBase NoInfo VName bareCase (CasePat pat e loc) = CasePat (barePat pat) (bareExp e) loc +bareTypeExp :: TypeExp Info VName -> TypeExp NoInfo VName+bareTypeExp = undefined+ -- | Remove all annotations from an expression, but retain the -- name/scope information. bareExp :: ExpBase Info VName -> ExpBase NoInfo VName@@ -447,7 +458,7 @@ bareExp (StringLit vs loc) = StringLit vs loc bareExp (RecordLit fields loc) = RecordLit (map bareField fields) loc bareExp (ArrayLit els _ loc) = ArrayLit (map bareExp els) NoInfo loc-bareExp (Ascript e te loc) = Ascript (bareExp e) te loc+bareExp (Ascript e te loc) = Ascript (bareExp e) (bareTypeExp te) loc bareExp (Negate x loc) = Negate (bareExp x) loc bareExp (Not x loc) = Not (bareExp x) loc bareExp (Update src slice v loc) =@@ -458,7 +469,7 @@   Project field (bareExp e) NoInfo loc bareExp (Assert e1 e2 _ loc) = Assert (bareExp e1) (bareExp e2) NoInfo loc bareExp (Lambda params body ret _ loc) =-  Lambda (map barePat params) (bareExp body) ret NoInfo loc+  Lambda (map barePat params) (bareExp body) (fmap bareTypeExp ret) NoInfo loc bareExp (OpSection name _ loc) = OpSection name NoInfo loc bareExp (OpSectionLeft name _ arg _ _ loc) =   OpSectionLeft name NoInfo (bareExp arg) (NoInfo, NoInfo) (NoInfo, NoInfo) loc@@ -496,16 +507,20 @@           BinOp fname NoInfo (bareExp x, NoInfo) (bareExp y, NoInfo) loc         If c texp fexp loc ->           If (bareExp c) (bareExp texp) (bareExp fexp) loc-        Apply f arg _ loc ->-          Apply (bareExp f) (bareExp arg) NoInfo loc+        Apply f args loc ->+          Apply (bareExp f) (fmap ((NoInfo,) . bareExp . snd) args) loc         LetPat sizes pat e body loc ->           LetPat sizes (barePat pat) (bareExp e) (bareExp body) loc         LetFun name (fparams, params, ret, _, e) body loc ->-          LetFun name (fparams, map barePat params, ret, NoInfo, bareExp e) (bareExp body) loc+          LetFun+            name+            (fparams, map barePat params, fmap bareTypeExp ret, NoInfo, bareExp e)+            (bareExp body)+            loc         Range start next end loc ->           Range (bareExp start) (fmap bareExp next) (fmap bareExp end) loc         Coerce e te loc ->-          Coerce (bareExp e) te loc+          Coerce (bareExp e) (bareTypeExp te) loc         Index arr slice loc ->           Index (bareExp arr) (map bareDimIndex slice) loc bareExp (Attr attr e loc) =

@@ -203,7 +203,7 @@  checkTypeDecl ::   UncheckedTypeExp ->-  TypeM ([VName], TypeExp VName, StructType, Liftedness)+  TypeM ([VName], TypeExp Info VName, StructType, Liftedness) checkTypeDecl te = do   (te', svars, RetType dims st, l) <- checkTypeExp te   pure (svars ++ dims, te', st, l)@@ -593,7 +593,7 @@           TypeBind name' l tps' te' (Info elab_t) doc loc         ) -entryPoint :: [Pat] -> Maybe (TypeExp VName) -> StructRetType -> EntryPoint+entryPoint :: [Pat] -> Maybe (TypeExp Info VName) -> StructRetType -> EntryPoint entryPoint params orig_ret_te (RetType ret orig_ret) =   EntryPoint (map patternEntry params ++ more_params) rettype'   where@@ -617,10 +617,10 @@      pname (Named v) = baseName v     pname Unnamed = "_"-    onRetType (Just (TEArrow p t1_te t2_te _)) (Scalar (Arrow _ _ t1 (RetType _ t2))) =+    onRetType (Just (TEArrow p t1_te t2_te _)) (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =       let (xs, y) = onRetType (Just t2_te) t2        in (EntryParam (maybe "_" baseName p) (EntryType t1 (Just t1_te)) : xs, y)-    onRetType _ (Scalar (Arrow _ p t1 (RetType _ t2))) =+    onRetType _ (Scalar (Arrow _ p _ t1 (RetType _ t2))) =       let (xs, y) = onRetType Nothing t2        in (EntryParam (pname p) (EntryType t1 Nothing) : xs, y)     onRetType te t =@@ -630,7 +630,7 @@   SrcLoc ->   [TypeParam] ->   [Pat] ->-  Maybe (TypeExp VName) ->+  Maybe (TypeExp Info VName) ->   StructRetType ->   TypeM () checkEntryPoint loc tparams params maybe_tdecl rettype@@ -639,8 +639,7 @@         withIndexLink           "polymorphic-entry"           "Entry point functions may not be polymorphic."-  | not (all patternOrderZero params)-      || not (all orderZero rettype_params)+  | not (all orderZero param_ts)       || not (orderZero rettype') =       typeError loc mempty $         withIndexLink@@ -649,7 +648,7 @@   | sizes_only_in_ret <-       S.fromList (map typeParamName tparams)         `S.intersection` freeInType rettype'-        `S.difference` foldMap freeInType (map patternStructType params ++ rettype_params),+        `S.difference` foldMap freeInType param_ts,     not $ S.null sizes_only_in_ret =       typeError loc mempty $         withIndexLink@@ -670,6 +669,7 @@   where     (RetType _ rettype_t) = rettype     (rettype_params, rettype') = unfoldFunType rettype_t+    param_ts = map patternStructType params ++ map snd rettype_params  checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind) checkValBind (ValBind entry fname maybe_tdecl NoInfo tparams params body doc attrs loc) = do@@ -704,10 +704,10 @@ nastyType t@Array {} = nastyType $ stripArray 1 t nastyType _ = True -nastyReturnType :: Monoid als => Maybe (TypeExp VName) -> TypeBase dim als -> Bool-nastyReturnType Nothing (Scalar (Arrow _ _ t1 (RetType _ t2))) =+nastyReturnType :: Monoid als => Maybe (TypeExp Info VName) -> TypeBase dim als -> Bool+nastyReturnType Nothing (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =   nastyType t1 || nastyReturnType Nothing t2-nastyReturnType (Just (TEArrow _ te1 te2 _)) (Scalar (Arrow _ _ t1 (RetType _ t2))) =+nastyReturnType (Just (TEArrow _ te1 te2 _)) (Scalar (Arrow _ _ _ t1 (RetType _ t2))) =   (not (niceTypeExp te1) && nastyType t1)     || nastyReturnType (Just te2) t2 nastyReturnType (Just te) _@@ -729,9 +729,9 @@     ascripted (PatParens p' _) = ascripted p'     ascripted _ = False -niceTypeExp :: TypeExp VName -> Bool+niceTypeExp :: TypeExp Info VName -> Bool niceTypeExp (TEVar (QualName [] _) _) = True-niceTypeExp (TEApply te TypeArgExpDim {} _) = niceTypeExp te+niceTypeExp (TEApply te TypeArgExpSize {} _) = niceTypeExp te niceTypeExp (TEArray _ te _) = niceTypeExp te niceTypeExp (TEUnique te _) = niceTypeExp te niceTypeExp _ = False

@@ -149,8 +149,8 @@           Scalar $ Sum $ (fmap . fmap) substituteInType ts         substituteInType (Array () u shape t) =           arrayOf u (substituteInShape shape) (substituteInType $ Scalar t)-        substituteInType (Scalar (Arrow als v t1 (RetType dims t2))) =-          Scalar $ Arrow als v (substituteInType t1) $ RetType dims $ substituteInType t2+        substituteInType (Scalar (Arrow als v d1 t1 (RetType dims t2))) =+          Scalar $ Arrow als v d1 (substituteInType t1) $ RetType dims $ substituteInType t2          substituteInShape (Shape ds) =           Shape $ map substituteInDim ds

@@ -151,9 +151,9 @@       <+> dquotes (pretty name)         <> ": variables prefixed with underscore may not be accessed." --- | A mapping from import strings to 'Env's.  This is used to resolve--- @import@ declarations.-type ImportTable = M.Map String Env+-- | A mapping from import import names to 'Env's.  This is used to+-- resolve @import@ declarations.+type ImportTable = M.Map ImportName Env  data Context = Context   { contextEnv :: Env,@@ -239,18 +239,18 @@     explode = unknownVariable Signature qn loc  -- | Look up an import.-lookupImport :: SrcLoc -> FilePath -> TypeM (FilePath, Env)+lookupImport :: SrcLoc -> FilePath -> TypeM (ImportName, Env) lookupImport loc file = do   imports <- asks contextImportTable   my_path <- asks contextImportName-  let canonical_import = includeToString $ mkImportFrom my_path file loc+  let canonical_import = mkImportFrom my_path file   case M.lookup canonical_import imports of     Nothing ->       typeError loc mempty $         "Unknown import"-          <+> dquotes (pretty canonical_import)+          <+> dquotes (pretty (includeToText canonical_import))           </> "Known:"-          <+> commasep (map pretty (M.keys imports))+          <+> commasep (map (pretty . includeToText) (M.keys imports))     Just scope -> pure (canonical_import, scope)  -- | Evaluate a 'TypeM' computation within an extended (/not/@@ -432,8 +432,8 @@       Record $ M.map (onType except) m     onScalar except (Sum m) =       Sum $ M.map (map $ onType except) m-    onScalar except (Arrow as p t1 (RetType dims t2)) =-      Arrow as p (onType except' t1) $ RetType dims (onType except' t2)+    onScalar except (Arrow as p d t1 (RetType dims t2)) =+      Arrow as p d (onType except' t1) $ RetType dims (onType except' t2)       where         except' = case p of           Named p' -> S.insert p' except

@@ -16,11 +16,12 @@ import Control.Monad.Reader import Control.Monad.State import Data.Either-import Data.List (find, foldl', partition)+import Data.List (find, foldl', genericLength, partition) import Data.List.NonEmpty qualified as NE import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S+import Futhark.Util (mapAccumLM) import Futhark.Util.Pretty hiding (space) import Language.Futhark import Language.Futhark.Primitive (intByteSize)@@ -146,7 +147,7 @@   SrcLoc ->   UncheckedTypeExp ->   UncheckedExp ->-  TermTypeM (TypeExp VName, Exp)+  TermTypeM (TypeExp Info VName, Exp) checkAscript loc te e = do   (te', decl_t, _) <- checkTypeExpNonrigid te   e' <- checkExp e@@ -161,7 +162,7 @@   SrcLoc ->   UncheckedTypeExp ->   UncheckedExp ->-  TermTypeM (TypeExp VName, StructType, Exp, [VName])+  TermTypeM (TypeExp Info VName, StructType, Exp, [VName]) checkCoerce loc te e = do   (te', te_t, ext) <- checkTypeExpRigid te RigidCoerce   e' <- checkExp e@@ -204,37 +205,12 @@     unAlias (AliasBound v) | v `M.member` unscoped = AliasFree v     unAlias a = a --- 'checkApplyExp' is like 'checkExp', but tries to find the "root--- function", for better error messages.-checkApplyExp :: UncheckedExp -> TermTypeM (Exp, ApplyOp)-checkApplyExp (AppExp (Apply e1 e2 _ loc) _) = do-  arg <- checkArg e2-  (e1', (fname, i)) <- checkApplyExp e1-  t <- expType e1'-  (t1, rt, argext, exts) <- checkApply loc (fname, i) t arg-  pure-    ( AppExp-        (Apply e1' (argExp arg) (Info (diet t1, argext)) loc)-        (Info $ AppRes rt exts),-      (fname, i + 1)-    )-checkApplyExp e = do-  e' <- checkExp e-  pure-    ( e',-      ( case e' of-          Var qn _ _ -> Just qn-          _ -> Nothing,-        0-      )-    )- checkExp :: UncheckedExp -> TermTypeM Exp checkExp (Literal val loc) =   pure $ Literal val loc checkExp (Hole _ loc) = do   t <- newTypeVar loc "t"-  pure $ Hole (Info $ t `setUniqueness` Unique) loc+  pure $ Hole (Info t) loc checkExp (StringLit vs loc) =   pure $ StringLit vs loc checkExp (IntLit val NoInfo loc) = do@@ -281,14 +257,14 @@   case all_es of     [] -> do       et <- newTypeVar loc "t"-      t <- arrayOfM loc et (Shape [ConstSize 0]) Unique+      t <- arrayOfM loc et (Shape [ConstSize 0]) Nonunique       pure $ ArrayLit [] (Info t) loc     e : es -> do       e' <- checkExp e       et <- expType e'       es' <- mapM (unifies "type of first array element" (toStruct et) <=< checkExp) es       et' <- normTypeFully et-      t <- arrayOfM loc et' (Shape [ConstSize $ length all_es]) Unique+      t <- arrayOfM loc et' (Shape [ConstSize $ genericLength all_es]) Nonunique       pure $ ArrayLit (e' : es') (Info t) loc checkExp (AppExp (Range start maybe_step end loc) _) = do   start' <- require "use in range expression" anySignedType =<< checkExp start@@ -328,7 +304,7 @@         d <- newDimVar loc (Rigid RigidRange) "range_dim"         pure (NamedSize $ qualName d, Just d) -  t <- arrayOfM loc start_t (Shape [dim]) Unique+  t <- arrayOfM loc start_t (Shape [dim]) Nonunique   let res = AppRes (t `setAliases` mempty) (maybeToList retext)    pure $ AppExp (Range start' maybe_step' end' loc) (Info res)@@ -347,8 +323,8 @@    -- Note that the application to the first operand cannot fix any   -- existential sizes, because it must by necessity be a function.-  (p1_t, rt, p1_ext, _) <- checkApply loc (Just op', 0) ftype e1_arg-  (p2_t, rt', p2_ext, retext) <- checkApply loc (Just op', 1) rt e2_arg+  (_, p1_t, rt, p1_ext, _) <- checkApply loc (Just op', 0) ftype e1_arg+  (_, p2_t, rt', p2_ext, retext) <- checkApply loc (Just op', 1) rt e2_arg    pure $     AppExp@@ -437,7 +413,23 @@ checkExp (Not arg loc) = do   arg' <- require "logical negation" (Bool : anyIntType) =<< checkExp arg   pure $ Not arg' loc-checkExp e@(AppExp Apply {} _) = fst <$> checkApplyExp e+checkExp (AppExp (Apply fe args loc) NoInfo) = do+  fe' <- checkExp fe+  args' <- mapM (checkArg . snd) args+  t <- expType fe'+  let fname =+        case fe' of+          Var v _ _ -> Just v+          _ -> Nothing+  ((_, exts, rt), args'') <- mapAccumLM (onArg fname) (0, [], t) args'+  pure $ AppExp (Apply fe' args'' loc) $ Info $ AppRes rt exts+  where+    onArg fname (i, all_exts, t) arg' = do+      (d1, _, rt, argext, exts) <- checkApply loc (fname, i) t arg'+      pure+        ( (i + 1, all_exts <> exts, rt),+          (Info (d1, argext), argExp arg')+        ) checkExp (AppExp (LetPat sizes pat e body loc) _) =   sequentially (checkExp e) $ \e' e_occs -> do     -- Not technically an ascription, but we want the pattern to have@@ -471,8 +463,7 @@       bindSpaced [(Term, name)] $ do         name' <- checkName Term name loc -        let arrow (xp, xt) yt = RetType [] $ Scalar $ Arrow () xp xt yt-            RetType _ ftype = foldr (arrow . patternParam) rettype params'+        let ftype = funType params' rettype             entry = BoundV Local tparams' $ ftype `setAliases` closure'             bindF scope =               scope@@ -627,12 +618,11 @@     -- are parameters, or are used in parameters.     inferReturnSizes params' ret = do       cur_lvl <- curLevel-      let named (Named x, _) = Just x-          named (Unnamed, _) = Nothing+      let named (Named x, _, _) = Just x+          named (Unnamed, _, _) = Nothing           param_names = mapMaybe (named . patternParam) params'           pos_sizes =-            sizeNamesPos . foldFunType (map patternStructType params') $-              RetType [] ret+            sizeNamesPos $ foldFunTypeFromParams params' $ RetType [] ret           hide k (lvl, _) =             lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes @@ -650,9 +640,9 @@ checkExp (OpSectionLeft op _ e _ _ loc) = do   (op', ftype) <- lookupVar loc op   e_arg <- checkArg e-  (t1, rt, argext, retext) <- checkApply loc (Just op', 0) ftype e_arg+  (_, t1, rt, argext, retext) <- checkApply loc (Just op', 0) ftype e_arg   case (ftype, rt) of-    (Scalar (Arrow _ m1 _ _), Scalar (Arrow _ m2 t2 rettype)) ->+    (Scalar (Arrow _ m1 _ _ _), Scalar (Arrow _ m2 _ t2 rettype)) ->       pure $         OpSectionLeft           op'@@ -668,12 +658,12 @@   (op', ftype) <- lookupVar loc op   e_arg <- checkArg e   case ftype of-    Scalar (Arrow as1 m1 t1 (RetType [] (Scalar (Arrow as2 m2 t2 (RetType dims2 ret))))) -> do-      (t2', ret', argext, _) <-+    Scalar (Arrow as1 m1 d1 t1 (RetType [] (Scalar (Arrow as2 m2 d2 t2 (RetType dims2 ret))))) -> do+      (_, t2', ret', argext, _) <-         checkApply           loc           (Just op', 1)-          (Scalar $ Arrow as2 m2 t2 $ RetType [] $ Scalar $ Arrow as1 m1 t1 $ RetType [] ret)+          (Scalar $ Arrow as2 m2 d2 t2 $ RetType [] $ Scalar $ Arrow as1 m1 d1 t1 $ RetType [] ret)           e_arg       pure $         OpSectionRight@@ -690,13 +680,13 @@   a <- newTypeVar loc "a"   let usage = mkUsage loc "projection at"   b <- foldM (flip $ mustHaveField usage) a fields-  let ft = Scalar $ Arrow mempty Unnamed (toStruct a) $ RetType [] b+  let ft = Scalar $ Arrow mempty Unnamed Observe (toStruct a) $ RetType [] b   pure $ ProjectSection fields (Info ft) loc checkExp (IndexSection slice NoInfo loc) = do   slice' <- checkSlice slice   (t, _) <- newArrayType loc "e" $ sliceDims slice'   (t', retext) <- sliceShape Nothing slice' t-  let ft = Scalar $ Arrow mempty Unnamed t $ RetType retext $ fromStruct t'+  let ft = Scalar $ Arrow mempty Unnamed Observe t $ RetType retext $ fromStruct t'   pure $ IndexSection slice' (Info ft) loc checkExp (AppExp (DoLoop _ mergepat mergeexp form loopbody loc) _) = do   ((sparams, mergepat', mergeexp', form', loopbody'), appres) <-@@ -841,7 +831,7 @@     f TypeArgDim {} = mempty boundInsideType (Scalar (Record fs)) = foldMap boundInsideType fs boundInsideType (Scalar (Sum cs)) = foldMap (foldMap boundInsideType) cs-boundInsideType (Scalar (Arrow _ pn t1 (RetType dims t2))) =+boundInsideType (Scalar (Arrow _ pn _ t1 (RetType dims t2))) =   pn' <> boundInsideType t1 <> S.fromList dims <> boundInsideType t2   where     pn' = case pn of@@ -863,11 +853,11 @@   ApplyOp ->   PatType ->   Arg ->-  TermTypeM (StructType, PatType, Maybe VName, [VName])+  TermTypeM (Diet, StructType, PatType, Maybe VName, [VName]) checkApply   loc   (fname, _)-  (Scalar (Arrow as pname tp1 tp2))+  (Scalar (Arrow as pname d1 tp1 tp2))   (argexp, argtype, dflow, argloc) =     onFailure (CheckingApply fname argexp tp1 (toStruct argtype)) $ do       expect (mkUsage argloc "use as function argument") (toStruct tp1) (toStruct argtype)@@ -896,14 +886,14 @@            in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1         _ -> pure () -      arg_consumed <- consumedByArg argloc argtype' (diet tp1')+      arg_consumed <- consumedByArg (locOf argloc) argtype' d1       checkIfConsumable loc $ mconcat arg_consumed       occur $ dflow `seqOccurrences` map (`consumption` argloc) arg_consumed        -- Unification ignores uniqueness in higher-order arguments, so       -- we check for that here.       unless (toStructural argtype' `subtypeOf` setUniqueness (toStructural tp1') Nonunique) $-        typeError loc mempty "Consumption/aliasing does not match."+        typeError loc mempty "Difference in whether argument is consumed."        (argext, parsubst) <-         case pname of@@ -923,18 +913,16 @@       v <- newID "internal_app_result"       modify $ \s -> s {stateNames = M.insert v (NameAppRes fname loc) $ stateNames s}       let appres = S.singleton $ AliasFree v-      let tp2'' = applySubst parsubst $ returnType appres tp2' (diet tp1') argtype'+      let tp2'' = applySubst parsubst $ returnType appres tp2' d1 argtype' -      pure (tp1', tp2'', argext, ext)+      pure (d1, tp1', tp2'', argext, ext) checkApply loc fname tfun@(Scalar TypeVar {}) arg = do   tv <- newTypeVar loc "b"   -- Change the uniqueness of the argument type because we never want   -- to infer that a function is consuming.   let argt_nonunique = toStruct (argType arg) `setUniqueness` Nonunique   unify (mkUsage loc "use as function") (toStruct tfun) $-    Scalar $-      Arrow mempty Unnamed argt_nonunique $-        RetType [] tv+    Scalar (Arrow mempty Unnamed Observe argt_nonunique $ RetType [] tv)   tfun' <- normPatType tfun   checkApply loc fname tfun' arg checkApply loc (fname, prev_applied) ftype (argexp, _, _, _) = do@@ -962,27 +950,21 @@       | prev_applied == 1 = "argument"       | otherwise = "arguments" -consumedByArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Aliasing]-consumedByArg loc (Scalar (Record ets)) (RecordDiet ds) =-  mconcat . M.elems <$> traverse (uncurry $ consumedByArg loc) (M.intersectionWith (,) ets ds)-consumedByArg loc (Scalar (Sum ets)) (SumDiet ds) =-  mconcat <$> traverse (uncurry $ consumedByArg loc) (concat $ M.elems $ M.intersectionWith zip ets ds)-consumedByArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _)-  | not $ contravariantArg t1 d =-      typeError loc mempty . withIndexLink "consuming-argument" $-        "Non-consuming higher-order parameter passed consuming argument."+aliasParts :: PatType -> [Aliasing]+aliasParts (Scalar (Record ts)) = foldMap aliasParts $ M.elems ts+aliasParts t = [aliases t]++consumedByArg :: Loc -> PatType -> Diet -> TermTypeM [Aliasing]+consumedByArg loc at Consume = do+  let parts = aliasParts at+  foldM_ check mempty parts+  pure parts   where-    contravariantArg (Array _ Unique _ _) Observe =-      False-    contravariantArg (Scalar (TypeVar _ Unique _ _)) Observe =-      False-    contravariantArg (Scalar (Record ets)) (RecordDiet ds) =-      and (M.intersectionWith contravariantArg ets ds)-    contravariantArg (Scalar (Arrow _ _ tp (RetType _ tr))) (FuncDiet dp dr) =-      contravariantArg tp dp && contravariantArg tr dr-    contravariantArg _ _ =-      True-consumedByArg _ at Consume = pure [aliases at]+    check seen als+      | any (`S.member` seen) als =+          typeError loc mempty . withIndexLink "self-aliasing-arg" $+            "Argument passed for consuming parameter is self-aliased."+      | otherwise = pure $ als <> seen consumedByArg _ _ _ = pure []  -- | Type-check a single expression in isolation.  This expression may@@ -1051,9 +1033,17 @@       onExp known e@(AppExp (LetPat _ _ bindee_e body_e _) (Info res)) = do         sequencePoint known bindee_e body_e $ appResExt res         pure e-      onExp known e@(AppExp (Apply f arg (Info (_, p)) _) (Info res)) = do-        sequencePoint known arg f $ maybeToList p ++ appResExt res+      onExp known e@(AppExp (Apply f args _) (Info res)) = do+        seqArgs known $ reverse $ NE.toList args         pure e+        where+          seqArgs known' [] = do+            void $ onExp known' f+            modify (S.fromList (appResExt res) <>)+          seqArgs known' ((Info (_, p), x) : xs) = do+            new_known <- lift $ execStateT (onExp known' x) mempty+            void $ seqArgs (new_known <> known') xs+            modify ((new_known <> S.fromList (maybeToList p)) <>)       onExp         known         e@(AppExp (BinOp (f, floc) ft (x, Info (_, xp)) (y, Info (_, yp)) _) (Info res)) = do@@ -1188,7 +1178,7 @@     ( VName,       [TypeParam],       [Pat],-      Maybe (TypeExp VName),+      Maybe (TypeExp Info VName),       StructRetType,       Exp     )@@ -1283,8 +1273,8 @@ hiddenParamNames params = hidden   where     param_all_names = mconcat $ map patNames params-    named (Named x, _) = Just x-    named (Unnamed, _) = Nothing+    named (Named x, _, _) = Just x+    named (Unnamed, _, _) = Nothing     param_names =       S.fromList $ mapMaybe (named . patternParam) params     hidden = param_all_names `S.difference` param_names@@ -1352,7 +1342,7 @@   TermTypeM     ( [TypeParam],       [Pat],-      Maybe (TypeExp VName),+      Maybe (TypeExp Info VName),       StructRetType,       Exp     )@@ -1382,7 +1372,7 @@         pure (Just retdecl', ret')       Nothing         | null params ->-            pure (Nothing, toStruct $ body_t `setUniqueness` Nonunique)+            pure (Nothing, toStruct body_t)         | otherwise -> do             body_t' <- inferredReturnType loc params'' body_t             pure (Nothing, body_t')@@ -1399,7 +1389,7 @@ -- | Extract all the shape names that occur in positive position -- (roughly, left side of an arrow) in a given type. sizeNamesPos :: TypeBase Size als -> S.Set VName-sizeNamesPos (Scalar (Arrow _ _ t1 (RetType _ t2))) = onParam t1 <> sizeNamesPos t2+sizeNamesPos (Scalar (Arrow _ _ _ t1 (RetType _ t2))) = onParam t1 <> sizeNamesPos t2   where     onParam :: TypeBase Size als -> S.Set VName     onParam (Scalar Arrow {}) = mempty@@ -1511,7 +1501,7 @@       where         forbidden' =           case patternParam p of-            (Named v, _) -> forbidden `S.difference` S.singleton v+            (Named v, _, _) -> forbidden `S.difference` S.singleton v             _ -> forbidden     verifyParams _ [] = pure () @@ -1536,8 +1526,8 @@     deeper (Scalar (Prim t)) = Scalar $ Prim t     deeper (Scalar (Record fs)) = Scalar $ Record $ M.map deeper fs     deeper (Scalar (Sum cs)) = Scalar $ Sum $ M.map (map deeper) cs-    deeper (Scalar (Arrow als p t1 (RetType t2_ext t2))) =-      Scalar $ Arrow als p t1 $ injectExt (ext_there <> t2_ext) t2+    deeper (Scalar (Arrow als p d1 t1 (RetType t2_ext t2))) =+      Scalar $ Arrow als p d1 t1 $ injectExt (ext_there <> t2_ext) t2     deeper (Scalar (TypeVar as u tn targs)) =       Scalar $ TypeVar as u tn $ map deeperArg targs     deeper t@Array {} = t@@ -1571,7 +1561,8 @@       injectExt (retext ++ mapMaybe mkExt (S.toList $ freeInType ret)) ret     )   where-    t = foldFunType paramts $ RetType [] ret+    -- Diet does not matter here.+    t = foldFunType (zip (repeat Observe) paramts) $ RetType [] ret     to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs     visible = typeVars t <> freeInType t 

@@ -553,7 +553,7 @@   TermTypeM ([VName], PatType) instantiateTypeScheme qn loc tparams t = do   let tnames = map typeParamName tparams-  (tparam_names, tparam_substs) <- unzip <$> mapM (instantiateTypeParam qn loc) tparams+  (tparam_names, tparam_substs) <- mapAndUnzipM (instantiateTypeParam qn loc) tparams   let substs = M.fromList $ zip tnames tparam_substs       t' = applySubst (`M.lookup` substs) t   pure (tparam_names, t')@@ -677,9 +677,9 @@         argtype <- newTypeVar loc "t"         equalityType usage argtype         pure $-          Scalar . Arrow mempty Unnamed argtype . RetType [] $+          Scalar . Arrow mempty Unnamed Observe argtype . RetType [] $             Scalar $-              Arrow mempty Unnamed argtype $+              Arrow mempty Unnamed Observe argtype $                 RetType [] $                   Scalar $                     Prim Bool@@ -687,7 +687,7 @@         argtype <- newTypeVar loc "t"         mustBeOneOf ts usage argtype         let (pts', rt') = instOverloaded argtype pts rt-            arrow xt yt = Scalar $ Arrow mempty Unnamed xt $ RetType [] yt+            arrow xt yt = Scalar $ Arrow mempty Unnamed Observe xt $ RetType [] yt         pure $ fromStruct $ foldr arrow rt' pts'      observe $ Ident name (Info t) loc@@ -812,7 +812,9 @@   mustBeOneOf ts (mkUsage (srclocOf e) why) . toStruct =<< expType e   pure e -termCheckTypeExp :: TypeExp Name -> TermTypeM (TypeExp VName, [VName], StructRetType)+termCheckTypeExp ::+  TypeExp NoInfo Name ->+  TermTypeM (TypeExp Info VName, [VName], StructRetType) termCheckTypeExp te = do   (te', svars, rettype, _l) <- checkTypeExp te @@ -830,7 +832,7 @@     observeDim v =       observe $ Ident v (Info $ Scalar $ Prim $ Signed Int64) mempty -checkTypeExpNonrigid :: TypeExp Name -> TermTypeM (TypeExp VName, StructType, [VName])+checkTypeExpNonrigid :: TypeExp NoInfo Name -> TermTypeM (TypeExp Info VName, StructType, [VName]) checkTypeExpNonrigid te = do   (te', svars, RetType dims st) <- termCheckTypeExp te   forM_ (svars ++ dims) $ \v ->@@ -838,9 +840,9 @@   pure (te', st, svars ++ dims)  checkTypeExpRigid ::-  TypeExp Name ->+  TypeExp NoInfo Name ->   RigidSource ->-  TermTypeM (TypeExp VName, StructType, [VName])+  TermTypeM (TypeExp Info VName, StructType, [VName]) checkTypeExpRigid te rsrc = do   (te', svars, RetType dims st) <- termCheckTypeExp te   forM_ (svars ++ dims) $ \v ->@@ -1002,7 +1004,7 @@     initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics      prim = Scalar . Prim-    arrow x y = Scalar $ Arrow mempty Unnamed x y+    arrow x y = Scalar $ Arrow mempty Unnamed Observe x y      addIntrinsicF (name, IntrinsicMonoFun pts t) =       Just (name, BoundV Global [] $ arrow pts' $ RetType [] $ prim t)@@ -1015,15 +1017,8 @@     addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) =       Just         ( name,-          BoundV Global tvs $-            fromStruct $-              Scalar $-                Arrow mempty Unnamed pts' rt+          BoundV Global tvs $ fromStruct $ foldFunType pts rt         )-      where-        pts' = case pts of-          [pt] -> pt-          _ -> Scalar $ tupleRecord pts     addIntrinsicF (name, IntrinsicEquality) =       Just (name, EqualityF)     addIntrinsicF _ = Nothing

@@ -72,7 +72,7 @@           M.fromList $             zip (S.toList $ freeInType ret) $               repeat True-   in S.fromList $ M.keys $ M.filter id $ alsoRet $ foldl' onType mempty ts+   in S.fromList $ M.keys $ M.filter id $ alsoRet $ foldl' onType mempty $ map snd ts   where     onType uses t = uses <> mustBeExplicitAux t -- Left-biased union. @@ -102,8 +102,8 @@   Scalar $ TypeVar (appres <> als <> arg_als) Unique t targs   where     arg_als = aliases $ maskAliases arg d-returnType _ (Scalar (Arrow old_als v t1 (RetType dims t2))) d arg =-  Scalar $ Arrow als v (t1 `setAliases` mempty) $ RetType dims $ t2 `setAliases` als+returnType _ (Scalar (Arrow old_als v pd t1 (RetType dims t2))) d arg =+  Scalar $ Arrow als v pd (t1 `setAliases` mempty) $ RetType dims $ t2 `setAliases` als   where     -- Make sure to propagate the aliases of an existing closure.     als = old_als <> aliases (maskAliases arg d)@@ -119,12 +119,6 @@   TypeBase shape as maskAliases t Consume = t `setAliases` mempty maskAliases t Observe = t-maskAliases (Scalar (Record ets)) (RecordDiet ds) =-  Scalar $ Record $ M.intersectionWith maskAliases ets ds-maskAliases (Scalar (Sum ets)) (SumDiet ds) =-  Scalar $ Sum $ M.intersectionWith (zipWith maskAliases) ets ds-maskAliases t FuncDiet {} = t-maskAliases _ _ = error "Invalid arguments passed to maskAliases."  -- | The two types are assumed to be structurally equal, but not -- necessarily regarding sizes.  Combines aliases.@@ -140,9 +134,9 @@     sort (M.keys ts1) == sort (M.keys ts2) =       Scalar $ Record $ M.intersectionWith addAliasesFromType ts1 ts2 addAliasesFromType-  (Scalar (Arrow als1 mn1 pt1 (RetType dims1 rt1)))-  (Scalar (Arrow als2 _ _ (RetType _ rt2))) =-    Scalar (Arrow (als1 <> als2) mn1 pt1 (RetType dims1 rt1'))+  (Scalar (Arrow als1 mn1 d1 pt1 (RetType dims1 rt1)))+  (Scalar (Arrow als2 _ _ _ (RetType _ rt2))) =+    Scalar (Arrow (als1 <> als2) mn1 d1 pt1 (RetType dims1 rt1'))     where       rt1' = addAliasesFromType rt1 rt2 addAliasesFromType (Scalar (Sum cs1)) (Scalar (Sum cs2))@@ -203,9 +197,9 @@           (M.intersectionWith (,) ts1 ts2) unifyScalarTypes   uf-  (Arrow as1 mn1 t1 (RetType dims1 t1'))-  (Arrow as2 _ t2 (RetType _ t2')) =-    Arrow (as1 <> as2) mn1+  (Arrow as1 mn1 d1 t1 (RetType dims1 t1'))+  (Arrow as2 _ _ t2 (RetType _ t2')) =+    Arrow (as1 <> as2) mn1 d1       <$> unifyTypesU (flip uf) t1 t2       <*> (RetType dims1 <$> unifyTypesU uf t1' t2') unifyScalarTypes uf (Sum cs1) (Sum cs2)@@ -242,10 +236,31 @@   | otherwise =       pure $ RetType dims st +checkExpForSize ::+  MonadTypeChecker m =>+  ExpBase NoInfo Name ->+  m (Exp, Size)+checkExpForSize (IntLit x NoInfo loc) =+  pure (IntLit x int64_info loc, ConstSize $ fromInteger x)+  where+    int64_info = Info (Scalar (Prim (Signed Int64)))+checkExpForSize (Literal (SignedValue (Int64Value x)) loc) =+  pure (Literal (SignedValue (Int64Value x)) loc, ConstSize x)+checkExpForSize (Var v NoInfo vloc) = do+  v' <- checkNamedSize vloc v+  pure (Var v' int64_info vloc, NamedSize v')+  where+    int64_info = Info (Scalar (Prim (Signed Int64)))+checkExpForSize e =+  typeError+    (locOf e)+    mempty+    "Only variables and i64 literals are allowed in size expressions."+ evalTypeExp ::   MonadTypeChecker m =>-  TypeExp Name ->-  m (TypeExp VName, [VName], StructRetType, Liftedness)+  TypeExp NoInfo Name ->+  m (TypeExp Info VName, [VName], StructRetType, Liftedness) evalTypeExp (TEVar name loc) = do   (name', ps, t, l) <- lookupType loc name   t' <- renameRetType t@@ -307,14 +322,12 @@           <+> dquotes (pretty t)           <+> "(might contain function)."   where-    checkSizeExp SizeExpAny = do+    checkSizeExp (SizeExpAny dloc) = do       dv <- newTypeName "d"-      pure ([dv], SizeExpAny, NamedSize $ qualName dv)-    checkSizeExp (SizeExpConst k dloc) =-      pure ([], SizeExpConst k dloc, ConstSize k)-    checkSizeExp (SizeExpNamed v dloc) = do-      v' <- checkNamedSize loc v-      pure ([], SizeExpNamed v' dloc, NamedSize v')+      pure ([dv], SizeExpAny dloc, NamedSize $ qualName dv)+    checkSizeExp (SizeExp e dloc) = do+      (e', sz) <- checkExpForSize e+      pure ([], SizeExp e' dloc, sz) -- evalTypeExp (TEUnique t loc) = do   (t', svars, RetType dims st, l) <- evalTypeExp t@@ -339,7 +352,7 @@       pure         ( TEArrow (Just v') t1' t2' loc,           svars1 ++ dims1 ++ svars2,-          RetType [] $ Scalar $ Arrow mempty (Named v') st1 (RetType dims2 st2),+          RetType [] $ Scalar $ Arrow mempty (Named v') (diet st1) st1 (RetType dims2 st2),           Lifted         ) --@@ -349,7 +362,9 @@   pure     ( TEArrow Nothing t1' t2' loc,       svars1 ++ dims1 ++ svars2,-      RetType [] $ Scalar $ Arrow mempty Unnamed st1 $ RetType dims2 st2,+      RetType [] . Scalar $+        Arrow mempty Unnamed (diet st1) (st1 `setUniqueness` Nonunique) $+          RetType dims2 st2,       Lifted     ) --@@ -424,7 +439,10 @@   where     tloc = srclocOf ote -    rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name])+    rootAndArgs ::+      MonadTypeChecker m =>+      TypeExp NoInfo Name ->+      m (QualName Name, SrcLoc, [TypeArgExp NoInfo Name])     rootAndArgs (TEVar qn loc) = pure (qn, loc, [])     rootAndArgs (TEApply op arg _) = do       (op', loc, args) <- rootAndArgs op@@ -433,26 +451,24 @@       typeError (srclocOf te') mempty $         "Type" <+> dquotes (pretty te') <+> "is not a type constructor." -    checkArgApply (TypeParamDim pv _) (TypeArgExpDim (SizeExpNamed v dloc) loc) = do-      v' <- checkNamedSize loc v-      pure-        ( TypeArgExpDim (SizeExpNamed v' dloc) loc,-          [],-          M.singleton pv $ SizeSubst $ NamedSize v'-        )-    checkArgApply (TypeParamDim pv _) (TypeArgExpDim (SizeExpConst x dloc) loc) =+    checkSizeExp (SizeExp e dloc) = do+      (e', sz) <- checkExpForSize e       pure-        ( TypeArgExpDim (SizeExpConst x dloc) loc,+        ( TypeArgExpSize (SizeExp e' dloc),           [],-          M.singleton pv $ SizeSubst $ ConstSize x+          SizeSubst sz         )-    checkArgApply (TypeParamDim pv _) (TypeArgExpDim SizeExpAny loc) = do+    checkSizeExp (SizeExpAny loc) = do       d <- newTypeName "d"       pure-        ( TypeArgExpDim SizeExpAny loc,+        ( TypeArgExpSize (SizeExpAny loc),           [d],-          M.singleton pv $ SizeSubst $ NamedSize $ qualName d+          SizeSubst $ NamedSize $ qualName d         )++    checkArgApply (TypeParamDim pv _) (TypeArgExpSize d) = do+      (d', svars, subst) <- checkSizeExp d+      pure (d', svars, M.singleton pv subst)     checkArgApply (TypeParamType _ pv _) (TypeArgExpType te) = do       (te', svars, RetType dims st, _) <- evalTypeExp te       pure@@ -475,8 +491,8 @@ -- * The liftedness of the type. checkTypeExp ::   MonadTypeChecker m =>-  TypeExp Name ->-  m (TypeExp VName, [VName], StructRetType, Liftedness)+  TypeExp NoInfo Name ->+  m (TypeExp Info VName, [VName], StructRetType, Liftedness) checkTypeExp te = do   checkForDuplicateNamesInType te   evalTypeExp te@@ -521,7 +537,7 @@ -- it (normal name shadowing). checkForDuplicateNamesInType ::   MonadTypeChecker m =>-  TypeExp Name ->+  TypeExp NoInfo Name ->   m () checkForDuplicateNamesInType = check mempty   where@@ -547,7 +563,7 @@     check seen (TESum cs _) = mapM_ (mapM (check seen) . snd) cs     check seen (TEApply t1 (TypeArgExpType t2) _) =       check seen t1 >> check seen t2-    check seen (TEApply t1 TypeArgExpDim {} _) =+    check seen (TEApply t1 TypeArgExpSize {} _) =       check seen t1     check seen (TEDim (v : vs) t loc)       | Just prev_loc <- M.lookup v seen =@@ -741,8 +757,8 @@           pure $ Scalar $ TypeVar als u v targs'     onType (Scalar (Record ts)) =       Scalar . Record <$> traverse onType ts-    onType (Scalar (Arrow als v t1 t2)) =-      Scalar <$> (Arrow als v <$> onType t1 <*> onRetType t2)+    onType (Scalar (Arrow als v d t1 t2)) =+      Scalar <$> (Arrow als v d <$> onType t1 <*> onRetType t2)     onType (Scalar (Sum ts)) =       Scalar . Sum <$> traverse (traverse onType) ts 

@@ -482,15 +482,15 @@         (_, Scalar (TypeVar _ _ (QualName [] v2) []))           | Just lvl <- nonrigid v2 ->               link (not ord) v2 lvl t1'-        ( Scalar (Arrow _ p1 a1 (RetType b1_dims b1)),-          Scalar (Arrow _ p2 a2 (RetType b2_dims b2))+        ( Scalar (Arrow _ p1 d1 a1 (RetType b1_dims b1)),+          Scalar (Arrow _ p2 d2 a2 (RetType b2_dims b2))           )-            | uncurry (<) $ swap ord (uniqueness a1) (uniqueness a2) -> do+            | uncurry (<) $ swap ord d1 d2 -> do                 unifyError usage mempty bcs . withIndexLink "unify-consuming-param" $-                  "Parameter types"-                    </> indent 2 (pretty a1)+                  "Parameters"+                    </> indent 2 (pretty d1 <> pretty a1)                     </> "and"-                    </> indent 2 (pretty a2)+                    </> indent 2 (pretty d2 <> pretty a2)                     </> "are incompatible regarding consuming their arguments."             | otherwise -> do                 -- Introduce the existentials as size variables so they

@@ -153,11 +153,18 @@ pScalarType = choice [try pFun, pScalarNonFun]   where     pFun =-      uncurry (Arrow ()) <$> pParam <* lexeme "->" <*> pStructRetType+      pParam <* lexeme "->" <*> pStructRetType     pParam =-      choice [try pNamedParam, (Unnamed,) <$> pNonFunType]-    pNamedParam =-      parens $ (,) <$> (Named <$> pVName) <* lexeme ":" <*> pStructType+      choice+        [ try pNamedParam,+          do+            t <- pNonFunType+            pure $ Arrow () Unnamed (diet t) t+        ]+    pNamedParam = parens $ do+      v <- pVName <* lexeme ":"+      t <- pStructType+      pure $ Arrow () (Named v) (diet t) t  pStructRetType :: Parser StructRetType pStructRetType =

@@ -15,7 +15,7 @@ import Test.Tasty import Test.Tasty.HUnit -evalTest :: TypeExp Name -> Either String ([VName], StructRetType) -> TestTree+evalTest :: TypeExp NoInfo Name -> Either String ([VName], StructRetType) -> TestTree evalTest te expected =   testCase (prettyString te) $     case (fmap (extract . fst) (run (checkTypeExp te)), expected) of